package pipeline

import (
	"reflect"
	"strings"
	"testing"
	"unicode"
	"unicode/utf8"
)

// testSeg is the ratified WS2 output-token budget (zh-ru defaults) used by the chunker tests.
func testSeg() SegBudget {
	return SegBudget{DraftBudgetOut: 1797, EditCeilingOut: 3200, FertCJK: 1.1978, FertOther: 0.3852}
}

// stripMeta zeroes the WS2 additive fields (EstOut/OversizedSentence/EditUnitID) so a
// {Chapter,ChunkIdx,Text}-level DeepEqual still expresses the intent of the pre-WS2 assertions.
func stripMeta(cs []Chunk) []Chunk {
	out := make([]Chunk, len(cs))
	for i, c := range cs {
		out[i] = Chunk{Chapter: c.Chapter, ChunkIdx: c.ChunkIdx, Text: c.Text}
	}
	return out
}

func TestSplitChunksSingleParagraph(t *testing.T) {
	got := SplitChunks([]string{"静かな図書館の朝。"}, testSeg())
	want := []Chunk{{Chapter: 1, ChunkIdx: 0, Text: "静かな図書館の朝。"}}
	if !reflect.DeepEqual(stripMeta(got), want) {
		t.Fatalf("single paragraph = %+v, want %+v", got, want)
	}
	// WS2: the additive fields are populated — est_out>0 and a single unit id.
	if got[0].EstOut <= 0 || got[0].EditUnitID != 0 || got[0].OversizedSentence {
		t.Fatalf("WS2 fields = est_out %.3f unit %d oversized %v", got[0].EstOut, got[0].EditUnitID, got[0].OversizedSentence)
	}
}

func TestSplitChunksChapters(t *testing.T) {
	got := SplitChunks([]string{"Глава один.", "Глава два.", "Глава три."}, testSeg())
	want := []Chunk{
		{Chapter: 1, ChunkIdx: 0, Text: "Глава один."},
		{Chapter: 2, ChunkIdx: 0, Text: "Глава два."},
		{Chapter: 3, ChunkIdx: 0, Text: "Глава три."},
	}
	if !reflect.DeepEqual(stripMeta(got), want) {
		t.Fatalf("chapters = %+v, want %+v", got, want)
	}
	// WS2: each chapter is its own edit unit → distinct monotone EditUnitIDs 0,1,2.
	for i, c := range got {
		if c.EditUnitID != i {
			t.Fatalf("chapter %d chunk edit_unit_id = %d, want %d (each short chapter is its own unit)", c.Chapter, c.EditUnitID, i)
		}
	}
}

func TestSplitChunksPacksToTarget(t *testing.T) {
	// Two paragraphs that together exceed the OUTPUT-token budget (each ~1156 out-tokens for
	// cyrillic: 3000 chars × 0.3852) must split into two draft chunks, each kept whole.
	p1 := strings.Repeat("а", 3000)
	p2 := strings.Repeat("б", 3000)
	got := SplitChunks([]string{p1 + "\n\n" + p2}, testSeg())
	want := []Chunk{
		{Chapter: 1, ChunkIdx: 0, Text: p1},
		{Chapter: 1, ChunkIdx: 1, Text: p2},
	}
	if !reflect.DeepEqual(stripMeta(got), want) {
		t.Fatalf("packing = %d chunks, want 2 whole paragraphs; got %+v", len(got), summarize(got))
	}
	// WS2: the two ~1156-out chunks (2311 together ≤ 3200) group into ONE edit unit.
	if got[0].EditUnitID != got[1].EditUnitID {
		t.Fatalf("two sub-ceiling draft chunks must share an edit unit, got ids %d and %d", got[0].EditUnitID, got[1].EditUnitID)
	}
}

func TestSplitChunksPacksSmallParagraphsTogether(t *testing.T) {
	got := SplitChunks([]string{"Абзац один.\n\nАбзац два.\n\nАбзац три."}, testSeg())
	if len(got) != 1 {
		t.Fatalf("small paragraphs must pack into one chunk, got %d: %+v", len(got), summarize(got))
	}
	if got[0].Text != "Абзац один.\n\nАбзац два.\n\nАбзац три." {
		t.Fatalf("packed text = %q", got[0].Text)
	}
}

