package reconcile

import (
	"context"
	"fmt"
	"strings"

	"git.d-ma.be/mathias/tapir/internal/domain"
	"git.d-ma.be/mathias/tapir/internal/progress"
)

// MaxContested is the default bound on the stored list of judged spots. A 90-minute
// recording has a few hundred, so this is a guard against a runaway, not a budget. The
// counts on the Reconciliation are always exact; only the list is capped.
const MaxContested = 2000

// beforeWords is how many words of left context a stored spot carries: enough to find
// the place again, small enough to stay out of the way.
const beforeWords = 8

// Options configure Reconcile. Zero Window and BreakAfter take their defaults.
type Options struct {
	// Window is how many words of the base to show on each side of a dispute. 100 is
	// what was measured (docs/research/local-model-adjudicator-sweep-2026-10-01.md).
	Window int
	// BreakAfter is how many consecutive model errors stop the reconciliation, so a
	// model that is down costs three timeouts, not one per dispute.
	BreakAfter int
	// MaxContested caps the stored judged spots; zero takes the package default.
	MaxContested int
	// BaseEngine, AltEngine and Reconciler are recorded on the result.
	BaseEngine, AltEngine, Reconciler string
}

// Reconcile merges two transcripts of the same speech. The base text stands except
// where the model chooses the alternative's reading of a dispute; trivial disputes
// (a compound split, a filler word) are resolved without the model. It never fails a
// transcription: if the model cannot be used the base text comes back unreconciled
// and the result says why. The only error it returns is a cancelled context, which
// is the job stopping, not a verdict on the model.
func Reconcile(ctx context.Context, base, alt domain.Transcription, adj Adjudicator, o Options) (domain.Transcription, error) {
	window, breakAfter, maxSpots := o.Window, o.BreakAfter, o.MaxContested
	if window <= 0 {
		window = 100
	}
	if breakAfter <= 0 {
		breakAfter = 3
	}
	if maxSpots <= 0 {
		maxSpots = MaxContested
	}
	rep := progress.From(ctx)

	rep.Begin("compare", 0)
	bt, at := Tokenize(base.Text), Tokenize(alt.Text)
	ds := FindDisputes(bt, at)
	rep.Done("compare")

	// The deciding phase counts what the MODEL decides: the trivial disputes are settled
	// without it, so counting them would leave the bar short of its end.
	deciding := 0
	for _, d := range ds {
		if !AutoResolve(d) {
			deciding++
		}
	}
	rec := &domain.Reconciliation{
		Status: domain.ReconcileDone, Base: o.BaseEngine, Alt: o.AltEngine, Reconciler: o.Reconciler, Disputes: len(ds),
	}
	seconds := positioner(base, len(bt))
	spot := func(d Dispute, used string, unsure bool) {
		before, _ := ContextAround(bt, d, window)
		if w := strings.Fields(before); len(w) > beforeWords {
			before = strings.Join(w[len(w)-beforeWords:], " ")
		}
		if len(rec.Contested) < maxSpots {
			rec.Contested = append(rec.Contested, domain.ContestedSpot{
				Seconds: seconds(d.BaseStart), Base: d.Base, Alt: d.Alt, Used: used, Unsure: unsure, Before: before,
			})
		}
	}

	rep.Begin("decide", deciding)
	verdicts := make([]Verdict, len(ds))
	consecutive, handled := 0, 0
	for i, d := range ds {
		if AutoResolve(d) {
			rec.AutoResolved++
			verdicts[i] = VerdictBase
			continue
		}
		if err := ctx.Err(); err != nil {
			return domain.Transcription{}, err
		}
		before, after := ContextAround(bt, d, window)
		v, err := adj.Adjudicate(ctx, Request{Before: before, After: after, Base: d.Base, Alt: d.Alt, AltFirst: i%2 == 1})
		if err != nil {
			if ctxErr := ctx.Err(); ctxErr != nil {
				return domain.Transcription{}, ctxErr
			}
			rec.Errors++
			consecutive++
			if consecutive >= breakAfter {
				failure := fmt.Errorf("the model failed %d times in a row: %w", consecutive, err)
				rep.Fail("decide", failure)
				return skipped(base, rec, failure.Error()), nil
			}
			v = VerdictUnsure
		} else {
			consecutive = 0
			rec.Adjudicated++
		}
		verdicts[i] = v
		switch v {
		case VerdictAlt:
			rec.Changed++
			spot(d, "alt", false)
		case VerdictUnsure:
			rec.Unsure++
			spot(d, "base", true)
		default:
			spot(d, "base", false)
		}
		handled++
		rep.Advance("decide", handled)
	}
	rep.Done("decide")

	out := base
	out.Text, out.Segments = merge(base, bt, at, ds, verdicts)
	out.Reconciliation = rec
	return out, nil
}

