package project

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/xerrors"
	"gorm.io/gorm"

	"github.com/pingcap-inc/security/baseline/api/v1/controller"
	"github.com/pingcap-inc/security/baseline/api/v1/controller/errors"
	"github.com/pingcap-inc/security/baseline/api/v1/model"
	"github.com/pingcap-inc/security/baseline/internal/trivy"
	"github.com/pingcap-inc/security/baseline/pkg/log"
	"github.com/pingcap-inc/security/baseline/pkg/utils"
)

type createTaskReq struct {
	Type string `json:"type" validate:"required,min=5,max=64"`
	Url  string `json:"url" validate:"required,min=5,max=256"`
}

func (req *createTaskReq) ParseUrl() (task map[string]string, err error) {
	defer func() {
		if task != nil {
			task["type"] = req.Type
			task["status"] = model.TaskCreate
		}
	}()

	switch req.Type {
	case "image":
		return utils.ParseImageUrl(req.Url)
	case "repository":
		task, _, err := utils.ParseRepoUrl(req.Url)
		return task, err
	default:
		return nil, errors.NewErrorf(errors.BL_PARAMETER_INVALID, "Task type `%s` not supported", req.Type)
	}
}

type createTaskRsp struct {
	Task *model.TaskModel `json:"task"`
}

func createNewTask(tx *gorm.DB, project_id string, req *createTaskReq) (*createTaskRsp, error) {
	taskInfo, err := req.ParseUrl()
	if err != nil {
		return nil, err
	}
	project := model.ProjectModel{}
	task := model.TaskModel{
		Type:    taskInfo["type"],
		Url:     taskInfo["url"],
		Name:    taskInfo["name"],
		Version: taskInfo["version"],
		Status:  taskInfo["status"],
	}

	err = tx.Transaction(func(tx *gorm.DB) error {
		if err := tx.Where(
			"id == ?", project_id,
		).First(&project).Error; err != nil {
			return err
		} else {
			task.Project = project
			task.ProjectID = project.ID
			return tx.Create(&task).Error
		}
	})

	return &createTaskRsp{
		Task: &task,
	}, err
}

func runImageTask(task *model.TaskModel, opt *trivy.ImageOption) ([]byte, error) {
	if err := trivy.UpdateDB(opt.BaseOption); err != nil {
		return nil, err
	}

	return trivy.ExecTrivyImage(opt)
}

func runRepoTask(task *model.TaskModel, opt *trivy.RepoOption) ([]byte, error) {
	if err := trivy.UpdateDB(opt.BaseOption); err != nil {
		return nil, err
	}

	return trivy.ExecTrivyRepo(opt)
}

func checkTaskStatus(tx *gorm.DB, task *model.TaskModel) (string, error) {
	err := tx.Transaction(func(tx *gorm.DB) error {
		return tx.Model(task).First(&task).Error
	})

	return task.Status, err
}

func updateTaskStatus(tx *gorm.DB, task *model.TaskModel, taskStatus string) error {
	return tx.Transaction(func(tx *gorm.DB) error {
		return tx.Model(task).Update("Status", taskStatus).Error
	})
}

