package vulnerability

import (
	"strings"
	"time"
)

type Metadata struct {
	ID             string
	DataSource     string // the primary reference URL, i.e. where the data originated
	Namespace      string
	Severity       string
	URLs           []string // secondary reference URLs a vulnerability may provide
	Description    string
	Cvss           []Cvss
	KnownExploited []KnownExploited
	EPSS           []EPSS
	CWEs           []CWE

	// calculated as-needed
	risk float64
}

// RiskScore computes a basic quantitative risk by combining threat and severity.
// Threat is represented by epss (likelihood of exploitation), and severity by the cvss base score + string severity.
// Impact is currently fixed at 1 and may be integrated into the calculation in future versions.
// Raw risk is epss * (cvss / 10) * impact, then scaled to 0–100 for readability.
// If a vulnerability appears in the KEV list, apply an additional boost to reflect known exploitation.
// Known ransomware campaigns receive a further, distinct boost.
func (m *Metadata) RiskScore() float64 {
	if m == nil {
		return 0
	}
	if m.risk != 0 {
		return m.risk
	}
	m.risk = riskScore(*m)
	return m.risk
}

func riskScore(m Metadata) float64 {
	return min(threat(m)*severity(m)*kevModifier(m), 1.0) * 100.0
}

func kevModifier(m Metadata) float64 {
	if len(m.KnownExploited) > 0 {
		for _, kev := range m.KnownExploited {
			if strings.ToLower(kev.KnownRansomwareCampaignUse) == "known" {
				// consider ransomware campaigns to be a greater kevModifier than other KEV threats
				return 1.1
			}
		}
		return 1.05 // boost the final result, as if there is a greater kevModifier inherently from KEV threats
	}
	return 1.0
}

func threat(m Metadata) float64 {
	if len(m.KnownExploited) > 0 {
		// per the EPSS guidance, any evidence of exploitation in the wild (not just PoC) should be considered over EPSS data
		return 1.0
	}
	if len(m.EPSS) == 0 {
		return 0.0
	}
	return m.EPSS[0].EPSS
}

// severity returns a 0-1 value, which is a combination of the string severity and the average of the cvss base scores.
// If there are no cvss scores, the string severity is used. Some vendors only update the string severity and not the
// cvss scores, so it's important to consider all sources. We are also not biasing towards any one source (multiple
// cvss scores won't over-weigh the string severity).
func severity(m Metadata) float64 {
	// TODO: summarization should take a policy: prefer NVD over CNA or vice versa...

	stringSeverityScore := severityToScore(m.Severity) / 10.0
	avgBaseScore := average(validBaseScores(m.Cvss...)...) / 10.0
	if avgBaseScore == 0 {
		return stringSeverityScore
	}
	return average(stringSeverityScore, avgBaseScore)
}

func severityToScore(severity string) float64 {
	// use the middle of the range for each severity
	switch strings.ToLower(severity) {
	case "negligible":
		return 0.5
	case "low":
		return 3.0
	case "medium":
		return 5.0
	case "high":
		return 7.5
	case "critical":
		return 9.0
	}
	// the severity value might be "unknown" or an unexpected value. These should not be lost
	// in the noise and placed at the bottom of the list... instead we compromise to the middle of the list.
	return 5.0
}

func validBaseScores(as ...Cvss) []float64 {
	var out []float64
	for _, a := range as {
		if a.Metrics.BaseScore == 0 {
			// this is a mistake... base scores cannot be 0. Don't include this value and bring down the average
			continue
		}
		out = append(out, a.Metrics.BaseScore)
	}
	return out
}

func average(as ...float64) float64 {
	if len(as) == 0 {
		return 0
	}
	sum := 0.0
	for _, a := range as {
		sum += a
	}
	return sum / float64(len(as))
}

type Cvss struct {
	Source         string
	Type           string
	Version        string
	Vector         string
	Metrics        CvssMetrics
	VendorMetadata any
}

type CvssMetrics struct {
	BaseScore           float64
	ExploitabilityScore *float64
	ImpactScore         *float64
}

type KnownExploited struct {
	CVE                        string
	VendorProject              string
	Product                    string
	DateAdded                  *time.Time
	RequiredAction             string
	DueDate                    *time.Time
	KnownRansomwareCampaignUse string
	Notes                      string
	URLs                       []string
	CWEs                       []string
}

type EPSS struct {
	CVE        string
	EPSS       float64
	Percentile float64
	Date       time.Time
}

type CWE struct {
	CVE    string
	CWE    string
	Source string
	Type   string
}
