package table

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"slices"
	"sort"
	"strings"
	"sync"

	"github.com/fatih/color"
	"github.com/samber/lo"
	"github.com/xlab/treeprint"

	"github.com/aquasecurity/table"
	"github.com/aquasecurity/tml"
	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
	ftypes "github.com/aquasecurity/trivy/pkg/fanal/types"
	"github.com/aquasecurity/trivy/pkg/log"
	"github.com/aquasecurity/trivy/pkg/types"
	"github.com/aquasecurity/trivy/pkg/version/doc"
)

const (
	vexNotice = `
For OSS Maintainers: VEX Notice
--------------------------------
If you're an OSS maintainer and Trivy has detected vulnerabilities in your project that you believe are not actually exploitable, consider issuing a VEX (Vulnerability Exploitability eXchange) statement.
VEX allows you to communicate the actual status of vulnerabilities in your project, improving security transparency and reducing false positives for your users.
Learn more and start using VEX: %s

To disable this notice, set the TRIVY_DISABLE_VEX_NOTICE environment variable.

`
	envDisableNotice = "TRIVY_DISABLE_VEX_NOTICE"
)

var (
	showVEXNoticeOnce  = &sync.Once{}
	showSuppressedOnce = sync.OnceFunc(func() {
		log.Info(`Some vulnerabilities have been ignored/suppressed. Use the "--show-suppressed" flag to display them.`)
	})
)

type vulnerabilityRenderer struct {
	w              *bytes.Buffer
	result         types.Result
	isTerminal     bool
	tree           bool // Show dependency tree
	showSuppressed bool // Show suppressed vulnerabilities
	severities     []dbTypes.Severity
	once           *sync.Once
}

func NewVulnerabilityRenderer(result types.Result, isTerminal, tree, suppressed bool, severities []dbTypes.Severity) *vulnerabilityRenderer {
	buf := bytes.NewBuffer([]byte{})
	if !isTerminal {
		tml.DisableFormatting()
	}
	return &vulnerabilityRenderer{
		w:              buf,
		result:         result,
		isTerminal:     isTerminal,
		tree:           tree,
		showSuppressed: suppressed,
		severities:     severities,
		once:           new(sync.Once),
	}
}

func (r *vulnerabilityRenderer) Render() string {
	// There are 3 cases when we show the vulnerability table (or only target and `Total: 0...`):
	// When Result contains vulnerabilities;
	// When Result target is OS packages even if no vulnerabilities are found;
	// When we show non-empty `Suppressed Vulnerabilities` table.
	if len(r.result.Vulnerabilities) > 0 || r.result.Class == types.ClassOSPkg || (r.showSuppressed && len(r.result.ModifiedFindings) > 0) {
		r.renderDetectedVulnerabilities()

		if r.tree {
			r.renderDependencyTree()
		}
	}

	if r.showSuppressed {
		r.renderModifiedVulnerabilities()
	} else if len(r.result.ModifiedFindings) > 0 {
		showSuppressedOnce()
	}

	return r.w.String()
}

func (r *vulnerabilityRenderer) renderDetectedVulnerabilities() {
	// Show VEX notice only on CI
	showVEXNoticeOnce.Do(func() {
		if os.Getenv(envDisableNotice) != "" || os.Getenv("CI") == "" {
			return
		}
		_, _ = color.New(color.FgCyan).Fprintf(r.w, vexNotice, doc.URL("docs/supply-chain/vex/repo", "publishing-vex-documents"))
	})

	tw := newTableWriter(r.w, r.isTerminal)
	r.setHeaders(tw)
	r.setVulnerabilityRows(tw, r.result.Vulnerabilities)

	severityCount := r.countSeverities(r.result.Vulnerabilities)
	total, summaries := summarize(r.severities, severityCount)

	target := r.result.Target
	if r.result.Class == types.ClassLangPkg {
		target += fmt.Sprintf(" (%s)", r.result.Type)
	}
	RenderTarget(r.w, target, r.isTerminal)
	r.printf("Total: %d (%s)\n\n", total, strings.Join(summaries, ", "))

	tw.Render()
}

func (r *vulnerabilityRenderer) setHeaders(tw *table.Table) {
	if len(r.result.Vulnerabilities) == 0 {
		return
	}
	header := []string{
		"Library",
		"Vulnerability",
		"Severity",
		"Status",
		"Installed Version",
		"Fixed Version",
		"Title",
	}
	tw.SetHeaders(header...)
}