func updateTaskReport(tx *gorm.DB, task *model.TaskModel, report []byte) error {
	reportMap := make(map[string]interface{})
	if err := json.Unmarshal(report, &reportMap); err != nil {
		return err
	}

	results, ok := reportMap["Results"].([]interface{})
	if !ok {
		return xerrors.Errorf("parse trivy report failed: %s", string(report))
	}

	vulnerabilityMap := map[string]model.VulnerabilityModel{}

	for _, item := range results {
		t := item.(map[string]interface{})
		target := "Unknown"
		if s, ok := t["Target"].(string); ok {
			target = s
		}
		target_class := "Unknown"
		if s, ok := t["Class"].(string); ok {
			target_class = s
		}

		target_type := "Unknown"
		if s, ok := t["Type"].(string); ok {
			target_type = s
		}

		if vulnerabilities, ok := t["Vulnerabilities"].([]interface{}); ok {
			for _, vulnerability := range vulnerabilities {
				v := vulnerability.(map[string]interface{})

				vulnerabilityID := "Unknown"
				if s, ok := v["VulnerabilityID"].(string); ok {
					vulnerabilityID = s
				}

				packageName := "Unknown"
				if s, ok := v["PkgName"].(string); ok {
					packageName = s
				}

				severity := "Unknown"
				if s, ok := v["Severity"].(string); ok {
					severity = s
				}

				title := "Unknown"
				if s, ok := v["Title"].(string); ok {
					title = s
				}

				description := "Unknown"
				if s, ok := v["Description"].(string); ok {
					description = s
				}

				installedVersion := "Unknown"
				if s, ok := v["InstalledVersion"].(string); ok {
					installedVersion = s
				}

				fixedVersion := "Unknown"
				if s, ok := v["FixedVersion"].(string); ok {
					fixedVersion = s
				}

				key := vulnerabilityID + packageName + target + target_class + target_type
				vulnerabilityMap[key] = model.VulnerabilityModel{
					Target: target,
					Class:  target_class,
					Type:   target_type,

					VulnerabilityID:  vulnerabilityID,
					PackageName:      packageName,
					Severity:         severity,
					Title:            title,
					Description:      description,
					InstalledVersion: installedVersion,
					FixedVersion:     fixedVersion,

					Task:   *task,
					TaskID: task.ID,
				}
			}
		}
	}

	return tx.Transaction(func(tx *gorm.DB) error {
		if err := tx.Where("task_id == ?", task.ID).Delete(&[]model.VulnerabilityModel{}).Error; err != nil {
			return err
		}

		for _, vulnerability := range vulnerabilityMap {
			if err := tx.Create(&vulnerability).Error; err != nil {
				return err
			}
		}

		return nil
	})
}

func runTask(tx *gorm.DB, task *model.TaskModel, opt *trivy.BaseOption, severity string) {
	if err := updateTaskStatus(tx, task, model.TaskRunning); err != nil {
		log.Logger.Errorf("Update task status failed, %+v", err)
	} else {
		report, err := func(task *model.TaskModel, opt *trivy.BaseOption, severity string) ([]byte, error) {
			switch task.Type {
			case "image":
				opt.CacheDir = utils.TrivyCacheDir(opt.CacheDir)
				return runImageTask(task, &trivy.ImageOption{
					BaseOption: opt,
					Severity:   severity,
					Targets:    []string{task.Url},
				})
			case "repository":
				opt.CacheDir = utils.TrivyCacheDir(opt.CacheDir)
				return runRepoTask(task, &trivy.RepoOption{
					BaseOption: opt,
					Severity:   severity,
					Targets:    []string{task.Url},
				})
			default:
				return nil, errors.NewErrorf(
					errors.BL_INTERNAL_SERVER_ERROR, "Run task failed, task type `%s` not supported", task.Type)
			}
		}(task, opt, severity)

		if err != nil {
			log.Logger.Errorf("Exec task `%+v` failed, %+v", task, err)

			if err := updateTaskStatus(tx, task, model.TaskFailed); err != nil {
				log.Logger.Errorf("Update task status failed, %+v", err)
			}
		} else {
			if report != nil {
				if err := updateTaskReport(tx, task, report); err != nil {
					log.Logger.Errorf("Update task report failed, %+v", err)
					if err := updateTaskStatus(tx, task, model.TaskFailed); err != nil {
						log.Logger.Errorf("Update task status failed, %+v", err)
					}
				} else {
					if err := updateTaskStatus(tx, task, model.TaskSuccess); err != nil {
						log.Logger.Errorf("Update task status failed, %+v", err)
					}
				}
			}
		}
	}
}

