// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package vulnerability

import (
	"context"
	"sync"
	"time"

	"github.com/goharbor/harbor/src/common/rbac"
	"github.com/goharbor/harbor/src/controller/artifact"
	scanCtl "github.com/goharbor/harbor/src/controller/scan"
	"github.com/goharbor/harbor/src/jobservice/job"
	"github.com/goharbor/harbor/src/jobservice/logger"
	"github.com/goharbor/harbor/src/lib/errors"
	"github.com/goharbor/harbor/src/lib/log"
	"github.com/goharbor/harbor/src/lib/orm"
	"github.com/goharbor/harbor/src/pkg/permission/types"
	"github.com/goharbor/harbor/src/pkg/robot/model"
	scanJob "github.com/goharbor/harbor/src/pkg/scan"
	"github.com/goharbor/harbor/src/pkg/scan/dao/scan"
	"github.com/goharbor/harbor/src/pkg/scan/dao/scanner"
	"github.com/goharbor/harbor/src/pkg/scan/postprocessors"
	"github.com/goharbor/harbor/src/pkg/scan/report"
	v1 "github.com/goharbor/harbor/src/pkg/scan/rest/v1"
	"github.com/goharbor/harbor/src/pkg/task"
)

func init() {
	scanJob.RegisterScanHanlder(v1.ScanTypeVulnerability, &scanHandler{
		reportConverter:    postprocessors.Converter,
		ReportMgrFunc:      func() report.Manager { return report.Mgr },
		TaskMgrFunc:        func() task.Manager { return task.Mgr },
		ScanControllerFunc: func() scanCtl.Controller { return scanCtl.DefaultController },
		cloneCtx:           orm.Clone,
	})
}

// scanHandler defines the handler for scan vulnerability
type scanHandler struct {
	reportConverter    postprocessors.NativeScanReportConverter
	ReportMgrFunc      func() report.Manager
	TaskMgrFunc        func() task.Manager
	ScanControllerFunc func() scanCtl.Controller
	cloneCtx           func(ctx context.Context) context.Context
}

func (h *scanHandler) MakePlaceHolder(ctx context.Context, art *artifact.Artifact,
	r *scanner.Registration) (rps []*scan.Report, err error) {
	mimeTypes := r.GetProducesMimeTypes(art.ManifestMediaType, v1.ScanTypeVulnerability)
	reportMgr := h.ReportMgrFunc()
	oldReports, err := reportMgr.GetBy(h.cloneCtx(ctx), art.Digest, r.UUID, mimeTypes)
	if err != nil {
		return nil, err
	}

	if err := h.assembleReports(ctx, oldReports...); err != nil {
		return nil, err
	}

	if len(oldReports) > 0 {
		for _, oldReport := range oldReports {
			if !job.Status(oldReport.Status).Final() {
				return nil, errors.ConflictError(nil).WithMessagef("a previous scan process is %s", oldReport.Status)
			}
		}

		for _, oldReport := range oldReports {
			if err := reportMgr.Delete(ctx, oldReport.UUID); err != nil {
				return nil, err
			}
		}
	}

	var reports []*scan.Report

	for _, pm := range r.GetProducesMimeTypes(art.ManifestMediaType, v1.ScanTypeVulnerability) {
		rpt := &scan.Report{
			Digest:           art.Digest,
			RegistrationUUID: r.UUID,
			MimeType:         pm,
		}

		create := func(ctx context.Context) error {
			reportUUID, err := reportMgr.Create(ctx, rpt)
			if err != nil {
				return err
			}
			rpt.UUID = reportUUID

			return nil
		}

		if err := orm.WithTransaction(create)(orm.SetTransactionOpNameToContext(ctx, "tx-make-report-placeholder")); err != nil {
			return nil, err
		}

		reports = append(reports, rpt)
	}

	return reports, nil
}

func (h *scanHandler) assembleReports(ctx context.Context, reports ...*scan.Report) error {
	reportUUIDs := make([]string, len(reports))
	for i, report := range reports {
		reportUUIDs[i] = report.UUID
	}

	tasks, err := h.listScanTasks(ctx, reportUUIDs)
	if err != nil {
		return err
	}

	reportUUIDToTasks := map[string]*task.Task{}
	for _, task := range tasks {
		for _, reportUUID := range scanCtl.GetReportUUIDs(task.ExtraAttrs) {
			reportUUIDToTasks[reportUUID] = task
		}
	}

	for _, report := range reports {
		if task, ok := reportUUIDToTasks[report.UUID]; ok {
			report.Status = task.Status
			report.StartTime = task.StartTime
			report.EndTime = task.EndTime
		} else {
			report.Status = job.ErrorStatus.String()
		}

		completeReport, err := h.reportConverter.FromRelationalSchema(ctx, report.UUID, report.Digest, report.Report)
		if err != nil {
			return err
		}
		report.Report = completeReport
	}

	return nil
}