func (r *vulnerabilityRenderer) setVulnerabilityRows(tw *table.Table, vulns []types.DetectedVulnerability) {
	for _, v := range vulns {
		lib := v.PkgName
		if v.PkgPath != "" {
			// get path to root jar
			// for other languages return unchanged path
			pkgPath := rootJarFromPath(v.PkgPath)
			fileName := filepath.Base(pkgPath)
			lib = fmt.Sprintf("%s (%s)", v.PkgName, fileName)
			r.once.Do(func() {
				log.Info("Table result includes only package filenames. Use '--format json' option to get the full path to the package file.")
			})
		}

		title := v.Title
		if title == "" {
			title = v.Description
		}
		splitTitle := strings.Split(title, " ")
		if len(splitTitle) >= 12 {
			title = strings.Join(splitTitle[:12], " ") + "..."
		}

		if v.PrimaryURL != "" {
			if r.isTerminal {
				title = tml.Sprintf("%s\n<blue>%s</blue>", title, v.PrimaryURL)
			} else {
				title = fmt.Sprintf("%s\n%s", title, v.PrimaryURL)
			}
		}

		var row []string
		if r.isTerminal {
			row = []string{
				lib,
				v.VulnerabilityID,
				ColorizeSeverity(v.Severity, v.Severity),
				v.Status.String(),
				v.InstalledVersion,
				v.FixedVersion,
				strings.TrimSpace(title),
			}
		} else {
			row = []string{
				lib,
				v.VulnerabilityID,
				v.Severity,
				v.Status.String(),
				v.InstalledVersion,
				v.FixedVersion,
				strings.TrimSpace(title),
			}
		}

		tw.AddRow(row...)
	}
}

func (r *vulnerabilityRenderer) countSeverities(vulns []types.DetectedVulnerability) map[string]int {
	severityCount := make(map[string]int)
	for _, v := range vulns {
		severityCount[v.Severity]++
	}
	return severityCount
}

func (r *vulnerabilityRenderer) renderModifiedVulnerabilities() {
	tw := newTableWriter(r.w, r.isTerminal)
	header := []string{
		"Library",
		"Vulnerability",
		"Severity",
		"Status",
		"Statement",
		"Source",
	}
	tw.SetHeaders(header...)

	var total int
	for _, m := range r.result.ModifiedFindings {
		if m.Type != types.FindingTypeVulnerability {
			continue
		}
		vuln := m.Finding.(types.DetectedVulnerability)
		total++

		stmt := lo.Ternary(m.Statement != "", m.Statement, "N/A")
		tw.AddRow(vuln.PkgName, vuln.VulnerabilityID, vuln.Severity, string(m.Status), stmt, m.Source)
	}

	if total == 0 {
		return
	}

	title := fmt.Sprintf("Suppressed Vulnerabilities (Total: %d)", total)
	if r.isTerminal {
		// nolint
		_ = tml.Fprintf(r.w, "\n<underline>%s</underline>\n\n", title)
	} else {
		_, _ = fmt.Fprintf(r.w, "\n%s\n", title)
		_, _ = fmt.Fprintf(r.w, "%s\n", strings.Repeat("=", len(title)))
	}

	tw.Render()
}

func (r *vulnerabilityRenderer) renderDependencyTree() {
	// Get parents of each dependency
	parents := ftypes.Packages(r.result.Packages).ParentDeps()
	if len(parents) == 0 {
		return
	}
	ancestors := traverseAncestors(r.result.Packages, parents)

	root := treeprint.NewWithRoot(fmt.Sprintf(`
Dependency Origin Tree (Reversed)
=================================
%s`, r.result.Target))

	// This count is next to the package ID.
	// e.g. node-fetch@1.7.3 (MEDIUM: 2, HIGH: 1, CRITICAL: 3)
	pkgSeverityCount := make(map[string]map[string]int)
	for _, vuln := range r.result.Vulnerabilities {
		cnts, ok := pkgSeverityCount[vuln.PkgID]
		if !ok {
			cnts = make(map[string]int)
		}

		cnts[vuln.Severity]++
		pkgSeverityCount[vuln.PkgID] = cnts
	}

	// Extract vulnerable packages
	vulnPkgs := lo.Filter(r.result.Packages, func(pkg ftypes.Package, _ int) bool {
		return lo.ContainsBy(r.result.Vulnerabilities, func(vuln types.DetectedVulnerability) bool {
			return pkg.ID == vuln.PkgID
		})
	})

	// Render tree
	for _, vulnPkg := range vulnPkgs {
		_, summaries := summarize(r.severities, pkgSeverityCount[vulnPkg.ID])
		topLvlID := tml.Sprintf("<red>%s, (%s)</red>", vulnPkg.ID, strings.Join(summaries, ", "))

		branch := root.AddBranch(topLvlID)
		addParents(branch, vulnPkg, parents, ancestors, map[string]struct{}{vulnPkg.ID: {}}, 1)

	}
	r.printf(root.String())
}