func createTaskHandler(ctx *gin.Context, rpcReq *controller.RpcRequest) (*controller.RpcResponse, error) {
	req := rpcReq.Request.(*createTaskReq)
	rsp, err := createNewTask(ctx.MustGet("db").(*gorm.DB), ctx.Param("project_id"), req)
	if err != nil {
		return nil, err
	} else {
		if response, err := json.Marshal(rsp); err != nil {
			return nil, err
		} else {
			go runTask(ctx.MustGet("db").(*gorm.DB), rsp.Task, &trivy.BaseOption{
				GITHUB_TOKEN: ctx.MustGet("GITHUB_TOKEN").(string),
				CacheDir:     ctx.MustGet("CACHE_DIR").(string),
				Timeout:      ctx.MustGet("TIMEOUT").(time.Duration),
			}, utils.DefaultTrivySeverity())

			return &controller.RpcResponse{
				Code:     errors.BL_SUCCESS,
				Response: string(response),
			}, err
		}
	}
}

// Profile Project profile
// @Summary Create Task
// @Description Create a new scan task
// @Tags Project
// @Accept application/json
// @Produce application/json
// @Param project_id path int true "project id"
// @Param createTaskReq body createTaskReq true "create task request"
// @Success 200 {object} controller.CommonResult{data=createTaskRsp}
// @Failure 500 {object} controller.CommonResult
// @Router /project/{project_id}/task [post]
func CreateTaskHandler(ctx *gin.Context) {
	if requestBody, ok := controller.HandleJsonRequestFromBody(ctx, &createTaskReq{}); ok {
		controller.InvokeHandler(ctx, createTaskHandler, requestBody, &createTaskReq{}, &createTaskRsp{})
	}
}

type getTaskInfoReq struct{}

type taskInfo struct {
	*model.TaskModel

	Critical int `json:"critical"`
	High     int `json:"high"`
	Medium   int `json:"medium"`
	Low      int `json:"low"`
	Unknown  int `json:"unknown"`
	Total    int `json:"total"`
}

type getTaskInfoRsp struct {
	Task *taskInfo `json:"task"`
}

func findTaskInfoById(tx *gorm.DB, project_id string, task_id string) (*getTaskInfoRsp, error) {
	task := model.TaskModel{}
	vulnerabilities := []model.VulnerabilityModel{}

	err := tx.Transaction(func(tx *gorm.DB) error {
		if err := tx.Where("project_id == ? and id == ?", project_id, task_id).First(&task).Error; err != nil {
			return err
		} else {
			if task.Status == model.TaskSuccess {
				return tx.Where("task_id == ?", task_id).Find(&vulnerabilities).Error
			}
		}

		return nil
	})

	if err != nil {
		return nil, err
	} else {
		info := taskInfo{
			TaskModel: &task,
		}

		for _, v := range vulnerabilities {
			info.Total++
			switch v.Severity {
			case "CRITICAL":
				info.Critical++
			case "HIGH":
				info.High++
			case "MEDIUM":
				info.Medium++
			case "LOW":
				info.Low++
			case "UNKNOWN":
				info.Unknown++
			}
		}

		return &getTaskInfoRsp{
			Task: &info,
		}, err
	}
}

func getTaskInfoHandler(ctx *gin.Context, rpcReq *controller.RpcRequest) (*controller.RpcResponse, error) {
	rsp, err := findTaskInfoById(ctx.MustGet("db").(*gorm.DB), ctx.Param("project_id"), ctx.Param("task_id"))
	if err != nil {
		return nil, err
	} else {
		if response, err := json.Marshal(rsp); err != nil {
			return nil, err
		} else {
			return &controller.RpcResponse{
				Code:     errors.BL_SUCCESS,
				Response: string(response),
			}, err
		}
	}
}

