package services

import (
	"context"
	"errors"
	"log"
	"time"

	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/db"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/models"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/vulndb"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

// FindingDiff is what one server's scan changes.
type FindingDiff struct {
	Upserts   []models.VulnFinding
	FixedIDs  []bson.ObjectID
	ReopenIDs []bson.ObjectID
	// NewlyOpened is what the digest reports: findings that were not open
	// before this scan. A finding that was already open must not re-alert every
	// tick, or the digest becomes noise and stops being read.
	NewlyOpened []models.VulnFinding
}

func findingKey(cveID, pkg string) string { return cveID + "\x00" + pkg }

// DiffFindings computes the state changes for one server's scan.
//
// Pure by design: no database, no clock of its own. The ordering below is
// load-bearing — see the comment above the second loop.
func DiffFindings(existing []models.VulnFinding, results []vulndb.Result, now time.Time) FindingDiff {
	var d FindingDiff

	byKey := make(map[string]models.VulnFinding, len(existing))
	for _, f := range existing {
		byKey[findingKey(f.CVEID, f.PackageName)] = f
	}

	seen := make(map[string]bool, len(results))
	for _, r := range results {
		key := findingKey(r.CVEID, r.PackageName)
		seen[key] = true

		prev, had := byKey[key]

		f := models.VulnFinding{
			CVEID:       r.CVEID,
			PackageName: r.PackageName,
			Installed:   r.Installed,
			FixedIn:     r.FixedIn,
			Severity:    r.Severity,
			CVSSScore:   r.CVSSScore,
			Title:       r.Title,
			References:  r.References,
			State:       models.FindingOpen,
			FirstSeen:   now,
			LastSeen:    now,
		}

		if had {
			f.ID = prev.ID
			// Preserved, never overwritten: an upsert that moves first_seen
			// forward makes every finding look discovered today.
			f.FirstSeen = prev.FirstSeen

			// A live acceptance survives the scan untouched: it is suppressed
			// from counts and alerts until its expiry, then reopens on its own.
			if prev.State == models.FindingAccepted && prev.Accepted != nil {
				if now.Before(prev.Accepted.Until) {
					continue
				}
				d.ReopenIDs = append(d.ReopenIDs, prev.ID)
				continue
			}

			if prev.State != models.FindingOpen {
				d.NewlyOpened = append(d.NewlyOpened, f)
			}
		} else {
			d.NewlyOpened = append(d.NewlyOpened, f)
		}

		d.Upserts = append(d.Upserts, f)
	}

	// Anything we hold that this scan did not produce is fixed. This runs AFTER
	// the loop above, and the ordering matters: a finding that is both absent
	// and past its acceptance expiry must settle as fixed rather than reopening
	// on a package that no longer carries it.
	for _, f := range existing {
		if seen[findingKey(f.CVEID, f.PackageName)] {
			continue
		}
		if f.State == models.FindingFixed {
			continue
		}
		d.FixedIDs = append(d.FixedIDs, f.ID)
	}

	return d
}

// ErrFindingNotFound is returned for a finding that does not exist in this
// instance. Callers turn it into a 404 — never a 403, which would confirm the
// finding exists in someone else's instance.
var ErrFindingNotFound = errors.New("finding not found")

// FindingFilter narrows a fleet-wide finding query. An empty field is no
// filter.
type FindingFilter struct {
	Severity string
	State    string
	ServerID string
	Tags     map[string]string
	// HasFix nil is no filter. true is "a vendor fix exists, this is
	// patchable"; false is the unfixable set — remove the package, disable the
	// service, or accept it, but do not wait for an update.
	HasFix *bool
}

// ListInstanceFindings returns findings across the whole fleet.
func ListInstanceFindings(instanceID string, f FindingFilter) ([]models.VulnFinding, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	filter := bson.M{"instance_id": instanceID}
	if f.State != "" {
		filter["state"] = f.State
	}
	if f.Severity != "" {
		filter["severity"] = f.Severity
	}
	if f.ServerID != "" {
		filter["server_id"] = f.ServerID
	}

	// fixed_in is omitempty, so a finding with no vendor fix carries no such
	// field at all rather than an empty string. Both forms must be matched, or
	// the unfixable set reads as empty on any document written before this.
	if f.HasFix != nil {
		if *f.HasFix {
			filter["fixed_in"] = bson.M{"$nin": bson.A{"", nil}}
		} else {
			filter["fixed_in"] = bson.M{"$in": bson.A{"", nil}}
		}
	}

	// The tag selector resolves through ResolveTargets, the single answer to
	// which servers a selector touches. A second matcher here could disagree
	// with what a workflow means by env:prod.
	if len(f.Tags) > 0 {
		servers, err := ResolveTargets(instanceID, nil, f.Tags)
		if err != nil {
			return nil, err
		}
		ids := make([]string, 0, len(servers))
		for _, s := range servers {
			ids = append(ids, s.ServerID)
		}
		if len(ids) == 0 {
			return []models.VulnFinding{}, nil
		}
		filter["server_id"] = bson.M{"$in": ids}
	}

	cur, err := db.Col("vuln_findings").Find(ctx, filter)
	if err != nil {
		return nil, err
	}
	defer cur.Close(ctx)

	out := []models.VulnFinding{}
	if err := cur.All(ctx, &out); err != nil {
		return nil, err
	}
	return out, nil
}

// CountOpenFindingsBySeverity powers the summary tiles. Accepted findings are
// excluded: they are suppressed from counts until their expiry, which is the
// whole point of accepting one.
func CountOpenFindingsBySeverity(instanceID string) (map[string]int, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	cur, err := db.Col("vuln_findings").Aggregate(ctx, []bson.M{
		{"$match": bson.M{"instance_id": instanceID, "state": models.FindingOpen}},
		{"$group": bson.M{"_id": "$severity", "n": bson.M{"$sum": 1}}},
	})
	if err != nil {
		return nil, err
	}
	defer cur.Close(ctx)

	var rows []struct {
		Severity string `bson:"_id"`
		N        int    `bson:"n"`
	}
	if err := cur.All(ctx, &rows); err != nil {
		return nil, err
	}

	counts := map[string]int{}
	for _, r := range rows {
		counts[r.Severity] = r.N
	}
	return counts, nil
}

// MarkInstanceForRescan flags every server in an instance for rescanning and
// returns how many were flagged.
//
// It does not scan. vulnsched picks the flags up on its next tick, which keeps
// matching on the leader and means this endpoint cannot become a second
// scanning path.
func MarkInstanceForRescan(instanceID string) (int64, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	res, err := db.Col("server_packages").UpdateMany(ctx,
		bson.M{"instance_id": instanceID},
		bson.M{"$set": bson.M{"scan_pending": true}},
	)
	if err != nil {
		log.Printf("vulnsched: mark rescan for instance %s: %v", instanceID, err)
		return 0, err
	}
	log.Printf("vulnsched: rescan requested for instance %s, %d of %d server(s) flagged",
		instanceID, res.ModifiedCount, res.MatchedCount)
	return res.ModifiedCount, nil
}

// AcceptFinding suppresses a finding until a date, with a reason.
//
// The expiry is mandatory at the API layer. A finding reopens on its own when
// it passes, which is what stops the accepted list becoming where risk goes to
// be forgotten.
func AcceptFinding(instanceID, findingID, actor, reason string, until time.Time) (*models.VulnFinding, error) {
	id, err := bson.ObjectIDFromHex(findingID)
	if err != nil {
		return nil, ErrFindingNotFound
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var f models.VulnFinding
	err = db.Col("vuln_findings").FindOneAndUpdate(ctx,
		bson.M{"_id": id, "instance_id": instanceID},
		bson.M{"$set": bson.M{
			"state": models.FindingAccepted,
			"accepted": models.Acceptance{
				By:     actor,
				Reason: reason,
				Until:  until,
				At:     time.Now(),
			},
		}},
		options.FindOneAndUpdate().SetReturnDocument(options.After),
	).Decode(&f)
	if err == mongo.ErrNoDocuments {
		return nil, ErrFindingNotFound
	}
	if err != nil {
		return nil, err
	}
	return &f, nil
}

// UnacceptFinding returns an accepted finding to open before its expiry.
func UnacceptFinding(instanceID, findingID string) (*models.VulnFinding, error) {
	id, err := bson.ObjectIDFromHex(findingID)
	if err != nil {
		return nil, ErrFindingNotFound
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var f models.VulnFinding
	err = db.Col("vuln_findings").FindOneAndUpdate(ctx,
		bson.M{"_id": id, "instance_id": instanceID},
		bson.M{
			"$set":   bson.M{"state": models.FindingOpen},
			"$unset": bson.M{"accepted": ""},
		},
		options.FindOneAndUpdate().SetReturnDocument(options.After),
	).Decode(&f)
	if err == mongo.ErrNoDocuments {
		return nil, ErrFindingNotFound
	}
	if err != nil {
		return nil, err
	}
	return &f, nil
}

// ListFindings returns every finding held for one server.
func ListFindings(ctx context.Context, instanceID, serverID string) ([]models.VulnFinding, error) {
	cur, err := db.Col("vuln_findings").Find(ctx, bson.M{
		"instance_id": instanceID,
		"server_id":   serverID,
	})
	if err != nil {
		return nil, err
	}
	defer cur.Close(ctx)

	var out []models.VulnFinding
	if err := cur.All(ctx, &out); err != nil {
		return nil, err
	}
	return out, nil
}

// ApplyFindingDiff writes a diff. Thin on purpose — the logic worth reading
// twice is all in DiffFindings.
func ApplyFindingDiff(ctx context.Context, instanceID, serverID string, d FindingDiff, now time.Time) error {
	col := db.Col("vuln_findings")

	// One BulkWrite per batch, not one UpdateOne per finding. A freshly scanned
	// Ubuntu host opens tens of thousands of findings, and at one round trip each
	// that is minutes of sequential latency during which the tick holds the
	// leader and every other server waits its turn. Unordered, because the
	// upserts are independent and one duplicate-key race must not abandon the
	// rest of the batch.
	const bulkBatch = 1000
	ops := make([]mongo.WriteModel, 0, bulkBatch)

	flush := func() error {
		if len(ops) == 0 {
			return nil
		}
		_, err := col.BulkWrite(ctx, ops, options.BulkWrite().SetOrdered(false))
		ops = ops[:0]
		return err
	}

	for _, f := range d.Upserts {
		ops = append(ops, mongo.NewUpdateOneModel().
			SetFilter(bson.M{
				"instance_id":  instanceID,
				"server_id":    serverID,
				"cve_id":       f.CVEID,
				"package_name": f.PackageName,
			}).
			SetUpdate(bson.M{
				"$set": bson.M{
					"installed_version": f.Installed,
					"fixed_in":          f.FixedIn,
					"severity":          f.Severity,
					"cvss_score":        f.CVSSScore,
					"title":             f.Title,
					"references":        f.References,
					"state":             models.FindingOpen,
					"last_seen":         now,
				},
				// first_seen is written only on insert, so a rescan cannot move
				// it forward.
				"$setOnInsert": bson.M{
					"instance_id":  instanceID,
					"server_id":    serverID,
					"cve_id":       f.CVEID,
					"package_name": f.PackageName,
					"first_seen":   f.FirstSeen,
				},
				"$unset": bson.M{"fixed_at": "", "accepted": ""},
			}).
			SetUpsert(true))

		if len(ops) >= bulkBatch {
			if err := flush(); err != nil {
				return err
			}
		}
	}
	if err := flush(); err != nil {
		return err
	}

	if len(d.FixedIDs) > 0 {
		if _, err := col.UpdateMany(ctx,
			bson.M{"_id": bson.M{"$in": d.FixedIDs}},
			bson.M{"$set": bson.M{"state": models.FindingFixed, "fixed_at": now}},
		); err != nil {
			return err
		}
	}

	if len(d.ReopenIDs) > 0 {
		if _, err := col.UpdateMany(ctx,
			bson.M{"_id": bson.M{"$in": d.ReopenIDs}},
			bson.M{"$set": bson.M{"state": models.FindingOpen, "last_seen": now}, "$unset": bson.M{"accepted": ""}},
		); err != nil {
			return err
		}
	}

	return nil
}

// StartVulnSweeper deletes old FIXED findings. Open and accepted findings are
// never swept at any setting: retention is about history, and an unresolved
// vulnerability is not history.
func StartVulnSweeper(ctx context.Context) {
	go func() {
		ticker := time.NewTicker(6 * time.Hour)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				sweepFixedFindings(ctx)
			}
		}
	}()
}

// defaultVulnRetentionDays is what an unset setting means. A pointer field and
// this constant together give absent-means-90 and 0-means-forever, the same
// shape as workflow log retention.
const defaultVulnRetentionDays = 90

// sweepFixedFindings deletes fixed findings past each instance's retention.
//
// Only "fixed" is ever swept. An open or accepted finding is not history, it is
// an outstanding decision, and deleting one on a timer would quietly shrink the
// fleet's risk picture.
func sweepFixedFindings(ctx context.Context) {
	instanceIDs, err := ListInstanceIDs()
	if err != nil {
		log.Printf("vuln sweeper: list instances: %v", err)
		return
	}

	for _, instanceID := range instanceIDs {
		if ctx.Err() != nil {
			return
		}

		days := defaultVulnRetentionDays
		if s, err := GetSettings(instanceID); err == nil && s != nil && s.VulnFindingRetentionDays != nil {
			days = *s.VulnFindingRetentionDays
		}
		if days <= 0 {
			continue // 0 means keep forever
		}

		cutoff := time.Now().AddDate(0, 0, -days)
		res, err := db.Col("vuln_findings").DeleteMany(ctx, bson.M{
			"instance_id": instanceID,
			"state":       models.FindingFixed,
			"fixed_at":    bson.M{"$lt": cutoff},
		})
		if err != nil {
			log.Printf("vuln sweeper: delete for %s: %v", instanceID, err)
			continue
		}
		if res.DeletedCount > 0 {
			log.Printf("vuln sweeper: removed %d fixed findings for %s", res.DeletedCount, instanceID)
		}
	}
}