package dbsearch

import (
	"fmt"
	"strings"
	"time"

	"github.com/hashicorp/go-multierror"
	"github.com/scylladb/go-set/strset"

	v6 "github.com/anchore/grype/grype/db/v6"
)

type canonicalVulnerability interface {
	getCVEs() []string
	decorate(kevs []KnownExploited, epss []EPSS)
}

// vulnerabilityDecorations are separate model elements (not from VulnerabilityHandle) that is fetched on
// the provider.GetMetadata() path instead of the provider.GetVulnerabilities() path. I hope for these two paths
// to be merged into the same path come grype 1.0, in which case these elements would already be on the
// store get methods when crafting a vulnerability.
type vulnerabilityDecorations struct {
	KnownExploited []KnownExploited `json:"knownExploited,omitempty"`
	EPSS           []EPSS           `json:"epss,omitempty"`
}

func decorateVulnerabilities(reader v6.VulnerabilityDecoratorStoreReader, cvs ...canonicalVulnerability) error {
	for _, cv := range cvs {
		cves := cv.getCVEs()
		if len(cves) == 0 {
			continue
		}

		knownExploited, err := fetchKnownExploited(reader, cves)
		if err != nil {
			return fmt.Errorf("unable to get known exploited vulnerabilities: %w", err)
		}

		epss, err := fetchEpss(reader, cves)
		if err != nil {
			return fmt.Errorf("unable to get EPSS scores: %w", err)
		}

		cv.decorate(knownExploited, epss)
	}
	return nil
}

func (afj *vulnerabilityAffectedPackageJoin) getCVEs() []string {
	if afj == nil {
		return nil
	}
	return getCVEs(&afj.Vulnerability)
}

func getCVEs(v *v6.VulnerabilityHandle) []string {
	var cves []string
	set := strset.New()

	addCVE := func(id string) {
		lower := strings.ToLower(id)
		if strings.HasPrefix(lower, "cve-") {
			if !set.Has(lower) {
				cves = append(cves, id)
				set.Add(lower)
			}
		}
	}

	if v == nil {
		return cves
	}

	addCVE(v.Name)

	if v.BlobValue == nil {
		return cves
	}

	addCVE(v.BlobValue.ID)

	for _, alias := range v.BlobValue.Aliases {
		addCVE(alias)
	}

	return cves
}

func (vd *vulnerabilityDecorations) decorate(kevs []KnownExploited, epss []EPSS) {
	if vd == nil {
		return
	}

	vd.KnownExploited = kevs
	vd.EPSS = epss
}

func fetchKnownExploited(reader v6.VulnerabilityDecoratorStoreReader, cves []string) ([]KnownExploited, error) {
	var out []KnownExploited
	var errs error
	for _, cve := range cves {
		kevs, err := reader.GetKnownExploitedVulnerabilities(cve)
		if err != nil {
			errs = multierror.Append(errs, err)
			continue
		}
		for _, kev := range kevs {
			out = append(out, KnownExploited{
				CVE:                        kev.Cve,
				VendorProject:              kev.BlobValue.VendorProject,
				Product:                    kev.BlobValue.Product,
				DateAdded:                  kev.BlobValue.DateAdded.Format(time.DateOnly),
				RequiredAction:             kev.BlobValue.RequiredAction,
				DueDate:                    kev.BlobValue.DueDate.Format(time.DateOnly),
				KnownRansomwareCampaignUse: kev.BlobValue.KnownRansomwareCampaignUse,
				Notes:                      kev.BlobValue.Notes,
				URLs:                       kev.BlobValue.URLs,
				CWEs:                       kev.BlobValue.CWEs,
			})
		}
	}
	return out, errs
}

func fetchEpss(reader v6.VulnerabilityDecoratorStoreReader, cves []string) ([]EPSS, error) {
	var out []EPSS
	var errs error
	for _, cve := range cves {
		entries, err := reader.GetEpss(cve)
		if err != nil {
			errs = multierror.Append(errs, err)
			continue
		}
		for _, entry := range entries {
			out = append(out, EPSS{
				CVE:        entry.Cve,
				EPSS:       entry.Epss,
				Percentile: entry.Percentile,
				Date:       entry.Date.Format(time.DateOnly),
			})
		}
	}
	return out, errs
}
