package checker

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"sort"
	"time"

	sdk "git.happydns.org/checker-sdk-go/checker"
)

const virusTotalEndpoint = "https://www.virustotal.com/api/v3/domains/"

func init() { Register(&virusTotalSource{endpoint: virusTotalEndpoint}) }

type virusTotalSource struct {
	endpoint string
}

func (*virusTotalSource) ID() string   { return "virustotal" }
func (*virusTotalSource) Name() string { return "VirusTotal" }

func (*virusTotalSource) Options() SourceOptions {
	return SourceOptions{
		Admin: []sdk.CheckerOptionField{
			{
				Id:          "virustotal_api_key",
				Type:        "string",
				Label:       "VirusTotal API key",
				Description: "VirusTotal v3 API key. Free tier is limited to 4 req/min and 500 req/day.",
				Secret:      true,
			},
		},
	}
}

// vtDetails persists the structured VT response so the rich detail
// renderer can show the per-vendor verdict table and the
// {malicious,suspicious,harmless,undetected} counts.
type vtDetails struct {
	Malicious  int               `json:"malicious"`
	Suspicious int               `json:"suspicious"`
	Harmless   int               `json:"harmless"`
	Undetected int               `json:"undetected"`
	Total      int               `json:"total"`
	Reputation int               `json:"reputation"`
	Vendors    []vtVendorVerdict `json:"vendors"`
}

type vtVendorVerdict struct {
	Engine   string `json:"engine"`
	Category string `json:"category"`
	Result   string `json:"result"`
}

