package vulnerability

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestRiskScore(t *testing.T) {
	tests := []struct {
		name     string
		metadata Metadata
		expected float64
	}{
		{
			name:     "nil metadata",
			metadata: Metadata{},
			expected: 0,
		},
		{
			name: "already calculated risk",
			metadata: Metadata{
				risk: 42.5,
			},
			expected: 42.5,
		},
		{
			name: "no EPSS data, no KEV",
			metadata: Metadata{
				Severity: "high",
				Cvss: []Cvss{
					{
						Metrics: CvssMetrics{
							BaseScore: 7.5,
						},
					},
				},
			},
			expected: 0, // threat is 0 without EPSS or KEV
		},
		{
			name: "with EPSS data, no KEV",
			metadata: Metadata{
				Severity: "high",
				EPSS: []EPSS{
					{
						EPSS:       0.5,
						Percentile: 0.95,
					},
				},
				Cvss: []Cvss{
					{
						Metrics: CvssMetrics{
							BaseScore: 7.5,
						},
					},
				},
			},
			expected: 37.5, // 0.5 * (7.5/10) * 1 * 100
		},
		{
			name: "with KEV, no EPSS",
			metadata: Metadata{
				Severity: "high",
				KnownExploited: []KnownExploited{
					{
						CVE:                        "CVE-2023-1234",
						KnownRansomwareCampaignUse: "No",
					},
				},
				Cvss: []Cvss{
					{
						Metrics: CvssMetrics{
							BaseScore: 7.5,
						},
					},
				},
			},
			expected: 78.75, // 1.0 * (7.5/10) * 1.05* 100
		},
		{
			name: "with KEV ransomware",
			metadata: Metadata{
				Severity: "high",
				KnownExploited: []KnownExploited{
					{
						CVE:                        "CVE-2023-1234",
						KnownRansomwareCampaignUse: "Known",
					},
				},
				Cvss: []Cvss{
					{
						Metrics: CvssMetrics{
							BaseScore: 7.5,
						},
					},
				},
			},
			expected: 82.5, // 1.0 * (7.5/10) * 1.1 * 100
		},
		{
			name: "with severity string only",
			metadata: Metadata{
				Severity: "critical",
				EPSS: []EPSS{
					{
						EPSS:       0.8,
						Percentile: 0.99,
					},
				},
			},
			expected: 72, // 0.8 * (9.0/10) * 1.0 * 100
		},
		{
			name: "with multiple CVSS scores + string severity",
			metadata: Metadata{
				Severity: "medium",
				EPSS: []EPSS{
					{
						EPSS:       0.6,
						Percentile: 0.90,
					},
				},
				Cvss: []Cvss{
					{
						Source: "NVD",
						Metrics: CvssMetrics{
							BaseScore: 6.5,
						},
					},
					{
						Source: "Vendor",
						Metrics: CvssMetrics{
							BaseScore: 5.5,
						},
					},
				},
			},
			expected: 33, // 0.6 * ( (((6.5+5.5)/2)+5)/2 /10) * 1.0 * 100
		},
		{
			name: "with some invalid CVSS scores + string severity",
			metadata: Metadata{
				Severity: "medium",
				EPSS: []EPSS{
					{
						EPSS:       0.4,
						Percentile: 0.85,
					},
				},
				Cvss: []Cvss{
					{
						Source: "NVD",
						Metrics: CvssMetrics{
							BaseScore: 0, // invalid, should be ignored
						},
					},
					{
						Source: "Vendor",
						Metrics: CvssMetrics{
							BaseScore: 6.0,
						},
					},
				},
			},
			expected: 22, // 0.4 * ((6.0+5)/2 /10) * 1.0 * 100
		},
		{
			name: "unknown severity",
			metadata: Metadata{
				Severity: "unknown",
				EPSS: []EPSS{
					{
						EPSS:       0.3,
						Percentile: 0.80,
					},
				},
			},
			expected: 15, // 0.3 * (5.0/10) * 1.0 * 100
		},
		{
			name: "maximum risk clamp",
			metadata: Metadata{
				Severity: "critical",
				KnownExploited: []KnownExploited{
					{
						CVE:                        "CVE-2023-1234",
						KnownRansomwareCampaignUse: "Known",
					},
				},
				Cvss: []Cvss{
					{
						Metrics: CvssMetrics{
							BaseScore: 10.0,
						},
					},
				},
			},
			expected: 100, // clamped to 100 as it would be 1.0 * 1.0 * 1.1 * 100 = 120
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := tt.metadata.RiskScore()
			assert.InDelta(t, tt.expected, result, 0.01, "RiskScore method returned incorrect value")

			// test the calculated value is cached
			if tt.name != "already calculated risk" && tt.name != "nil metadata" {
				require.InDelta(t, tt.expected, tt.metadata.risk, 0.01, "risk was not cached")
			}

			// test the standalone function
			if tt.name != "nil metadata" && tt.name != "already calculated risk" {
				funcResult := riskScore(tt.metadata)
				assert.InDelta(t, tt.expected, funcResult, 0.0001, "riskScore function returned incorrect value")
			}
		})
	}
}