// Profile Project profile
// @Summary Get Task Info
// @Description Get specific task info
// @Tags Project
// @Accept application/json
// @Produce application/json
// @Param project_id path int true "project id"
// @Param task_id path int true "task id"
// @Success 200 {object} controller.CommonResult{data=getTaskInfoRsp}
// @Failure 500 {object} controller.CommonResult
// @Router /project/{project_id}/task/{task_id} [get]
func GetTaskInfoHandler(ctx *gin.Context) {
	if requestBody, ok := controller.HandleJsonRequestFromQuery(ctx, &getTaskInfoReq{}); ok {
		controller.InvokeHandler(ctx, getTaskInfoHandler, requestBody, &getTaskInfoReq{}, &getTaskInfoRsp{})
	}
}

type getTaskListReq struct {
	controller.PageRequest
}

type getTaskListRsp struct {
	Tasks []taskInfo `json:"tasks"`
}

func findAllTasksInPage(tx *gorm.DB, project_id string, req *getTaskListReq) (*getTaskListRsp, *controller.Page, error) {
	tasks := []model.TaskModel{}
	tasksInfo := []taskInfo{}

	page := controller.Page{
		PageRequest: req.PageRequest,
	}

	err := tx.Transaction(func(tx *gorm.DB) error {
		if allTasks := tx.Where("project_id == ?", project_id).Find(&[]model.TaskModel{}); allTasks.Error != nil {
			return allTasks.Error
		} else {
			allTasks.Count(&page.Total)

			return allTasks.Offset((req.Page - 1) * req.PageSize).Limit(req.PageSize).Order("updated_at desc").Find(&tasks).Error
		}
	})

	for _, task := range tasks {
		rsp, err := findTaskInfoById(tx, project_id, fmt.Sprintf("%d", task.ID))
		if err != nil {
			return nil, &page, err
		} else {
			tasksInfo = append(tasksInfo, *rsp.Task)
		}
	}

	return &getTaskListRsp{
		Tasks: tasksInfo,
	}, &page, err
}

func getTaskListHandler(ctx *gin.Context, rpcReq *controller.RpcRequest) (*controller.RpcResponse, error) {
	req := rpcReq.Request.(*getTaskListReq)
	rsp, page, err := findAllTasksInPage(ctx.MustGet("db").(*gorm.DB), ctx.Param("project_id"), req)
	if err != nil {
		return nil, err
	} else {
		if response, err := json.Marshal(rsp); err != nil {
			return nil, err
		} else {
			return &controller.RpcResponse{
				Code:     errors.BL_SUCCESS,
				Response: string(response),
				Page:     page,
			}, err
		}
	}
}

// Profile Project profile
// @Summary Get Task List
// @Description Get all tasks in specific project
// @Tags Project
// @Accept application/json
// @Produce application/json
// @Param project_id path int true "project id"
// @Param getTaskListReq query getTaskListReq true "get task list request"
// @Success 200 {object} controller.CommonResult{data=getTaskListRsp}
// @Failure 500 {object} controller.CommonResult
// @Router /project/{project_id}/task [get]
func GetTaskListHandler(ctx *gin.Context) {
	if requestBody, ok := controller.HandleJsonRequestFromQuery(ctx, &getTaskListReq{
		PageRequest: controller.PageRequest{
			PageSize: 20,
		},
	}); ok {
		controller.InvokeHandler(ctx, getTaskListHandler, requestBody, &getTaskListReq{}, &getTaskListRsp{})
	}
}

type getTaskReportReq struct{}

type getTaskReportRsp struct {
	Task            *taskInfo                  `json:"task"`
	Vulnerabilities []model.VulnerabilityModel `json:"vulnerabilities"`
}

func findTaskReportById(tx *gorm.DB, project_id string, task_id string) (*getTaskReportRsp, error) {
	vulnerabilities := []model.VulnerabilityModel{}

	err := tx.Transaction(func(tx *gorm.DB) error {
		task := model.TaskModel{}
		if err := tx.Where("project_id == ? and id == ?", project_id, task_id).First(&task).Error; err != nil {
			return err
		} else {
			if task.Status == model.TaskSuccess {
				return tx.Where("task_id == ?", task_id).Find(&vulnerabilities).Error
			}
		}

		return nil
	})

	if err != nil {
		return nil, err
	} else {
		rsp, err := findTaskInfoById(tx, project_id, task_id)
		if err != nil {
			return nil, err
		} else {
			return &getTaskReportRsp{
				Task:            rsp.Task,
				Vulnerabilities: vulnerabilities,
			}, err
		}
	}
}

