// Compact-labels command: one-time (but safely re-runnable) rebalance of
// labels.jsonl against the same per-feed negative ceiling `export` now
// enforces. Existing labels were written before the ceiling existed, so a
// handful of firehoses can already own a large share of the negative class;
// this trims them back.
//
// This is the one deliberate exception to labels.jsonl's append-only rule
// (see core/labels.go). It exists because the ceiling only bounds the
// future: without a compaction, pre-cap history keeps distorting training
// until it ages out — which it never does. Rules that keep it honest:
// positives are never touched; a feed's newest negatives are the ones kept
// (recency carries the signal); the original file is copied to a backup
// before anything is rewritten; and running it twice is a no-op.
package cmds

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"sort"

	"fluxrec/core"
)

type CompactLabelsCommand struct {
	Labels        string
	NegRatio      int
	NegFeedCapPct int
	Backup        string
	DryRun        bool
}

func (c *CompactLabelsCommand) Name() string { return "compact-labels" }

func (c *CompactLabelsCommand) Init(args []string) error {
	fs := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	fs.Usage = func() {
		fmt.Fprint(fs.Output(), `Usage: fluxrec compact-labels [flags]

Trim per-feed negative over-representation in labels.jsonl down to the
same ceiling export enforces (--neg-feed-cap-pct of the target, floor 50).
Positives are untouched; each trimmed feed keeps its NEWEST negatives; the
original is copied to --backup first; re-running is a no-op. Use --dry-run
to see the effect without writing anything.
`)
		fs.PrintDefaults()
	}
	fs.StringVar(&c.Labels, "labels", "labels.jsonl", "Path to labels.jsonl")
	fs.IntVar(&c.NegRatio, "neg-ratio", 20, "Negatives-per-positive ratio the ceiling is computed against")
	fs.IntVar(&c.NegFeedCapPct, "neg-feed-cap-pct", 5, "Per-feed negative ceiling as % of target (floor 50)")
	fs.StringVar(&c.Backup, "backup", "", "Backup path for the original file (default: LABELS.bak)")
	fs.BoolVar(&c.DryRun, "dry-run", false, "Report what would change without writing")
	if err := fs.Parse(args); err != nil {
		return err
	}
	if c.NegRatio < 1 {
		return fmt.Errorf("--neg-ratio must be >= 1, got %d", c.NegRatio)
	}
	if c.NegFeedCapPct < 1 {
		return fmt.Errorf("--neg-feed-cap-pct must be >= 1, got %d", c.NegFeedCapPct)
	}
	return nil
}

func (c *CompactLabelsCommand) Run(stdin io.Reader, stdout io.Writer) error {
	rows, err := core.ReadLabels(c.Labels)
	if err != nil {
		return err
	}
	if len(rows) == 0 {
		fmt.Fprintf(stdout, "compact-labels: %s is empty or missing; nothing to do\n", c.Labels)
		return nil
	}

	posTotal := 0
	for _, row := range rows {
		if row.Label == core.LabelPos {
			posTotal++
		}
	}
	cap := feedNegCap(c.NegFeedCapPct, c.NegRatio, posTotal)
	if cap <= 0 {
		return fmt.Errorf("computed per-feed ceiling is 0 (positives=%d); check --neg-ratio/--neg-feed-cap-pct", posTotal)
	}

	// Per-feed negative row indexes, to decide what each overfull feed keeps.
	negIdx := make(map[int64][]int)
	for i, row := range rows {
		if row.Label == core.LabelNeg {
			negIdx[row.FeedID] = append(negIdx[row.FeedID], i)
		}
	}

	drop := make(map[int]bool)
	trimmed := 0
	for _, idxs := range negIdx {
		if len(idxs) <= cap {
			continue
		}
		trimmed++
		// Keep the newest `cap` rows for this feed; newest-first sort with a
		// deterministic tiebreak, so repeated runs agree.
		sorted := make([]int, len(idxs))
		copy(sorted, idxs)
		sort.Slice(sorted, func(a, b int) bool {
			ra, rb := rows[sorted[a]], rows[sorted[b]]
			if ra.CapturedAt.Equal(rb.CapturedAt) {
				return ra.EntryID > rb.EntryID
			}
			return ra.CapturedAt.After(rb.CapturedAt)
		})
		for _, idx := range sorted[cap:] {
			drop[idx] = true
		}
	}

	kept := make([]core.LabelRow, 0, len(rows)-len(drop))
	for i, row := range rows {
		if !drop[i] {
			kept = append(kept, row)
		}
	}

	negTotal := len(rows) - posTotal
	fmt.Fprintf(stdout, "compact-labels: rows %d → %d (positives %d kept; negatives %d → %d; feeds trimmed %d; cap %d/feed)\n",
		len(rows), len(kept), posTotal, negTotal, negTotal-len(drop), trimmed, cap)
	if len(drop) == 0 {
		fmt.Fprintln(stdout, "nothing over the ceiling; labels.jsonl unchanged")
		return nil
	}
	if c.DryRun {
		fmt.Fprintf(stdout, "dry-run: would drop %d negatives (no files written)\n", len(drop))
		return nil
	}

	backup := c.Backup
	if backup == "" {
		backup = c.Labels + ".bak"
	}
	orig, err := os.ReadFile(c.Labels)
	if err != nil {
		return fmt.Errorf("reading %s for backup: %w", c.Labels, err)
	}
	if err := os.WriteFile(backup, orig, 0o644); err != nil {
		return fmt.Errorf("writing backup %s: %w", backup, err)
	}

	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	for _, row := range kept {
		if err := enc.Encode(row); err != nil {
			return fmt.Errorf("encoding label row: %w", err)
		}
	}
	if err := os.WriteFile(c.Labels, buf.Bytes(), 0o644); err != nil {
		return fmt.Errorf("rewriting %s: %w", c.Labels, err)
	}
	fmt.Fprintf(stdout, "wrote %s (backup: %s); re-running export will sample fresh negatives under the ceiling\n",
		c.Labels, backup)
	return nil
}