func (s *virusTotalSource) Query(ctx context.Context, domain, registered string, opts sdk.CheckerOptions) []SourceResult {
	apiKey := stringOpt(opts, "virustotal_api_key")
	if apiKey == "" {
		return []SourceResult{{SourceID: s.ID(), SourceName: s.Name(), Enabled: false}}
	}
	if registered == "" {
		return []SourceResult{{SourceID: s.ID(), SourceName: s.Name(), Enabled: true}}
	}

	res := SourceResult{
		SourceID: s.ID(), SourceName: s.Name(), Enabled: true,
		Reference: "https://www.virustotal.com/gui/domain/" + registered,
	}

	reqCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
	defer cancel()

	req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, s.endpoint+registered, nil)
	if err != nil {
		res.Error = err.Error()
		return []SourceResult{res}
	}
	req.Header.Set("x-apikey", apiKey)
	req.Header.Set("Accept", "application/json")

	body, status, err := httpDo(req, 4<<20)
	if err != nil {
		res.Error = err.Error()
		return []SourceResult{res}
	}
	if status == http.StatusNotFound {
		// VT has never seen this domain → quiet "not listed".
		return []SourceResult{res}
	}
	if status != http.StatusOK {
		res.Error = fmt.Sprintf("HTTP %d: %s", status, truncate(string(body), 200))
		return []SourceResult{res}
	}

	var parsed struct {
		Data struct {
			Attributes struct {
				LastAnalysisStats struct {
					Harmless   int `json:"harmless"`
					Malicious  int `json:"malicious"`
					Suspicious int `json:"suspicious"`
					Undetected int `json:"undetected"`
					Timeout    int `json:"timeout"`
				} `json:"last_analysis_stats"`
				Reputation      int `json:"reputation"`
				LastAnalysisRes map[string]struct {
					Category   string `json:"category"`
					Result     string `json:"result"`
					EngineName string `json:"engine_name"`
				} `json:"last_analysis_results"`
			} `json:"attributes"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &parsed); err != nil {
		res.Error = "decode: " + err.Error()
		return []SourceResult{res}
	}

	stats := parsed.Data.Attributes.LastAnalysisStats
	d := vtDetails{
		Malicious:  stats.Malicious,
		Suspicious: stats.Suspicious,
		Harmless:   stats.Harmless,
		Undetected: stats.Undetected,
		Total:      stats.Harmless + stats.Malicious + stats.Suspicious + stats.Undetected + stats.Timeout,
		Reputation: parsed.Data.Attributes.Reputation,
	}
	for engine, v := range parsed.Data.Attributes.LastAnalysisRes {
		if v.Category != "malicious" && v.Category != "suspicious" {
			continue
		}
		name := v.EngineName
		if name == "" {
			name = engine
		}
		d.Vendors = append(d.Vendors, vtVendorVerdict{Engine: name, Category: v.Category, Result: v.Result})
	}
	sort.Slice(d.Vendors, func(i, j int) bool {
		if d.Vendors[i].Category != d.Vendors[j].Category {
			return d.Vendors[i].Category == "malicious"
		}
		return d.Vendors[i].Engine < d.Vendors[j].Engine
	})
	res.Details = mustJSON(d)

	if d.Malicious == 0 && d.Suspicious == 0 {
		// Clean.
		return []SourceResult{res}
	}
	res.Listed = true
	if d.Malicious > 0 {
		res.Severity = SeverityCrit
	} else {
		res.Severity = SeverityWarn
	}
	for _, v := range d.Vendors {
		res.Reasons = append(res.Reasons, v.Engine)
		res.Evidence = append(res.Evidence, Evidence{
			Label: "Engine", Value: v.Engine, Status: v.Category,
			Extra: map[string]string{"result": v.Result},
		})
	}
	return []SourceResult{res}
}

func (*virusTotalSource) Diagnose(res SourceResult) Diagnosis {
	var d vtDetails
	_ = json.Unmarshal(res.Details, &d)
	previewN := min(len(d.Vendors), 5)
	preview := make([]string, 0, previewN)
	for _, v := range d.Vendors[:previewN] {
		preview = append(preview, v.Engine)
	}
	gravity := "Suspicious"
	sev := SeverityWarn
	if d.Malicious > 0 {
		gravity = "Malicious"
		sev = SeverityCrit
	}
	return Diagnosis{
		Severity: sev,
		Title:    fmt.Sprintf("VirusTotal: %d/%d engine(s) flagged the domain (%s)", d.Malicious+d.Suspicious, d.Total, gravity),
		Detail: fmt.Sprintf(
			"Reputation %d. Vendors flagging this domain include: %s. Open the VirusTotal page to see the per-engine verdicts and the related URLs/downloads. If you believe the verdicts are stale, request a re-scan from the VirusTotal page; for false positives, contact each engine vendor directly (VT does not arbitrate).",
			d.Reputation, joinNonEmpty(preview, ", "),
		),
		Fix:      res.Reference,
		FixIsURL: res.Reference != "",
	}
}

func (*virusTotalSource) RenderDetail(res SourceResult) (template.HTML, error) {
	var d vtDetails
	if len(res.Details) > 0 {
		if err := json.Unmarshal(res.Details, &d); err != nil {
			return "", fmt.Errorf("virustotal: decode details: %w", err)
		}
	}
	if d.Total == 0 && len(d.Vendors) == 0 {
		return "", nil
	}
	var b bytes.Buffer
	if err := vtDetailTpl.Execute(&b, d); err != nil {
		return "", err
	}
	return template.HTML(b.String()), nil
}

var vtDetailTpl = template.Must(template.New("vt_detail").Parse(`
<p>Engines: <strong{{if gt .Malicious 0}} class="warn"{{end}}>{{.Malicious}} malicious</strong>, <strong>{{.Suspicious}} suspicious</strong>, {{.Harmless}} harmless, {{.Undetected}} undetected (total {{.Total}}). Reputation score: <strong>{{.Reputation}}</strong>.</p>
{{if .Vendors}}<table>
<thead><tr><th>Engine</th><th>Verdict</th><th>Result</th></tr></thead>
<tbody>{{range .Vendors}}<tr class="row-{{if eq .Category "malicious"}}crit{{else}}warn{{end}}">
<td>{{.Engine}}</td><td>{{.Category}}</td><td>{{.Result}}</td>
</tr>{{end}}</tbody>
</table>{{end}}`))