// listScanTasks returns the tasks of the reports
func (h *scanHandler) listScanTasks(ctx context.Context, reportUUIDs []string) ([]*task.Task, error) {
	if len(reportUUIDs) == 0 {
		return nil, nil
	}

	tasks := make([]*task.Task, len(reportUUIDs))
	errs := make([]error, len(reportUUIDs))

	var wg sync.WaitGroup
	for i, reportUUID := range reportUUIDs {
		wg.Add(1)

		go func(ix int, reportUUID string) {
			defer wg.Done()

			task, err := h.getScanTask(h.cloneCtx(ctx), reportUUID)
			if err == nil {
				tasks[ix] = task
			} else if !errors.IsNotFoundErr(err) {
				errs[ix] = err
			} else {
				log.G(ctx).Warningf("task for the scan report %s not found", reportUUID)
			}
		}(i, reportUUID)
	}
	wg.Wait()

	for _, err := range errs {
		if err != nil {
			return nil, err
		}
	}

	var results []*task.Task
	for _, task := range tasks {
		if task != nil {
			results = append(results, task)
		}
	}

	return results, nil
}

func (h *scanHandler) getScanTask(ctx context.Context, reportUUID string) (*task.Task, error) {
	// NOTE: the method uses the postgres' unique operations and should consider here if support other database in the future.
	taskMgr := h.TaskMgrFunc()
	tasks, err := taskMgr.ListScanTasksByReportUUID(ctx, reportUUID)
	if err != nil {
		return nil, err
	}

	if len(tasks) == 0 {
		return nil, errors.NotFoundError(nil).WithMessagef("task for report %s not found", reportUUID)
	}

	return tasks[0], nil
}

func (h *scanHandler) GetPlaceHolder(ctx context.Context, _ string, artDigest, scannerUUID string,
	mimeType string) (rp *scan.Report, err error) {
	reportMgr := h.ReportMgrFunc()
	reports, err := reportMgr.GetBy(ctx, artDigest, scannerUUID, []string{mimeType})
	if err != nil {
		logger.Errorf("failed to get report for artifact %s of mimetype %s, error %v", artDigest, mimeType, err)
		return nil, err
	}
	if len(reports) == 0 {
		logger.Errorf("no report found for artifact %s of mimetype %s, error %v", artDigest, mimeType, err)
		return nil, errors.NotFoundError(nil).WithMessage("no report found to update data")
	}
	return reports[0], nil
}

// RequestProducesMineTypes returns the produces mime types
func (h *scanHandler) RequestProducesMineTypes() []string {
	return []string{v1.MimeTypeGenericVulnerabilityReport}
}

// RequestParameters defines the parameters for scan request
func (h *scanHandler) RequestParameters() map[string]interface{} {
	return nil
}

// RequiredPermissions defines the permission used by the scan robot account
func (h *scanHandler) RequiredPermissions() []*types.Policy {
	return []*types.Policy{
		{
			Resource: rbac.ResourceRepository,
			Action:   rbac.ActionPull,
		},
		{
			Resource: rbac.ResourceRepository,
			Action:   rbac.ActionScannerPull,
		},
	}
}

// PostScan ...
func (h *scanHandler) PostScan(ctx job.Context, _ *v1.ScanRequest, origRp *scan.Report, rawReport string,
	_ time.Time, _ *model.Robot) (string, error) {
	// use a new ormer here to use the short db connection
	_, refreshedReport, err := postprocessors.Converter.ToRelationalSchema(ctx.SystemContext(), origRp.UUID,
		origRp.RegistrationUUID, origRp.Digest, rawReport)
	return refreshedReport, err
}

// URLParameter vulnerability doesn't require any scan report parameters
func (h *scanHandler) URLParameter(_ *v1.ScanRequest) (string, error) {
	return "", nil
}

func (h *scanHandler) Update(ctx context.Context, uuid string, rpt string) error {
	reportMgr := h.ReportMgrFunc()
	if err := reportMgr.UpdateReportData(ctx, uuid, rpt); err != nil {
		return err
	}
	return nil
}

func (h *scanHandler) GetSummary(ctx context.Context, ar *artifact.Artifact, mimeTypes []string) (map[string]interface{}, error) {
	bc := h.ScanControllerFunc()
	if ar == nil {
		return nil, errors.New("no way to get report summaries for nil artifact")
	}
	// Get reports first
	rps, err := bc.GetReport(ctx, ar, mimeTypes)
	if err != nil {
		return nil, err
	}
	summaries := make(map[string]interface{}, len(rps))
	for _, rp := range rps {
		sum, err := report.GenerateSummary(rp)
		if err != nil {
			return nil, err
		}

		if s, ok := summaries[rp.MimeType]; ok {
			r, err := report.MergeSummary(rp.MimeType, s, sum)
			if err != nil {
				return nil, err
			}

			summaries[rp.MimeType] = r
		} else {
			summaries[rp.MimeType] = sum
		}
	}

	return summaries, nil
}

func (h *scanHandler) JobVendorType() string {
	return job.ImageScanJobVendorType
}