func TestSplitChunksOversizeParagraphSplitsAtSentences(t *testing.T) {
	// One paragraph far over the output budget, built of many whole sentences, must
	// split into >1 chunk AND never cut a sentence: every original sentence must
	// survive intact and in order across the chunks.
	sentence := strings.Repeat("слово ", 40) + "конец." // ~79 out-tokens each
	var sb strings.Builder
	for i := 0; i < 40; i++ { // ~3160 out-tokens total ≫ 1797
		sb.WriteString(sentence)
		sb.WriteByte(' ')
	}
	para := strings.TrimSpace(sb.String())
	got := SplitChunks([]string{para}, testSeg())
	if len(got) < 2 {
		t.Fatalf("oversize paragraph must split into >1 chunk, got %d", len(got))
	}
	wantSentences := trimAll(splitSourceSentences(para))
	var gotSentences []string
	for _, c := range got {
		gotSentences = append(gotSentences, trimAll(splitSourceSentences(c.Text))...)
	}
	if !reflect.DeepEqual(gotSentences, wantSentences) {
		t.Fatalf("sentence integrity broken:\n got %d sentences\nwant %d sentences", len(gotSentences), len(wantSentences))
	}
}

func TestSplitChunksShortLinesPackByTrueTokenBudget(t *testing.T) {
	// Many short dialogue paragraphs must pack by TRUE output-token content (est_out on the
	// joined text), not by summing a per-unit floor. 150 lines of 「はい。」 are ~533 out-tokens,
	// so they belong in ONE chunk (a per-unit floor would wrongly flush into sub-500-char pieces).
	var paras []string
	for i := 0; i < 150; i++ {
		paras = append(paras, "「はい。」")
	}
	got := SplitChunks([]string{strings.Join(paras, "\n\n")}, testSeg())
	if len(got) != 1 {
		t.Fatalf("short lines must pack by true token budget into 1 chunk, got %d chunks", len(got))
	}
	if got[0].EstOut > testSeg().DraftBudgetOut {
		t.Fatalf("packed chunk overshoots the budget: est_out %.1f", got[0].EstOut)
	}
}

func TestSplitChunksOversizeSentenceIsOwnChunk(t *testing.T) {
	// A single sentence (no internal terminator) larger than the budget can only be its own
	// chunk — never split — and is flagged OversizedSentence (passthrough, WS2 §2б).
	big := strings.Repeat("ы", 6000) // ~2311 out-tokens, one un-terminated run
	got := SplitChunks([]string{big}, testSeg())
	if len(got) != 1 || got[0].Text != big {
		t.Fatalf("an oversize sentence must be one un-split chunk, got %+v", summarize(got))
	}
	if !got[0].OversizedSentence {
		t.Fatalf("a single over-budget sentence must be flagged OversizedSentence")
	}
}

func TestSplitChunksDropsEmpties(t *testing.T) {
	// A blank/whitespace-only chapter vanishes and does NOT consume a chapter number.
	got := SplitChunks([]string{"A", "   \n\n  ", "B"}, testSeg())
	want := []Chunk{
		{Chapter: 1, ChunkIdx: 0, Text: "A"},
		{Chapter: 2, ChunkIdx: 0, Text: "B"},
	}
	if !reflect.DeepEqual(stripMeta(got), want) {
		t.Fatalf("empties = %+v, want %+v", got, want)
	}
}

func TestSplitChunksDeterministic(t *testing.T) {
	chapters := []string{
		"Абзац.\n\nЕщё абзац.",
		strings.Repeat("слово ", 400) + "конец.",
	}
	a, b := SplitChunks(chapters, testSeg()), SplitChunks(chapters, testSeg())
	if !reflect.DeepEqual(a, b) {
		t.Fatal("SplitChunks must be deterministic")
	}
}

