textmachine/backend/internal/pipeline/chunker.go

package pipeline

import (
	"strings"
	"unicode"
)

// chunker.go: the шаг-3a SOURCE segmenter. It turns the per-chapter normalized
// text produced by ingest.go into an ordered list of (chapter, chunk) units for
// the runner loop. WS2 (слой 1) moved the packing budget to OUTPUT (ru) tokens via
// the fertility estimate (est_out = f_cjk·cjk + f_other·other over source char-classes),
// decoupling draft sizing (small chunk, COGS/coverage/alignment) from the EDIT unit
// (coarse chapter-scale unit for cross-chunk cohesion): pack whole paragraphs to
// DraftBudgetOut, and when a single paragraph exceeds it descend to SENTENCE boundaries
// (never splitting a sentence; a lone over-budget sentence is a passthrough
// OversizedSentence chunk). Edit units then group WHOLE draft chunks to EditCeilingOut.
//
// Contract kept intact for the loop (bookrun.go/chunkrun.go depend on it): SplitChunks emits
// Chunk{Chapter,ChunkIdx,Text,…}; ChunkIdx resets to 0 per chapter; the new fields (EstOut,
// OversizedSentence, EditUnitID) are ADDITIVE; the function is a PURE, deterministic function of
// (chapters, SegBudget) (no time/rand, no map iteration). Its behavior is versioned by
// chunkerVersion (render.go); the SegBudget is snapshot-folded via segmentationSnap — so changing
// the segmentation OR the budget/fertility is a loud --resnapshot re-translation, never a silent
// cache miss (§3.4/R6/D5.2/D30.9).
//
// IMPORTANT: this sentence splitter is a SEPARATE, independent piece of logic from
// coverage.go's splitSentences. That one is a bit-for-bit port of the Python oracle
// (refusal_bench.py) over the RUSSIAN OUTPUT and is Python-locked (D12 Q1). THIS one
// segments the SOURCE (zh/ja/en), is versioned by chunkerVersion (Р2 = code), and is
// free to be more faithful (closing-bracket absorption, an abbreviation guard). Do
// not couple them.

// Chunk is one ordered unit of translation work.
type Chunk struct {
	Chapter  int // 1-based
	ChunkIdx int // 0-based within its chapter
	Text     string
	// EstOut is the fertility-estimated OUTPUT (ru) token cost of this chunk (WS2, слой 1):
	// est_out = fert.CJK·cjk_src + fert.Other·other_src over the source char-classes. It is the
	// unit the draft budget packs against (DraftBudgetOut), decoupling draft sizing from the
	// old input-token heuristic (L2-budget-wrong-unit). Additive field; the {Chapter,ChunkIdx,Text}
	// contract the loop/ingest/status depend on is untouched.
	EstOut float64
	// OversizedSentence marks a chunk that is a SINGLE source sentence whose est_out alone exceeds
	// DraftBudgetOut (WS2 §2б). It is passthrough (never clause-split — clause splitting is
	// boundary machinery, anti-scope), and the flag makes the un-budgetable unit OBSERVABLE
	// ("деградация не молчит") rather than silently oversized.
	OversizedSentence bool
	// EditUnitID is the 0-based book-global id of the coarse EDIT unit this draft chunk belongs to
	// (WS2 decoupling: draft = small chunk, edit = large unit). An edit unit is a greedy grouping of
	// WHOLE draft chunks up to EditCeilingOut (per chapter), so W2's editor reads a unit's draft as
	// the concatenation of its member chunks' W1 outputs — clean reconstruction by construction (no
	// straddle; verified: grouping whole chunks reproduces the paragraph-packed count 37 with 0
	// straddle on the rerun corpus). Draft chunks never cross an edit-unit boundary.
	EditUnitID int
}

