package pipeline

import (
	"encoding/json"
	"fmt"
	"strings"

	"textmachine/backend/internal/chunk"
	"textmachine/backend/internal/config"
	"textmachine/backend/internal/lang"
	"textmachine/backend/internal/membank"
)

// repin.go: POINTWISE re-edit by key (pack-20 point 5, D39.42 п.5) — the machinery that makes "the bank
// changed" cost only the units the change actually reaches.
//
// THE PROBLEM. The bank is folded into the wave snapshot, so signing one term moves the edit-wave
// snapshot, and the resume fast-path demands an exact snapshot match. Every already-paid unit of that
// wave therefore re-translates — a whole edit wave bought again for one word. Owner, 26.07: «если
// меняется банк — перезапускается редактура; вопрос насколько широко, на какие главы».
//
// THE OBSERVATION. A snapshot is a hash over many things, and only ONE of them is the bank. If the ONLY
// component that moved is the memory version, then every other input to the request — model, prompt SHA,
// temperature, reasoning, capability, budgets, gates — is bit-for-bit what it was. What the bank change
// does to a given unit is then fully observable in ONE place: the injected message. And the injected
// message is already hashed per unit, in chunk_status.content_hash, by the resume fast-path.
//
// So: bank-only snapshot move + unchanged content hash ⇒ the request this run would issue is byte-identical
// to the one already paid for, except for the snapshot id inside its address. Serving the stored result is
// then exactly as sound as an ordinary resume, and the unit is RE-PINNED for $0 instead of re-bought.
// A unit whose injected bytes DID change is the one where the term actually occurs — and that one is
// re-translated, which is the whole point of signing it.
//
// The same predicate drives the ESTIMATE: projectRebill uses it, so the "N units, ~$X" an operator is
// asked to consent to counts the units that will really be re-paid, not the whole wave.

// snapshotMoveKind classifies a difference between a stored snapshot and the current one.
type snapshotMoveKind int

const (
	// moveUnknown — the stored payload is unavailable or unreadable, so nothing can be concluded and the
	// conservative answer holds (a full re-payment).
	moveUnknown snapshotMoveKind = iota
	// moveBankOnly — the ONLY differing component is the memory version. Every wire-shaping input is
	// unchanged, so a unit whose injected bytes are unchanged is re-pinnable for $0.
	moveBankOnly
	// moveOther — something besides the bank moved (a prompt, a model, a gate). Nothing is re-pinnable:
	// the change can alter what is sent or how the result is judged, and this predicate does not model that.
	moveOther
)

// memoryVersionField is the payload key holding the bank's content hash — the one component a signature,
// a terminologist consolidation or any other bank edit moves.
const memoryVersionField = "memory_version"

// classifySnapshotMove compares the stored snapshot's payload with the current one and reports whether the
// difference is confined to the bank. Both payloads are the same fixed-field-order JSON object, so the
// comparison is a plain key-by-key equality over the decoded maps — no knowledge of the struct is needed,
// which is what keeps this from silently going stale when a new component is folded: an unrecognised new
// field simply counts as a difference, i.e. the conservative answer.
func classifySnapshotMove(storedPayload, currentPayload string) snapshotMoveKind {
	if storedPayload == "" || currentPayload == "" {
		return moveUnknown
	}
	var was, now map[string]json.RawMessage
	if json.Unmarshal([]byte(storedPayload), &was) != nil || json.Unmarshal([]byte(currentPayload), &now) != nil {
		return moveUnknown
	}
	bankMoved := false
	for k, v := range now {
		w, had := was[k]
		if k == memoryVersionField {
			if !had || !jsonEqual(w, v) {
				bankMoved = true
			}
			continue
		}
		if !had || !jsonEqual(w, v) {
			return moveOther
		}
	}
	for k := range was {
		if _, still := now[k]; !still {
			return moveOther // a component that was folded and no longer is — not a bank move
		}
	}
	if !bankMoved {
		return moveOther // the ids differ but no component does: unexplained, so never re-pinnable
	}
	return moveBankOnly
}

func jsonEqual(a, b json.RawMessage) bool {
	if string(a) == string(b) {
		return true
	}
	var x, y any
	if json.Unmarshal(a, &x) != nil || json.Unmarshal(b, &y) != nil {
		return false
	}
	xb, xe := json.Marshal(x)
	yb, ye := json.Marshal(y)
	return xe == nil && ye == nil && string(xb) == string(yb)
}