func TestSplitChunksNeverSplitsAbbreviation(t *testing.T) {
	// Two long en paragraphs each ending mid-flow with "Mr. Smith" must not chunk
	// between "Mr." and "Smith" — the abbreviation guard keeps the sentence whole.
	p := strings.Repeat("The road went on and on. ", 100) + "It was Mr. Smith who arrived."
	got := SplitChunks([]string{p}, testSeg())
	for _, c := range got {
		if strings.HasSuffix(strings.TrimSpace(c.Text), "Mr.") {
			t.Fatalf("chunk ended at an abbreviation 'Mr.' — a sentence was cut: %q", c.Text)
		}
	}
}

// TestSplitChunksEditUnitGroupsWholeChunks pins the WS2 decoupling: a chapter whose draft chunks
// together exceed EditCeilingOut is grouped into >1 edit unit at DRAFT-CHUNK boundaries (never
// splitting a chunk), so every chunk of a unit reconstructs the unit's draft losslessly.
func TestSplitChunksEditUnitGroupsWholeChunks(t *testing.T) {
	// Three ~1156-out paragraphs (0.3852×3000): draft chunks {p1},{p2},{p3} (each ~1156, two
	// together 2311 ≤ 1797? no — 2311>1797 so each is its own draft chunk). Edit ceiling 3200:
	// group whole chunks → {p1,p2}=2311 ≤3200, +p3=3467>3200 → unit0={p1,p2}, unit1={p3}.
	p := strings.Repeat("а", 3000)
	q := strings.Repeat("б", 3000)
	r := strings.Repeat("в", 3000)
	got := SplitChunks([]string{p + "\n\n" + q + "\n\n" + r}, testSeg())
	if len(got) != 3 {
		t.Fatalf("want 3 draft chunks, got %d", len(got))
	}
	ids := []int{got[0].EditUnitID, got[1].EditUnitID, got[2].EditUnitID}
	if ids[0] != ids[1] || ids[1] == ids[2] {
		t.Fatalf("edit-unit grouping = %v, want first two together and the third separate", ids)
	}
	// Every draft chunk belongs to exactly one unit; a unit's chunks are contiguous.
	for i := 1; i < len(got); i++ {
		if got[i].EditUnitID < got[i-1].EditUnitID {
			t.Fatalf("edit unit ids must be monotone, got %v", ids)
		}
	}
}

// --- source sentence splitter (zh/ja/en) ---------------------------------------

func TestSplitSourceSentencesTilesLosslessly(t *testing.T) {
	// The segments must TILE the input: concatenation reproduces it byte-for-byte.
	inputs := []string{
		"静かな朝。彼は歩いた。「止まれ！」と叫んだ。",
		"他说。“不要！”然后离开了。",
		"Hello world. This is a test!  Really?  Yes.",
		"Mr. Smith met Dr. J. R. R. Brown. They talked.",
		"末尾に句点なし",
		"",
	}
	for _, in := range inputs {
		segs := splitSourceSentences(in)
		if got := strings.Join(segs, ""); got != in {
			t.Fatalf("tiling broke for %q: rejoined to %q", in, got)
		}
	}
}

func TestSplitSourceSentencesJapanese(t *testing.T) {
	// CJK terminators split zero-width; a closing quote is absorbed into its sentence.
	got := trimAll(splitSourceSentences("静かな朝。「止まれ！」と彼は言った。"))
	want := []string{"静かな朝。", "「止まれ！」と彼は言った。"}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("ja split = %q, want %q", got, want)
	}
}

func TestSplitSourceSentencesEnglish(t *testing.T) {
	got := trimAll(splitSourceSentences(`He said "Stop!" She ran. Done.`))
	want := []string{`He said "Stop!"`, "She ran.", "Done."}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("en split = %q, want %q", got, want)
	}
}

func TestSplitSourceSentencesEllipsisNotOverSplitInCJK(t *testing.T) {
	// "……" mid-CJK-sentence must not split; the trailing 。 still ends it.
	got := trimAll(splitSourceSentences("そう……だった。次へ。"))
	want := []string{"そう……だった。", "次へ。"}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("ellipsis split = %q, want %q", got, want)
	}
}