// SegBudget is the resolved output-token segmentation budget (WS2 §2в): the draft-chunk and
// edit-unit ceilings in ru-OUTPUT tokens, plus the per-pair fertility coefficients that convert
// source char-classes to an output-token estimate. Resolved from config.Segmentation by the runner
// and snapshot-folded via segmentationSnap, so a budget/fertility edit is a loud --resnapshot
// (D30.9). The classifier (tokenClassCounts) uses the REAL unicode.RangeTable of EstimateTokens
// (Han|Hiragana|Katakana|Hangul) for generality (§0.1) — NOT script-specific ord ranges.
type SegBudget struct {
	DraftBudgetOut float64 // draft-chunk ceiling, ru-output tokens (default 1797)
	EditCeilingOut float64 // edit-unit ceiling, ru-output tokens (default 3200, gated arm at 8000)
	FertCJK        float64 // output tokens per CJK source char (default 1.1978)
	FertOther      float64 // output tokens per non-CJK/non-space source char (default 0.3852)
}

// estOut applies the fertility formula to pre-counted source char-classes — identical to
// est_out over the joined chunk text (whitespace is class-free, as tokenClassCounts drops it).
func (s SegBudget) estOut(cjk, other int) float64 {
	return s.FertCJK*float64(cjk) + s.FertOther*float64(other)
}

// SplitChunks segments the ordered, per-chapter NORMALIZED text of a Document
// (ingest.go already applied NormalizeSource and split chapters) into an ordered
// chunk list. Chapters that yield no text (blank/whitespace-only, e.g. an epub
// cover/nav page) do NOT consume a chapter number, so numbering stays dense over
// real content (Фаза-0 backward compat). Fully deterministic. Each chunk carries its
// EstOut + OversizedSentence flag + EditUnitID (WS2); the edit-unit id is monotone across
// the book so a W2 unit is uniquely addressable.
func SplitChunks(chapters []string, seg SegBudget) []Chunk {
	var out []Chunk
	chapterNo := 0
	editUnitID := 0
	for _, chapText := range chapters {
		paras := splitParagraphs(chapText)
		if len(paras) == 0 {
			continue // an empty chapter does not consume a chapter number
		}
		chapterNo++
		chapterChunks := chapterDraftChunks(chapterNo, paras, seg)
		assignEditUnits(chapterChunks, seg, &editUnitID)
		out = append(out, chapterChunks...)
	}
	return out
}

// chapterDraftChunks packs one chapter's paragraphs into DRAFT chunks (the fine tiling). Rule:
// pack whole paragraphs while the running OUTPUT-token estimate stays within DraftBudgetOut
// (prefer paragraph boundaries); a paragraph that alone exceeds the budget is flushed on its own
// and then SENTENCE-packed (never splitting a sentence; a single sentence over budget is its own
// OversizedSentence chunk). The running budget is tracked as ADDITIVE char-class counts (CJK vs
// other), applying est_out ONCE to the whole chunk — the "\n\n" separators are whitespace, which
// tokenClassCounts drops, so counting classes and estimating once reproduces est_out(joined text)
// exactly. Each emitted chunk gets its ChunkIdx (per-chapter) + EstOut; EditUnitID is set later.
func chapterDraftChunks(chapterNo int, paras []string, seg SegBudget) []Chunk {
	var out []Chunk
	chunkIdx := 0
	emit := func(text string, oversized bool, cjk, other int) {
		if t := strings.TrimSpace(text); t != "" {
			out = append(out, Chunk{
				Chapter: chapterNo, ChunkIdx: chunkIdx, Text: t,
				EstOut: seg.estOut(cjk, other), OversizedSentence: oversized,
			})
			chunkIdx++
		}
	}
	var buf []string // paragraphs accumulated for the current chunk
	var bufCJK, bufOther int
	flush := func() {
		if len(buf) == 0 {
			return
		}
		emit(strings.Join(buf, "\n\n"), false, bufCJK, bufOther)
		buf = buf[:0]
		bufCJK, bufOther = 0, 0
	}
	for _, p := range paras {
		pCJK, pOther := tokenClassCounts(p)
		if seg.estOut(pCJK, pOther) > seg.DraftBudgetOut {
			// Oversize paragraph: close any open chunk at this paragraph boundary,
			// then split the paragraph at sentence boundaries into its own chunks.
			flush()
			for _, sub := range packSentences(p, seg) {
				emit(sub.text, sub.oversizedSentence, sub.cjk, sub.other)
			}
			continue
		}
		// Normal paragraph: start a new chunk if adding it would exceed the budget
		// and the current chunk is non-empty (never split a paragraph here).
		if len(buf) > 0 && seg.estOut(bufCJK+pCJK, bufOther+pOther) > seg.DraftBudgetOut {
			flush()
		}
		buf = append(buf, p)
		bufCJK, bufOther = bufCJK+pCJK, bufOther+pOther
	}
	flush()
	return out
}