func (r *vulnerabilityRenderer) printf(format string, args ...any) {
	// nolint
	_ = tml.Fprintf(r.w, format, args...)
}

func addParents(topItem treeprint.Tree, pkg ftypes.Package, parentMap map[string]ftypes.Packages, ancestors map[string][]string,
	seen map[string]struct{}, depth int) {
	if pkg.Relationship == ftypes.RelationshipDirect {
		return
	}

	roots := make(map[string]struct{})
	for _, parent := range parentMap[pkg.ID] {
		if _, ok := seen[parent.ID]; ok {
			continue
		}
		seen[parent.ID] = struct{}{} // to avoid infinite loops

		if depth == 1 && parent.Relationship == ftypes.RelationshipDirect {
			topItem.AddBranch(parent.ID)
		} else {
			// We omit intermediate dependencies and show only direct dependencies
			// as this could make the dependency tree huge.
			for _, ancestor := range ancestors[parent.ID] {
				roots[ancestor] = struct{}{}
			}
		}
	}

	// Omitted
	rootIDs := lo.Filter(lo.Keys(roots), func(pkgID string, _ int) bool {
		_, ok := seen[pkgID]
		return !ok
	})
	sort.Strings(rootIDs)
	if len(rootIDs) > 0 {
		branch := topItem.AddBranch("...(omitted)...")
		for _, rootID := range rootIDs {
			branch.AddBranch(rootID)
		}
	}
}

func traverseAncestors(pkgs []ftypes.Package, parentMap map[string]ftypes.Packages) map[string][]string {
	ancestors := make(map[string][]string)
	for _, pkg := range pkgs {
		ancestors[pkg.ID] = findAncestor(pkg.ID, parentMap, make(map[string]struct{}))
	}
	return ancestors
}

func findAncestor(pkgID string, parentMap map[string]ftypes.Packages, seen map[string]struct{}) []string {
	ancestors := make(map[string]struct{})
	seen[pkgID] = struct{}{}
	for _, parent := range parentMap[pkgID] {
		if _, ok := seen[parent.ID]; ok {
			continue
		}
		switch {
		case parent.Relationship == ftypes.RelationshipDirect:
			ancestors[parent.ID] = struct{}{}
		case len(parentMap[parent.ID]) == 0:
			// Some package managers, such as "package-lock.json" v1, can retrieve package dependencies but not relationships.
			// We try to guess direct dependencies in this case. A dependency with no parents must be a direct dependency.
			//
			// e.g.
			//   -> styled-components
			//     -> fbjs
			//       -> isomorphic-fetch
			//         -> node-fetch
			//
			// Even if `styled-components` is not marked as a direct dependency, it must be a direct dependency
			// as it has no parents. Note that it doesn't mean `fbjs` is an indirect dependency.
			ancestors[parent.ID] = struct{}{}
		default:
			for _, a := range findAncestor(parent.ID, parentMap, seen) {
				ancestors[a] = struct{}{}
			}
		}
	}
	return lo.Keys(ancestors)
}

var jarExtensions = []string{
	".jar",
	".war",
	".par",
	".ear",
}

func rootJarFromPath(path string) string {
	// File paths are always forward-slashed in Trivy
	paths := strings.Split(path, "/")
	for i, p := range paths {
		if slices.Contains(jarExtensions, filepath.Ext(p)) {
			return strings.Join(paths[:i+1], "/")
		}
	}
	return path
}
