// ©AngelaMos | 2026
// vuln.go

package cli

import (
	"context"
	"os"
	"time"

	"github.com/spf13/cobra"

	"github.com/CarterPerez-dev/bomber/internal/graph"
	"github.com/CarterPerez-dev/bomber/internal/parser"
	"github.com/CarterPerez-dev/bomber/internal/report"
	"github.com/CarterPerez-dev/bomber/internal/scanner"
	"github.com/CarterPerez-dev/bomber/internal/ui"
	"github.com/CarterPerez-dev/bomber/internal/vuln"
	"github.com/CarterPerez-dev/bomber/pkg/types"
)

var vulnCmd = &cobra.Command{
	Use:   "vuln [path]",
	Short: "Scan for known vulnerabilities",
	Args:  cobra.MaximumNArgs(1),
	RunE:  runVuln,
}

func runVuln(cmd *cobra.Command, args []string) error {
	path := "."
	if len(args) > 0 {
		path = args[0]
	}

	reg := parser.NewRegistry()
	parser.RegisterAll(reg)

	s := scanner.New(reg)
	result, err := s.Scan(path)
	if err != nil {
		return err
	}

	if len(result.Graphs) == 0 {
		if formatFlag == "json" {
			return report.WriteJSON(os.Stdout, result, nil, nil)
		}
		ui.PrintBanner()
		report.PrintScanSummary(os.Stdout, result)
		return nil
	}

	sp := ui.NewSpinner("Querying vulnerability databases...")
	if formatFlag != "json" {
		sp.Start()
	}

	vulnReport, err := queryVulns(cmd.Context(), result)

	if formatFlag != "json" {
		sp.Stop()
	}
	if err != nil {
		return err
	}

	if formatFlag == "json" {
		return report.WriteJSON(os.Stdout, result, vulnReport, nil)
	}

	ui.PrintBanner()
	report.PrintScanSummary(os.Stdout, result)
	report.PrintVulnReport(os.Stdout, vulnReport)
	return nil
}

func queryVulns(ctx context.Context, result *types.ScanResult) (*types.VulnReport, error) {
	allPkgs := make([]types.Package, 0, len(result.Graphs))
	for _, g := range result.Graphs {
		allPkgs = append(allPkgs, graph.AllPackages(g)...)
	}

	var cache *vuln.Cache
	if !noCache {
		c, err := vuln.NewCache(vuln.DefaultCachePath(), 24*time.Hour)
		if err == nil {
			cache = c
			defer func() { _ = cache.Close() }()
		}
	}

	osvClient := vuln.NewOSVClient()
	var clients []vuln.Client
	clients = append(clients, osvClient)

	nvdKey := os.Getenv("BOMBER_NVD_API_KEY")
	if nvdKey != "" {
		clients = append(clients, vuln.NewNVDClient(vuln.WithNVDAPIKey(nvdKey)))
	}

	vulnReport := &types.VulnReport{
		TotalPkgs:  result.TotalPkgs,
		DirectPkgs: result.DirectPkgs,
		BySeverity: make(map[types.Severity]int),
	}

	for _, client := range clients {
		var uncached []types.Package
		for _, pkg := range allPkgs {
			if cache != nil {
				cached, ok, err := cache.Get(pkg.PURL, client.Source())
				if err == nil && ok {
					vulnReport.Matches = append(vulnReport.Matches, cached...)
					for _, m := range cached {
						vulnReport.BySeverity[m.Vulnerability.Severity]++
					}
					continue
				}
			}
			uncached = append(uncached, pkg)
		}

		if len(uncached) == 0 {
			continue
		}

		matches, err := client.Query(ctx, uncached)
		if err != nil {
			continue
		}

		matchesByPURL := make(map[string][]types.VulnMatch)
		for _, m := range matches {
			vulnReport.Matches = append(vulnReport.Matches, m)
			vulnReport.BySeverity[m.Vulnerability.Severity]++
			matchesByPURL[m.Package.PURL] = append(matchesByPURL[m.Package.PURL], m)
		}

		if cache != nil {
			for purl, pkgMatches := range matchesByPURL {
				_ = cache.Put(purl, client.Source(), pkgMatches)
			}
		}
	}

	vulnReport.Matches = deduplicateMatches(vulnReport.Matches)
	vulnReport.BySeverity = make(map[types.Severity]int)
	for _, m := range vulnReport.Matches {
		vulnReport.BySeverity[m.Vulnerability.Severity]++
	}

	return vulnReport, nil
}

func deduplicateMatches(matches []types.VulnMatch) []types.VulnMatch {
	seen := make(map[string]int)
	var deduped []types.VulnMatch

	for _, m := range matches {
		ids := make([]string, 0, 1+len(m.Vulnerability.Aliases))
		ids = append(ids, m.Vulnerability.ID)
		ids = append(ids, m.Vulnerability.Aliases...)

		existingIdx := -1
		for _, id := range ids {
			if idx, ok := seen[id]; ok {
				existingIdx = idx
				break
			}
		}

		if existingIdx >= 0 {
			existing := deduped[existingIdx]
			if m.Vulnerability.Score > existing.Vulnerability.Score ||
				(m.Vulnerability.FixVersion != "" && existing.Vulnerability.FixVersion == "") {
				deduped[existingIdx] = m
			}
			continue
		}

		idx := len(deduped)
		for _, id := range ids {
			seen[id] = idx
		}
		deduped = append(deduped, m)
	}

	return deduped
}