// assignEditUnits groups a chapter's already-built DRAFT chunks into EDIT units — greedy packing of
// WHOLE draft chunks until the accumulated est_out would exceed EditCeilingOut — and stamps each
// chunk's EditUnitID (monotone across the book via *nextID). Because units are formed from whole
// draft chunks, a unit's source/draft span is exactly the concatenation of its member chunks (no
// straddle) and W2 reconstructs the unit's draft losslessly. A single draft chunk larger than the
// ceiling (an OversizedSentence, or a chapter whose first chunk already exceeds the ceiling) is its
// own unit — grouping never splits a chunk (never-split-below-chunk mirrors never-split-paragraph).
func assignEditUnits(chunks []Chunk, seg SegBudget, nextID *int) {
	acc := 0.0
	started := false
	for i := range chunks {
		if started && acc+chunks[i].EstOut > seg.EditCeilingOut {
			*nextID++ // close the current unit at this chunk boundary
			acc = 0
		}
		chunks[i].EditUnitID = *nextID
		acc += chunks[i].EstOut
		started = true
	}
	if started {
		*nextID++ // advance past the last unit so the next chapter starts a fresh id
	}
}

// subChunk is one sentence-packed sub-chunk with its char-class counts and the oversized-sentence
// flag (a single sentence whose est_out alone exceeds DraftBudgetOut).
type subChunk struct {
	text              string
	cjk, other        int
	oversizedSentence bool
}

// packSentences splits one oversize paragraph into sentence segments (lossless tiling —
// concatenation reproduces the paragraph) and packs consecutive segments into ≤ DraftBudgetOut
// (est_out) sub-chunks, never splitting a sentence. A single sentence larger than the budget
// becomes its own sub-chunk flagged OversizedSentence (passthrough, never clause-split — the plan's
// §2б decision: "деградация не молчит" is satisfied by the FLAG, not emergency clause machinery).
func packSentences(paragraph string, seg SegBudget) []subChunk {
	segs := splitSourceSentences(paragraph)
	var chunks []subChunk
	var buf strings.Builder
	var bufCJK, bufOther int
	flush := func() {
		if buf.Len() == 0 {
			return
		}
		chunks = append(chunks, subChunk{text: buf.String(), cjk: bufCJK, other: bufOther})
		buf.Reset()
		bufCJK, bufOther = 0, 0
	}
	for _, s := range segs {
		sCJK, sOther := tokenClassCounts(s)
		// A lone sentence over budget: flush what precedes it, then emit it as its own
		// OversizedSentence sub-chunk (never split a sentence, and never merge it with a
		// neighbour so the oversized flag stays attributable to the single sentence).
		if seg.estOut(sCJK, sOther) > seg.DraftBudgetOut {
			flush()
			chunks = append(chunks, subChunk{text: s, cjk: sCJK, other: sOther, oversizedSentence: true})
			continue
		}
		if buf.Len() > 0 && seg.estOut(bufCJK+sCJK, bufOther+sOther) > seg.DraftBudgetOut {
			flush()
		}
		buf.WriteString(s)
		bufCJK, bufOther = bufCJK+sCJK, bufOther+sOther
	}
	flush()
	return chunks
}

// tokenClassCounts counts the two token-estimate character classes of EstimateTokens
// (render.go): CJK ideographs/kana/hangul, and "other" non-space runes (whitespace
// folds into neighbours). Kept separate from EstimateTokens so the packer can sum
// counts ADDITIVELY across units and floor once at the chunk level. Mirrors
// EstimateTokens' classes; a recalibration there that changes packing granularity
// is itself a chunkerVersion bump (both are behaviour changes covered by the snapshot).
func tokenClassCounts(s string) (cjk, other int) {
	for _, r := range s {
		switch {
		case unicode.In(r, unicode.Han, unicode.Hiragana, unicode.Katakana, unicode.Hangul):
			cjk++
		case unicode.IsSpace(r):
		default:
			other++
		}
	}
	return cjk, other
}