func getTaskReportHandler(ctx *gin.Context, rpcReq *controller.RpcRequest) (*controller.RpcResponse, error) {
	rsp, err := findTaskReportById(ctx.MustGet("db").(*gorm.DB), ctx.Param("project_id"), ctx.Param("task_id"))
	if err != nil {
		return nil, err
	} else {
		if response, err := json.Marshal(rsp); err != nil {
			return nil, err
		} else {
			return &controller.RpcResponse{
				Code:     errors.BL_SUCCESS,
				Response: string(response),
			}, err
		}
	}
}

// Profile Project profile
// @Summary Get Task Report
// @Description Get specific task report
// @Tags Project
// @Accept application/json
// @Produce application/json
// @Param project_id path int true "project id"
// @Param task_id path int true "task id"
// @Param getTaskReportReq query getTaskReportReq true "get task report request"
// @Success 200 {object} controller.CommonResult{data=getTaskReportRsp}
// @Failure 500 {object} controller.CommonResult
// @Router /project/{project_id}/task/{task_id}/report [get]
func GetTaskReportHandler(ctx *gin.Context) {
	if requestBody, ok := controller.HandleJsonRequestFromQuery(ctx, &getTaskReportReq{}); ok {
		controller.InvokeHandler(ctx, getTaskReportHandler, requestBody, &getTaskReportReq{}, &getTaskReportRsp{})
	}
}

type deleteTaskReq struct{}

type deleteTaskRsp struct{}

func deleteTaskById(tx *gorm.DB, project_id string, task_id string) (*deleteTaskRsp, error) {
	err := tx.Transaction(func(tx *gorm.DB) error {
		project := model.ProjectModel{}
		if err := tx.Where(
			"id == ?", project_id,
		).First(&project).Error; err != nil {
			return err
		} else {
			task := model.TaskModel{}
			if err := tx.Where("project_id == ? and id == ?", project_id, task_id).First(&task).Error; err != nil {
				return err
			} else {
				if err := tx.Where("task_id == ?", task_id).Delete(&[]model.VulnerabilityModel{}).Error; err != nil {
					return err
				} else {
					return tx.Delete(&task).Error
				}
			}
		}
	})

	return &deleteTaskRsp{}, err
}

func deleteTaskHandler(ctx *gin.Context, rpcReq *controller.RpcRequest) (*controller.RpcResponse, error) {
	rsp, err := deleteTaskById(ctx.MustGet("db").(*gorm.DB), ctx.Param("project_id"), ctx.Param("task_id"))
	if err != nil {
		return nil, err
	} else {
		if response, err := json.Marshal(rsp); err != nil {
			return nil, err
		} else {
			return &controller.RpcResponse{
				Code:     errors.BL_SUCCESS,
				Response: string(response),
			}, err
		}
	}
}

// Profile Project profile
// @Summary Delete Task
// @Description Delete specific task
// @Tags Project
// @Accept application/json
// @Produce application/json
// @Param project_id path int true "project id"
// @Param task_id path int true "task id"
// @Success 200 {object} controller.CommonResult{data=deleteTaskRsp}
// @Failure 500 {object} controller.CommonResult
// @Router /project/{project_id}/task/{task_id} [delete]
func DeleteTaskHandler(ctx *gin.Context) {
	if requestBody, ok := controller.HandleJsonRequestFromQuery(ctx, &deleteProjectReq{}); ok {
		controller.InvokeHandler(ctx, deleteTaskHandler, requestBody, &deleteProjectReq{}, &deleteProjectRsp{})
	}
}