// cachedRenderedContentHashes is renderedContentHashes memoized for the life of ONE materialized bank.
//
// The reproduction walks every stage of every position and re-renders it — the expensive half of every
// $0 projection in this engine. Since the volume ceiling landed, a single bounded run could pay for it
// THREE times: once in planVolume, and once in each of the consent gate's two projections (the book's and
// the run's). Nothing changed between them, so two of the three were pure waste, and the cost scales with
// the book — exactly where it is least affordable.
//
// ⚠ THE CACHE'S LIFETIME IS THE BANK'S, and that is what makes it safe rather than merely fast. Every
// hash here is rendered against the materialized memory, so the memo is only valid while that memory is;
// materializeBanks clears it, which is the ONE place the bank can change (a mid-run re-seed at the
// bank-mining stop goes through it like everything else). A memo that outlived its bank would hand a
// caller hashes for a bank the run no longer has — the precise class of silent wrongness the free
// estimate exists to avoid.
func (r *Runner) cachedRenderedContentHashes(chunks []chunk.Chunk, stickySel []membank.Selection) map[chunkKey]map[string]string {
	if r.contentHashes != nil {
		return r.contentHashes
	}
	r.contentHashes = r.renderedContentHashes(chunks, stickySel)
	return r.contentHashes
}

// renderedContentHashes reproduces, for every live position of every stage, the content hash the run
// WOULD compute — the same msgsContentHash the resume fast-path compares against. It is the $0 half of
// the honest estimate: with it, "the bank moved" can be answered per unit instead of per wave.
//
// The reproduction is deliberately faithful rather than approximate: the draft wave renders from the
// precomputed BASE-bank selection (what runDraftChunk uses), and the edit wave does a fresh Select over
// the unit's clean source against the ENRICHED bank with the unit's stored member drafts as {{draft}}
// (what runEditUnit uses). Where a needed input is missing — a member draft not yet resolved, a template
// absent — the position is simply OMITTED, and the caller treats an absent hash as "cannot conclude",
// i.e. the conservative answer.
func (r *Runner) renderedContentHashes(chunks []chunk.Chunk, stickySel []membank.Selection) map[chunkKey]map[string]string {
	out := map[chunkKey]map[string]string{}
	put := func(k chunkKey, stage, hash string) {
		if out[k] == nil {
			out[k] = map[string]string{}
		}
		out[k][stage] = hash
	}
	tx := lang.InjectionTextsFor(r.Book.TargetLang)

	// --- draft wave: one sequence per chunk, over the BASE bank selection ---
	draftStages := waveStages(r.Pipeline.Stages, waveDraft)
	for i, ch := range chunks {
		injection := map[string]string{}
		if r.baseMemory != nil && i < len(stickySel) {
			for role, render := range roleInjectionRenderers {
				injection[role] = render(stickySel[i].Injected, tx)
			}
		}
		prev := ""
		for _, st := range draftStages {
			tpl := r.templates[st.Name]
			if tpl == nil {
				break
			}
			msgs, err := MessagesWithInjection(tpl, RenderVars{Book: r.Book, Text: ch.Text, Draft: prev}, injection[st.Role])
			if err != nil {
				break
			}
			put(chunkKey{ch.Chapter, ch.ChunkIdx}, st.Name, msgsContentHash(msgs))
			txt, ok := r.storedStageText(ch.Chapter, ch.ChunkIdx, st.Name)
			if !ok {
				break // the next stage's {{draft}} is unknown → stop this chain
			}
			prev = txt
		}
	}

	// --- edit wave: one sequence per unit, over the ENRICHED bank ---
	editStages := waveStages(r.Pipeline.Stages, waveEdit)
	if len(editStages) == 0 || len(draftStages) == 0 {
		return out
	}
	lastDraft := draftStages[len(draftStages)-1].Name
	for _, u := range buildEditUnits(chunks) {
		var cleanSources, draftParts []string
		complete := true
		for _, m := range u.Members {
			cs, err := r.Store.GetChunkStatus(r.Book.BookID, m.Chapter, m.ChunkIdx, lastDraft)
			if err != nil || cs == nil {
				complete = false
				break
			}
			if cs.Disposition == string(DispFlagged) {
				continue // c-lite: a flagged member is DROPPED from the edit, exactly as runEditUnit does
			}
			txt, ok := r.storedStageText(m.Chapter, m.ChunkIdx, lastDraft)
			if !ok {
				complete = false
				break
			}
			cleanSources = append(cleanSources, m.Text)
			draftParts = append(draftParts, txt)
		}
		if !complete || len(draftParts) == 0 {
			continue
		}
		leader := chunk.Chunk{Chapter: u.Chapter, ChunkIdx: u.FirstChunkIdx, Text: strings.Join(cleanSources, unitJoinSeparator)}
		injection := map[string]string{}
		if r.memory != nil {
			sel := r.memory.Select(leader.Text, u.Chapter, nil, r.Pipeline.Context.GlossaryTokenBudget)
			for role, render := range roleInjectionRenderers {
				injection[role] = render(sel.Injected, tx)
			}
		}
		prev := strings.Join(draftParts, unitJoinSeparator)
		for _, st := range editStages {
			tpl := r.templates[st.Name]
			if tpl == nil {
				break
			}
			msgs, err := MessagesWithInjection(tpl, RenderVars{Book: r.Book, Text: leader.Text, Draft: prev}, injection[st.Role])
			if err != nil {
				break
			}
			put(chunkKey{u.Chapter, u.FirstChunkIdx}, st.Name, msgsContentHash(msgs))
			txt, ok := r.storedStageText(u.Chapter, u.FirstChunkIdx, st.Name)
			if !ok {
				break
			}
			prev = txt
		}
	}
	return out
}