func TestSeverityToScore(t *testing.T) {
	tests := []struct {
		severity string
		expected float64
	}{
		{"negligible", 0.5},
		{"NEGLIGIBLE", 0.5},
		{"low", 3.0},
		{"LOW", 3.0},
		{"medium", 5.0},
		{"MEDIUM", 5.0},
		{"high", 7.5},
		{"HIGH", 7.5},
		{"critical", 9.0},
		{"CRITICAL", 9.0},
		{"unknown", 5.0},
		{"", 5.0},
		{"something-else", 5.0},
	}

	for _, tt := range tests {
		t.Run(tt.severity, func(t *testing.T) {
			result := severityToScore(tt.severity)
			assert.Equal(t, tt.expected, result)
		})
	}
}

func TestAverageCVSS(t *testing.T) {
	tests := []struct {
		name     string
		cvss     []Cvss
		expected float64
	}{
		{
			name:     "empty slice",
			cvss:     []Cvss{},
			expected: 0,
		},
		{
			name: "single valid score",
			cvss: []Cvss{
				{Metrics: CvssMetrics{BaseScore: 7.5}},
			},
			expected: 7.5,
		},
		{
			name: "multiple valid scores",
			cvss: []Cvss{
				{Metrics: CvssMetrics{BaseScore: 7.5}},
				{Metrics: CvssMetrics{BaseScore: 8.5}},
				{Metrics: CvssMetrics{BaseScore: 9.0}},
			},
			expected: 8.33333,
		},
		{
			name: "with invalid scores",
			cvss: []Cvss{
				{Metrics: CvssMetrics{BaseScore: 0}}, // invalid
				{Metrics: CvssMetrics{BaseScore: 7.5}},
				{Metrics: CvssMetrics{BaseScore: 0}}, // invalid
				{Metrics: CvssMetrics{BaseScore: 8.5}},
			},
			expected: 8.0,
		},
		{
			name: "all invalid scores",
			cvss: []Cvss{
				{Metrics: CvssMetrics{BaseScore: 0}},
				{Metrics: CvssMetrics{BaseScore: 0}},
			},
			expected: 0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := average(validBaseScores(tt.cvss...)...)
			assert.InDelta(t, tt.expected, result, 0.00001)
		})
	}
}

func TestThreat(t *testing.T) {
	tests := []struct {
		name     string
		metadata Metadata
		expected float64
	}{
		{
			name:     "no EPSS, no KEV",
			metadata: Metadata{},
			expected: 0,
		},
		{
			name: "with EPSS, no KEV",
			metadata: Metadata{
				EPSS: []EPSS{
					{EPSS: 0.75},
				},
			},
			expected: 0.75,
		},
		{
			name: "with KEV, no EPSS",
			metadata: Metadata{
				KnownExploited: []KnownExploited{
					{CVE: "CVE-2023-1234"},
				},
			},
			expected: 1.0,
		},
		{
			name: "with KEV and EPSS",
			metadata: Metadata{
				EPSS: []EPSS{
					{EPSS: 0.5},
				},
				KnownExploited: []KnownExploited{
					{CVE: "CVE-2023-1234"},
				},
			},
			expected: 1.0, // KEV takes precedence
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := threat(tt.metadata)
			assert.Equal(t, tt.expected, result)
		})
	}
}

func TestImpact(t *testing.T) {
	tests := []struct {
		name     string
		metadata Metadata
		expected float64
	}{
		{
			name:     "no KEV",
			metadata: Metadata{},
			expected: 1.0,
		},
		{
			name: "KEV without ransomware",
			metadata: Metadata{
				KnownExploited: []KnownExploited{
					{KnownRansomwareCampaignUse: "No"},
				},
			},
			expected: 1.05,
		},
		{
			name: "KEV with ransomware",
			metadata: Metadata{
				KnownExploited: []KnownExploited{
					{KnownRansomwareCampaignUse: "Known"},
				},
			},
			expected: 1.1,
		},
		{
			name: "KEV with case insensitive ransomware",
			metadata: Metadata{
				KnownExploited: []KnownExploited{
					{KnownRansomwareCampaignUse: "KNOWN"},
				},
			},
			expected: 1.1,
		},
		{
			name: "multiple KEV entries, one with ransomware",
			metadata: Metadata{
				KnownExploited: []KnownExploited{
					{KnownRansomwareCampaignUse: "No"},
					{KnownRansomwareCampaignUse: "Known"},
				},
			},
			expected: 1.1, // highest wins
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := kevModifier(tt.metadata)
			assert.Equal(t, tt.expected, result)
		})
	}
}

func TestSeverity(t *testing.T) {
	tests := []struct {
		name     string
		metadata Metadata
		expected float64
	}{
		{
			name: "no CVSS, medium severity",
			metadata: Metadata{
				Severity: "medium",
			},
			expected: 0.5,
		},
		{
			name: "with CVSS + severity string",
			metadata: Metadata{
				Severity: "medium",
				Cvss: []Cvss{
					{Metrics: CvssMetrics{BaseScore: 8.0}},
				},
			},
			expected: 0.65,
		},
		{
			name: "multiple CVSS scores + severity string",
			metadata: Metadata{
				Severity: "medium",
				Cvss: []Cvss{
					{Metrics: CvssMetrics{BaseScore: 6.0}},
					{Metrics: CvssMetrics{BaseScore: 8.0}},
				},
			},
			expected: 0.6,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := severity(tt.metadata)
			assert.InDelta(t, tt.expected, result, 0.00001)
		})
	}
}