// splitParagraphs breaks a chapter into paragraphs on blank lines and trims each,
// dropping empties. Deterministic; whitespace-only paragraphs vanish. (Ingest emits
// a blank line between block-level epub elements, so paragraphs survive extraction.)
func splitParagraphs(chapter string) []string {
	var paras []string
	for _, raw := range strings.Split(chapter, "\n\n") {
		if p := strings.TrimSpace(raw); p != "" {
			paras = append(paras, p)
		}
	}
	return paras
}

// --- source sentence segmentation (zh/ja/en) — chunkerVersion-versioned ---------

// splitSourceSentences segments one paragraph at sentence boundaries, returning
// substrings that TILE the input (their concatenation is exactly the paragraph —
// no character is lost or moved), so a packed sub-chunk is an exact source slice.
// A boundary is a run of sentence terminators, plus any closing quotes/brackets, at:
//
//	(1) a CJK terminator (。！？) — zero-width, CJK prose has no inter-sentence space —
//	    but ONLY at quote-depth 0. A terminator inside a quote is dialogue-internal:
//	    「止まれ！」と言った。 is ONE sentence ending at the outer 。, not two (the same
//	    for zh “…！”）— this is the dialogue case §3.7/D12-Q1 defers for the coverage
//	    oracle, but our source splitter is free to get it right;
//	(2) an ASCII terminator (.!?) or ellipsis (…) — only when followed by whitespace
//	    or end-of-paragraph (so a decimal/URL dot mid-token does not split, and en
//	    dialogue "Stop!" She ran. still splits on the space after the quote).
//
// Trailing whitespace after a boundary is attached to the LEFT segment (the next
// segment starts at the next content char), keeping the tiling exact. An
// abbreviation guard suppresses a split after a lone "." following a known
// abbreviation or a single-letter initial (never split "Mr. Smith" / "J. R. R.").
// When unsure it prefers NOT to split (over-merge is safe — units just pack
// together; over-split cuts a sentence, the forbidden case). Unbalanced quotes are
// tolerated (depth floors at 0), erring toward merge.
func splitSourceSentences(paragraph string) []string {
	runes := []rune(paragraph)
	var segs []string
	start, i, depth := 0, 0, 0
	for i < len(runes) {
		r := runes[i]
		if isOpenQuote(r) {
			depth++
			i++
			continue
		}
		if isCloseQuote(r) {
			if depth > 0 {
				depth--
			}
			i++
			continue
		}
		if !isSourceTerminator(r) {
			i++
			continue
		}
		// Consume a run of terminators (collapse "。。。" / "!?" into one boundary).
		j := i
		cjk := false
		for j < len(runes) && isSourceTerminator(runes[j]) {
			if isCJKTerminator(runes[j]) {
				cjk = true
			}
			j++
		}
		// The terminator's quote-depth is the depth HERE — before absorbing the
		// closing quote that may follow it (「…！」: the ！ is at depth 1, the 」 that
		// closes the quote comes after and must not retro-promote it to depth 0).
		depthHere := depth
		// Absorb a run of closing quotes/brackets that belong to the sentence,
		// decrementing quote-depth for the depth-tracked closers among them.
		e := j
		for e < len(runes) && isClosingBracket(runes[e]) {
			if isCloseQuote(runes[e]) && depth > 0 {
				depth--
			}
			e++
		}
		boundary := false
		if cjk {
			boundary = depthHere == 0 // dialogue-internal terminators do not end a sentence
		} else {
			// ASCII terminator / ellipsis: a boundary only outside a tracked quote AND
			// before whitespace or EOS. The quote-depth guard must apply here too, not
			// just to the CJK branch — otherwise 「Stop. Now.」 over-splits mid-dialogue
			// (external-review #4). Straight-quote en ("Stop!" She ran.) is untracked
			// (depth 0), so it still splits on the space as before.
			boundary = depthHere == 0 && (e >= len(runes) || isASCIISpace(runes[e]))
			// Abbreviation guard: a lone "." after an abbreviation/initial is not a end.
			if boundary && j-i == 1 && runes[i] == '.' && isAbbrevBefore(runes[start:i]) {
				boundary = false
			}
		}
		if !boundary {
			i = e
			continue
		}
		// Attach the following whitespace run to the left segment (keeps tiling exact).
		w := e
		for w < len(runes) && isASCIISpace(runes[w]) {
			w++
		}
		segs = append(segs, string(runes[start:w]))
		start, i = w, w
	}
	if start < len(runes) {
		segs = append(segs, string(runes[start:]))
	}
	return segs
}