// merge applies the verdicts: the base's words stand except where the alternative's
// reading was chosen. It returns the merged text and the base's segments carrying the
// merged words (timing is the base's). Segments are rewritten only when their words line
// up with the text; a word wrapped across a segment edge makes them differ, and guessing
// would put words in the wrong place, so then the base's segments are returned as they
// were (ADR-035). The base's own segments are never edited in place.
func merge(base domain.Transcription, bt, at []Token, ds []Dispute, verdicts []Verdict) (string, []domain.Segment) {
	segOf, aligned := segmentIndex(base.Segments, len(bt))
	var perSeg [][]string
	var changed map[int]bool
	if aligned {
		perSeg, changed = make([][]string, len(base.Segments)), map[int]bool{}
	}

	var parts []string
	keep := func(from, to int) {
		for i := from; i < to; i++ {
			parts = append(parts, bt[i].Text)
			if aligned {
				perSeg[segOf[i]] = append(perSeg[segOf[i]], bt[i].Text)
			}
		}
	}
	pos := 0
	for i, d := range ds {
		keep(pos, d.BaseStart)
		if verdicts[i] == VerdictAlt {
			for _, t := range at[d.AltStart:d.AltEnd] {
				parts = append(parts, t.Text)
			}
			if aligned {
				anchor := min(d.BaseStart, len(bt)-1)
				for _, t := range at[d.AltStart:d.AltEnd] {
					perSeg[segOf[anchor]] = append(perSeg[segOf[anchor]], t.Text)
				}
				changed[segOf[anchor]] = true
				for k := d.BaseStart; k < d.BaseEnd; k++ {
					changed[segOf[k]] = true
				}
			}
		} else {
			keep(d.BaseStart, d.BaseEnd)
		}
		pos = d.BaseEnd
	}
	keep(pos, len(bt))

	if !aligned {
		return strings.Join(parts, " "), base.Segments
	}
	segs := append([]domain.Segment(nil), base.Segments...)
	for k := range changed {
		segs[k].Text = strings.Join(perSeg[k], " ")
	}
	return strings.Join(parts, " "), segs
}

// segmentIndex maps each base token to the segment it came from. aligned is false when
// there are no segments or their word count is not the text's.
func segmentIndex(segs []domain.Segment, tokens int) (segOf []int, aligned bool) {
	if len(segs) == 0 || tokens == 0 {
		return nil, false
	}
	segOf = make([]int, 0, tokens)
	for k, s := range segs {
		for range Tokenize(s.Text) {
			segOf = append(segOf, k)
		}
	}
	return segOf, len(segOf) == tokens
}

// skipped is the base text, none of the partial changes, and the reason.
func skipped(base domain.Transcription, rec *domain.Reconciliation, reason string) domain.Transcription {
	out := base
	out.Reconciliation = &domain.Reconciliation{
		Status: domain.ReconcileSkipped, Reason: reason, Base: rec.Base, Alt: rec.Alt, Reconciler: rec.Reconciler,
		Disputes: rec.Disputes, AutoResolved: rec.AutoResolved, Errors: rec.Errors,
	}
	return out
}

// positioner maps a token index of the base to the second it starts at: through the
// segments when there are any, proportionally otherwise.
func positioner(base domain.Transcription, tokens int) func(int) float64 {
	if len(base.Segments) == 0 {
		return func(i int) float64 {
			if tokens == 0 {
				return 0
			}
			return float64(i) / float64(tokens) * base.AudioSeconds
		}
	}
	starts := make([]int, len(base.Segments))
	n := 0
	for i, s := range base.Segments {
		starts[i] = n
		n += len(Tokenize(s.Text))
	}
	return func(i int) float64 {
		k := 0
		for k+1 < len(starts) && starts[k+1] <= i {
			k++
		}
		return base.Segments[k].Start
	}
}