// storedStageText returns a resolved stage's text from its chunk_status → checkpoint, as the run would
// feed it to the next stage. It follows final_hash, so a banknote-cleaned or cosmetically-stripped draft
// yields the SHIPPED bytes rather than the raw completion — otherwise the reproduction would diverge from
// the run precisely on the units where a derived export exists.
func (r *Runner) storedStageText(chapter, chunkIdx int, stage string) (string, bool) {
	cs, err := r.Store.GetChunkStatus(r.Book.BookID, chapter, chunkIdx, stage)
	if err != nil || cs == nil || cs.FinalHash == "" {
		return "", false
	}
	cp, err := r.Store.GetCheckpoint(cs.FinalHash)
	if err != nil || cp == nil {
		return "", false
	}
	return cp.ResponseText, true
}

// repinDecider answers "is this stored row re-pinnable at $0?" for one wave, caching the current payload
// and the per-stored-snapshot classification so a book-sized projection does a handful of store reads
// rather than one per row.
type repinDecider struct {
	r        *Runner
	payloads map[wave]string             // current payload per wave (rendered lazily)
	kinds    map[string]snapshotMoveKind // stored snapshot id → classification against the current one
}

func newRepinDecider(r *Runner) *repinDecider {
	return &repinDecider{r: r, payloads: map[wave]string{}, kinds: map[string]snapshotMoveKind{}}
}

// bankOnlyMove reports whether the stored snapshot differs from the current one ONLY in the bank.
func (d *repinDecider) bankOnlyMove(storedID string, w wave) (bool, error) {
	if k, done := d.kinds[storedID]; done {
		return k == moveBankOnly, nil
	}
	cur, ok := d.payloads[w]
	if !ok {
		_, payload, err := d.r.snapshotIDForWave(w)
		if err != nil {
			return false, fmt.Errorf("pipeline: render the current snapshot payload for the re-pin check: %w", err)
		}
		d.payloads[w] = payload
		cur = payload
	}
	stored, err := d.r.Store.SnapshotPayload(storedID)
	if err != nil {
		return false, fmt.Errorf("pipeline: read the stored snapshot payload %.12s: %w", storedID, err)
	}
	k := classifySnapshotMove(stored, cur)
	d.kinds[storedID] = k
	return k == moveBankOnly, nil
}

// waveOfStage reports which wave owns a stage by name (the same partition waveStages makes). An unknown
// name defaults to the edit wave, which is the conservative side: an unknown stage's stored payload will
// not match the edit-wave render, so the re-pin predicate answers "no".
func waveOfStage(stages []config.Stage, name string) wave {
	for _, st := range stages {
		if st.Name == name {
			if st.Role == roleTranslator {
				return waveDraft
			}
			return waveEdit
		}
	}
	return waveEdit
}

// repinnable is the runner-level predicate behind the $0 re-pin: the stored snapshot differs from the
// current one ONLY in the bank. It caches per (stored id, wave) for the run, so the wave executor's
// parallel workers do not each re-read the payload — the cache is written under a mutex because the draft
// wave calls this from N goroutines.
func (r *Runner) repinnable(storedID, currentID string, w wave) bool {
	if storedID == "" || storedID == currentID {
		return false
	}
	r.repinMu.Lock()
	defer r.repinMu.Unlock()
	if r.repinCache == nil {
		r.repinCache = map[string]bool{}
	}
	key := storedID + "\x00" + currentID
	if v, done := r.repinCache[key]; done {
		return v
	}
	ok, err := newRepinDecider(r).bankOnlyMove(storedID, w)
	if err != nil {
		// A read failure is never a licence to serve a stale result — but it is also the difference between
		// "this book has nothing re-pinnable" and "the store could not answer", and the operator pays the
		// second one in re-translated chunks. The negative is cached for the run, so this logs once.
		r.Log.Warn("re-pin check could not read the stored snapshot; this unit will be re-translated rather than re-pinned",
			"stored_snapshot", storedID, "err", err)
		ok = false
	}
	r.repinCache[key] = ok
	return ok
}