// isOpenQuote / isCloseQuote are the DEPTH-TRACKED paired quote/bracket delimiters
// (CJK brackets + curly quotes with distinct open/close glyphs). Straight quotes
// (" ') are NOT tracked — same glyph opens and closes, so depth is undecidable; en
// relies on the whitespace rule instead, which is correct for spaced prose.
func isOpenQuote(r rune) bool {
	switch r {
	case '「', '『', '（', '【', '《', '〈', '［', '｛', '“', '‘':
		return true
	}
	return false
}

func isCloseQuote(r rune) bool {
	switch r {
	case '」', '』', '）', '】', '》', '〉', '］', '｝', '”', '’':
		return true
	}
	return false
}

// isSourceTerminator is the sentence-terminator class for the SOURCE: CJK 。！？,
// ASCII .!?, and the ellipsis … (U+2026).
func isSourceTerminator(r rune) bool {
	switch r {
	case '。', '！', '？', '.', '!', '?', '…':
		return true
	}
	return false
}

// isCJKTerminator is the zero-width-splitting subset (fullwidth CJK). The ellipsis
// is deliberately NOT here: "……" mid-CJK-sentence must not split, but "……。" still
// splits on the 。 in the same run.
func isCJKTerminator(r rune) bool {
	switch r {
	case '。', '！', '？':
		return true
	}
	return false
}

// isClosingBracket is the closing-quote/bracket class absorbed into the left
// sentence after its terminator (ja/zh brackets + straight/curly quotes).
func isClosingBracket(r rune) bool {
	switch r {
	case '」', '』', '】', '）', ')', '》', '〉', '］', ']', '｝', '}',
		'"', '\'', '”', '’':
		return true
	}
	return false
}

// isASCIISpace is the whitespace used for ASCII-terminator boundary detection and
// tiling — the ordinary prose separators (space, tab, newline, CR, form feed,
// vertical tab, and the ideographic space U+3000).
func isASCIISpace(r rune) bool {
	switch r {
	case ' ', '\t', '\n', '\r', '\f', '\v', '　':
		return true
	}
	return false
}

// sourceAbbrevs are trailing tokens after which a lone "." is treated as an
// abbreviation, not a sentence end (case-insensitive). A pragmatic set for prose;
// the acceptance path is ja→ru where CJK terminators dominate, so this only guards
// the en source. Single-letter initials are handled separately (isAbbrevBefore).
var sourceAbbrevs = map[string]bool{
	"mr": true, "mrs": true, "ms": true, "dr": true, "prof": true, "st": true,
	"jr": true, "sr": true, "vs": true, "no": true, "vol": true, "ch": true,
	"fig": true, "col": true, "gen": true, "sgt": true, "capt": true, "lt": true,
	"rev": true, "gov": true, "sen": true, "rep": true, "etc": true, "inc": true,
	"ltd": true, "co": true, "mt": true, "ave": true, "rd": true,
}

// isAbbrevBefore reports whether the text immediately before a lone "." ends in a
// known abbreviation or a single-letter initial (so the "." is not a boundary).
func isAbbrevBefore(left []rune) bool {
	// Trailing run of ASCII letters (the "word" before the period).
	end := len(left)
	k := end
	for k > 0 && isASCIILetter(left[k-1]) {
		k--
	}
	word := left[k:end]
	if len(word) == 0 {
		return false
	}
	// Single-letter initial ("J.") — only when it stands alone (start or preceded by
	// a space/opening punctuation), never a one-letter word ending a real sentence
	// mid-flow would be rare; over-merge is the safe side.
	if len(word) == 1 {
		return k == 0 || !isASCIILetter(left[k-1])
	}
	return sourceAbbrevs[strings.ToLower(string(word))]
}

func isASCIILetter(r rune) bool {
	return (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z')
}