func TestSplitSourceSentencesASCIIRespectsQuoteDepth(t *testing.T) {
	// ASCII terminators inside a tracked quote must NOT split (like the CJK branch):
	// 「Stop. Now.」 is one dialogue unit, not two sentences (external-review #4).
	got := trimAll(splitSourceSentences("「Stop. Now.」次へ。"))
	want := []string{"「Stop. Now.」次へ。"}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("ASCII-in-quote split = %q, want %q (must not over-split dialogue)", got, want)
	}
	// Sanity: outside quotes ASCII still splits.
	if n := len(trimAll(splitSourceSentences("Stop. Now."))); n != 2 {
		t.Fatalf("plain ASCII must still split into 2, got %d", n)
	}
}

func TestSplitSourceSentencesAbbreviations(t *testing.T) {
	// Abbreviations and single-letter initials do not end a sentence.
	got := trimAll(splitSourceSentences("Dr. J. R. R. Tolkien wrote it. The end."))
	want := []string{"Dr. J. R. R. Tolkien wrote it.", "The end."}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("abbrev split = %q, want %q", got, want)
	}
}

// --- property/fuzz (external-review #7: would have caught #3) ----------------

// FuzzSplitSourceSentencesTiles asserts the segmenter is lossless: its segments
// always concatenate back to the input's rune sequence, and byte-exactly when the
// input is valid UTF-8 (the production invariant after NormalizeSource). A path that
// dropped/reordered/replaced a rune — like the size-dependent []rune→U+FFFD bug #3 —
// fails this over random input.
func FuzzSplitSourceSentencesTiles(f *testing.F) {
	for _, s := range []string{
		"静かな朝。「止まれ！」と言った。", "Hello. World!", "Mr. Smith went. Home.",
		"そう……だった。", "「Stop. Now.」", "abc\xe9def", "",
	} {
		f.Add(s)
	}
	f.Fuzz(func(t *testing.T, s string) {
		segs := splitSourceSentences(s)
		joined := strings.Join(segs, "")
		if joined != string([]rune(s)) {
			t.Fatalf("tiling lost/reordered runes: input=%q rejoined=%q", s, joined)
		}
		if utf8.ValidString(s) && joined != s {
			t.Fatalf("valid-UTF-8 tiling not byte-exact: %q → %q", s, joined)
		}
	})
}

// FuzzNormalizeSourceValidUTF8 guards fix #3 directly: NormalizeSource always yields
// valid UTF-8, so the chunker never sees a stray byte the two paths would treat
// differently.
func FuzzNormalizeSourceValidUTF8(f *testing.F) {
	for _, s := range []string{"abc\xe9def", "正常なテキスト", "\xff\xfe", ""} {
		f.Add(s)
	}
	f.Fuzz(func(t *testing.T, s string) {
		if out := NormalizeSource(s); !utf8.ValidString(out) {
			t.Fatalf("NormalizeSource must return valid UTF-8, invalid for %q", s)
		}
	})
}

// FuzzSplitChunksPreservesContent asserts chunking loses no non-space content: the
// count of non-space runes across all chunks equals that of the normalized source.
func FuzzSplitChunksPreservesContent(f *testing.F) {
	for _, s := range []string{
		"静かな朝。\n\n「止まれ！」と彼は言った。次へ。", strings.Repeat("слово ", 500), "abc\xe9def", "",
	} {
		f.Add(s)
	}
	f.Fuzz(func(t *testing.T, s string) {
		norm := NormalizeSource(s)
		want := countNonSpace(norm)
		got := 0
		for _, c := range SplitChunks([]string{norm}, testSeg()) {
			got += countNonSpace(c.Text)
		}
		if got != want {
			t.Fatalf("chunking lost non-space content: got %d want %d (input %q)", got, want, s)
		}
	})
}

func countNonSpace(s string) int {
	n := 0
	for _, r := range s {
		if !unicode.IsSpace(r) {
			n++
		}
	}
	return n
}

// trimAll TrimSpaces each segment and drops empties (test convenience).
func trimAll(segs []string) []string {
	var out []string
	for _, s := range segs {
		if t := strings.TrimSpace(s); t != "" {
			out = append(out, t)
		}
	}
	return out
}

func summarize(cs []Chunk) []string {
	out := make([]string, len(cs))
	for i, c := range cs {
		txt := c.Text
		if len(txt) > 12 {
			txt = txt[:12] + "…"
		}
		out[i] = strings.TrimSpace(txt)
	}
	return out
}