package table

import (
	"bytes"
	"fmt"
	"path/filepath"
	"strings"
	"sync"

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

	"github.com/aquasecurity/table"
	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"
)

type vulnerabilityRenderer struct {
	w           *bytes.Buffer
	tableWriter *table.Table
	result      types.Result
	isTerminal  bool
	tree        bool
	severities  []dbTypes.Severity
	once        *sync.Once
}

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

func (r vulnerabilityRenderer) Render() string {
	r.setHeaders()
	r.setVulnerabilityRows(r.result.Vulnerabilities)

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

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

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

	return r.w.String()
}

func (r vulnerabilityRenderer) setHeaders() {
	header := []string{"Library", "Vulnerability", "Severity", "Installed Version", "Fixed Version", "Title"}
	r.tableWriter.SetHeaders(header...)
}

func (r vulnerabilityRenderer) setVulnerabilityRows(vulns []types.DetectedVulnerability) {
	for _, v := range vulns {
		lib := v.PkgName
		if v.PkgPath != "" {
			fileName := filepath.Base(v.PkgPath)
			lib = fmt.Sprintf("%s (%s)", v.PkgName, fileName)
			r.once.Do(func() {
				log.Logger.Infof("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 len(v.PrimaryURL) > 0 {
			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.InstalledVersion, v.FixedVersion, strings.TrimSpace(title),
			}
		} else {
			row = []string{
				lib, v.VulnerabilityID, v.Severity, v.InstalledVersion, v.FixedVersion, strings.TrimSpace(title),
			}
		}

		r.tableWriter.AddRow(row...)
	}
}

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

func (r vulnerabilityRenderer) renderDependencyTree() {
	// Get parents of each dependency
	parents := reverseDeps(r.result.Packages)
	if len(parents) == 0 {
		return
	}

	root := treeprint.NewWithRoot(fmt.Sprintf(`
Dependency Origin Tree
======================
%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 := map[string]map[string]int{}
	for _, vuln := range r.result.Vulnerabilities {
		cnts, ok := pkgSeverityCount[vuln.PkgID]
		if !ok {
			cnts = map[string]int{}
		}

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

	// Render tree
	seen := map[string]struct{}{}
	for _, vuln := range r.result.Vulnerabilities {
		if _, ok := seen[vuln.PkgID]; ok {
			continue
		}

		_, summaries := summarize(r.severities, pkgSeverityCount[vuln.PkgID])
		topLvlID := fmt.Sprintf("%s, (%s)", vuln.PkgID, strings.Join(summaries, ", "))
		if r.isTerminal {
			topLvlID = color.HiRedString(topLvlID)
		}

		seen[vuln.PkgID] = struct{}{}
		branch := root.AddBranch(topLvlID)
		addParents(branch, vuln.PkgID, parents)

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

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

func addParents(topItem treeprint.Tree, pkgID string, parentMap map[string][]string) {
	parents, ok := parentMap[pkgID]
	if !ok {
		return
	}
	for _, parent := range parents {
		branch := topItem.AddBranch(parent)
		addParents(branch, parent, parentMap)
	}
}

func reverseDeps(libs []ftypes.Package) map[string][]string {
	reversed := make(map[string][]string)
	for _, lib := range libs {
		for _, dependOn := range lib.DependsOn {
			items, ok := reversed[dependOn]
			if !ok {
				reversed[dependOn] = []string{lib.ID}
			} else {
				reversed[dependOn] = append(items, lib.ID)
			}
		}
	}

	for k, v := range reversed {
		reversed[k] = lo.Uniq(v)
	}
	return reversed
}
