package ci

import (
	"github.com/google/osv-scanner/v2/internal/grouper"
	"github.com/google/osv-scanner/v2/internal/output"
	"github.com/google/osv-scanner/v2/pkg/models"
)

// DiffVulnerabilityResults will return any new vulnerabilities that are in `newRes`
// which is not present in `oldRes`, but not the reverse.
func DiffVulnerabilityResults(oldRes, newRes models.VulnerabilityResults) models.VulnerabilityResults {
	result := models.VulnerabilityResults{}
	// Initialize caches for quick lookup
	sourceToIndex, packageToIndex, vulnToIndex := initializeCaches(oldRes)

	for _, ps := range newRes.Results {
		sourceIdx, sourceExists := sourceToIndex[ps.Source]
		if !sourceExists {
			// Newly introduced source, so all results for this source are going to be new, add everything for this source
			result.Results = append(result.Results, ps)
			continue
		}
		// Otherwise the old source used to exist, so we need to find the difference in the packages
		result.Results = append(result.Results, models.PackageSource{
			Source: ps.Source,
		})
		resultPS := &result.Results[len(result.Results)-1]
		for _, pv := range ps.Packages {
			pkgIdx, packageExists := packageToIndex[sourceIdx][pv.Package]
			if !packageExists {
				// Newly introduced package, so all results for this package are going to be new, add everything for this package
				resultPS.Packages = append(resultPS.Packages, pv)
				continue
			}
			// Otherwise the old package used to exist, so we need to find the difference in the vulnerabilities
			// Only copy over packages as vulns and groups might change
			resultPS.Packages = append(resultPS.Packages, models.PackageVulns{
				Package: pv.Package,
			})
			resultPV := &resultPS.Packages[len(resultPS.Packages)-1]
			for _, v := range pv.Vulnerabilities {
				if !vulnToIndex[sourceIdx][pkgIdx][v.GetId()] {
					// Vulnerability is new, add it to the results
					resultPV.Vulnerabilities = append(resultPV.Vulnerabilities, v)
					continue
				}
			}
			if len(resultPV.Vulnerabilities) == 0 {
				// No vulns, so we can remove the PackageVulns entry entirely, and skip grouping
				resultPS.Packages = resultPS.Packages[:len(resultPS.Packages)-1]
				continue
			}
			// Rebuild the groups lost in the previous step
			groups := grouper.Group(grouper.ConvertVulnerabilityToIDAliases(resultPV.Vulnerabilities))
			for i, group := range groups {
				groups[i].MaxSeverity = output.MaxSeverity(group, *resultPV)
			}
			resultPV.Groups = groups
		}
		if len(resultPS.Packages) == 0 {
			// No packages, so we can remove the PackageSource entry entirely
			result.Results = result.Results[:len(result.Results)-1]
			continue
		}
	}

	return result
}

// initializeCaches sets up maps for quick lookup of sources, packages, and vulnerabilities by their indices.
func initializeCaches(oldRes models.VulnerabilityResults) (map[models.SourceInfo]int, []map[models.PackageInfo]int, [][]map[string]bool) {
	sourceToIndex := make(map[models.SourceInfo]int, len(oldRes.Results))
	// The index in the array corresponds to a source index, a query would look like packageToIndex[sourceIndex][packageInfo]
	packageToIndex := make([]map[models.PackageInfo]int, len(oldRes.Results))
	// The first index in the array corresponds to a source index, and the second index corresponds to a package index
	// a query would look like vulnToIndex[sourceIndex][packageIndex][vulnID]
	vulnToIndex := make([][]map[string]bool, len(oldRes.Results))

	// Populate index maps for sources, packages, and vulnerabilities
	for sourceIndex, vulnResult := range oldRes.Results {
		sourceToIndex[oldRes.Results[sourceIndex].Source] = sourceIndex
		if vulnToIndex[sourceIndex] == nil {
			vulnToIndex[sourceIndex] = make([]map[string]bool, len(vulnResult.Packages))
		}
		for packageIndex, pkg := range vulnResult.Packages {
			if packageToIndex[sourceIndex] == nil {
				packageToIndex[sourceIndex] = make(map[models.PackageInfo]int, len(vulnResult.Packages))
			}
			packageToIndex[sourceIndex][pkg.Package] = packageIndex
			if vulnToIndex[sourceIndex][packageIndex] == nil {
				vulnToIndex[sourceIndex][packageIndex] = make(map[string]bool, len(pkg.Vulnerabilities))
			}
			for _, vuln := range pkg.Vulnerabilities {
				vulnToIndex[sourceIndex][packageIndex][vuln.GetId()] = true // Mark the vulnerability as present
			}
		}
	}

	return sourceToIndex, packageToIndex, vulnToIndex
}

// DiffVulnerabilityResultsByOccurrences will return the occurrence of each vulnerability that are in `newRes`
// which is not present in `oldRes`, but not the reverse. This calculates the difference by vulnerability ID,
// while ignoring the source of the vulnerability.
//
// This prevents us reporting "new" vulnerabilities in a PR when a previously vulnerable file is being moved.
func DiffVulnerabilityResultsByOccurrences(oldRes, newRes models.VulnerabilityResults) map[string]int {
	oldResFlat := oldRes.Flatten()
	newResFlat := newRes.Flatten()

	oldResMap := map[string]int{}
	newResMap := map[string]int{}

	for _, vf := range oldResFlat {
		oldResMap[vf.Vulnerability.GetId()] += 1
	}

	for _, vf := range newResFlat {
		newResMap[vf.Vulnerability.GetId()] += 1
	}

	for k, oldVulnCount := range oldResMap {
		// If the new result has fewer vulnerabilities than the old result remove the entry from the new result.
		// `map`'s default value is 0 when empty, and delete also works fine when the entry is empty
		if newResMap[k] <= oldVulnCount {
			delete(newResMap, k)
		}
	}

	return newResMap
}