#!/usr/bin/env python3
"""S0 — deterministic pre-processing and chunking. No LLM, no cost.

Handles the three transcript shapes the pipeline actually produces:
  qwen3cam   `【谭凯】…` / `【speaker_3】…` turns, blank-line separated, plus a
             `<!-- voiceprint_session: … -->` header line.
  qwen3-plain `--- 第N段 [00:00 - 09:00] ---` segment banners then a wall of text
             with no speaker information at all.
  vibevoice  `[mm:ss] S0: …` lines, with `[复读截断]` / `·` noise artifacts.

Chunking rule: never split inside a speaker turn if avoidable, because the
cleaning model needs the whole turn to tell a filler word from a real answer.
Turns longer than the budget are split on sentence boundaries. Each chunk carries
`overlap_prefix` — the tail of the previous chunk, given to the model as
read-only context so it can resolve a pronoun or a name that started earlier
without being asked to re-emit it.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass, field

VOICEPRINT_HDR = re.compile(r"^<!--\s*voiceprint_session:")
SEGMENT_BANNER = re.compile(r"^---\s*第\d+段.*---\s*$")
TURN_TAG = re.compile(r"^【([^】]+)】\s*(.*)$", re.S)
VIBE_TAG = re.compile(r"^\[(\d{1,2}:\d{2})\]\s*(S\d+)\s*[:：]\s*(.*)$", re.S)
# Must include ASCII terminators: the qwen3-plain backend emits English stretches
# ("...so that's why they wanted us to set up company.") with no CJK punctuation at
# all, and a CJK-only rule left single 15k-char "sentences" that blew the budget.
SENT_END = re.compile(r"(?<=[。！？；;?!])(?![”’\"'）)])")
# Last-resort split points when a run has no terminator for thousands of chars.
SOFT_BREAK = re.compile(r"(?<=[，,、：:])")
NOISE = re.compile(r"\[复读截断\]|^[·\s]+$")

# ASR decoders loop when fed non-speech audio (hold music, a phone beep, dead air).
# Measured on the fixture set: one 24k-char call transcript is 77% "beep beep beep…"
# (3,698 repetitions) and a 58k-char meeting has a 3.7k-char run of "呃，". Left in
# place this garbage consumed 11 of 17 chunks, i.e. ~65% of the LLM spend on that
# recording, and provoked the cleaning model into emitting 26k chars of noise per
# chunk. Collapsing it is free, deterministic, and strictly beneficial.
DEGENERATE = re.compile(
    r"((?:[A-Za-z]{1,12}|[\u4e00-\u9fff]{1,4})[ ，,、。!?！？]{0,2})\1{7,}")
DEGEN_MIN_SPAN = 120        # only collapse runs long enough to be a real artifact
DEGEN_KEEP = 2              # keep a couple of instances so the context stays readable


def collapse_degenerate(text: str) -> tuple[str, list[dict]]:
    """Replace long ASR repetition loops with a short, explicit marker.

    Legitimate speech repetition (对对对, 好的好的) is short and must survive, so the
    rule requires both a high repeat count and a span over DEGEN_MIN_SPAN chars.
    The marker stays in the text rather than deleting silently: a recording with
    four minutes of hold music is a fact worth surfacing in the draft.
    """
    events: list[dict] = []

    def sub(m):
        span, unit = m.group(0), m.group(1)
        if len(span) < DEGEN_MIN_SPAN:
            return span
        n = len(span) // max(1, len(unit))
        events.append({"unit": unit, "repeats": n, "chars": len(span)})
        return unit * DEGEN_KEEP + f"[…重复噪声已折叠：{unit.strip()}×{n}…]"

    return DEGENERATE.sub(sub, text), events


@dataclass
class Turn:
    speaker: str      # "" when the backend gives no speaker
    text: str
    idx: int
    ts: str = ""


@dataclass
class Chunk:
    idx: int
    turns: list[Turn]
    overlap_prefix: str = ""
    meta: dict = field(default_factory=dict)

    @property
    def body(self) -> str:
        out = []
        for t in self.turns:
            out.append(f"【{t.speaker}】{t.text}" if t.speaker else t.text)
        return "\n\n".join(out)

    @property
    def chars(self) -> int:
        return len(self.body)


def detect_backend(raw: str) -> str:
    head = raw[:500]
    if "CAM++" in head or TURN_TAG.match(raw.strip().splitlines()[-1] if raw.strip() else ""):
        return "qwen3cam"
    if re.search(r"^\[\d{1,2}:\d{2}\]\s*S\d+\s*[:：]", raw, re.M):
        return "vibevoice"
    if SEGMENT_BANNER.search(raw):
        return "qwen3-plain"
    return "qwen3cam" if "【" in raw else "qwen3-plain"


def parse(raw: str) -> tuple[list[Turn], dict]:
    """Split a transcript into turns plus header metadata."""
    meta: dict = {
        "voiceprint_session": "", "voiceprint_named": {}, "banners": [],
        "intake": {},
    }
    lines = raw.splitlines()
    body_lines = []
    for ln in lines:
        if ln.startswith("<!-- intake:") and ln.endswith("-->"):
            try:
                meta["intake"] = json.loads(
                    ln[len("<!-- intake:"):-len("-->")].strip())
            except (TypeError, ValueError):
                pass
            continue
        if VOICEPRINT_HDR.match(ln):
            m = re.search(r"voiceprint_session:\s*(\S+)", ln)
            if m:
                meta["voiceprint_session"] = m.group(1)
            m = re.search(r"clusters:\s*(\d+)", ln)
            if m:
                meta["clusters"] = int(m.group(1))
            continue
        if ln.startswith("# 转写"):
            for pair in re.findall(r"(speaker_\d+)=([^,\s]+)", ln):
                meta["voiceprint_named"][pair[0]] = pair[1]
            continue
        if SEGMENT_BANNER.match(ln):
            meta["banners"].append(ln.strip())
            continue
        body_lines.append(ln)

    backend = detect_backend(raw)
    meta["backend"] = backend
    turns: list[Turn] = []

    if backend == "vibevoice":
        for ln in body_lines:
            m = VIBE_TAG.match(ln.strip())
            if m:
                txt = NOISE.sub("", m.group(3)).strip()
                if txt:
                    turns.append(Turn(m.group(2), txt, len(turns), m.group(1)))
        return turns, meta

    blocks = [b.strip() for b in re.split(r"\n\s*\n", "\n".join(body_lines)) if b.strip()]
    for b in blocks:
        m = TURN_TAG.match(b)
        if m:
            txt = NOISE.sub("", m.group(2)).strip()
            if txt:
                turns.append(Turn(m.group(1), txt, len(turns)))
        else:
            txt = NOISE.sub("", b).strip()
            if txt:
                turns.append(Turn("", txt, len(turns)))
    return turns, meta


def _pack(pieces: list[str], budget: int) -> list[str]:
    out, cur = [], ""
    for s in pieces:
        if len(cur) + len(s) > budget and cur:
            out.append(cur)
            cur = s
        else:
            cur += s
    if cur:
        out.append(cur)
    return out


def _split_long(text: str, budget: int) -> list[str]:
    """Break an over-long single turn, escalating through split granularities.

    A hard character cut is the final fallback. It is ugly, but an oversized chunk
    is worse: it is exactly the condition that makes the cleaning model start
    summarizing instead of rewriting, which is the failure mode this whole
    pipeline exists to prevent.
    """
    out = _pack([s for s in SENT_END.split(text) if s], budget)
    if all(len(p) <= budget for p in out):
        return out or [text]
    stage2: list[str] = []
    for p in out:
        if len(p) <= budget:
            stage2.append(p)
        else:
            stage2 += _pack([s for s in SOFT_BREAK.split(p) if s], budget)
    final: list[str] = []
    for p in stage2:
        if len(p) <= budget:
            final.append(p)
        else:
            final += [p[i:i + budget] for i in range(0, len(p), budget)]
    return final or [text]


def chunk(raw: str, budget: int = 1800, overlap_chars: int = 200,
          collapse_noise: bool = True) -> tuple[list[Chunk], dict]:
    events: list[dict] = []
    if collapse_noise:
        raw, events = collapse_degenerate(raw)
    turns, meta = parse(raw)
    meta["degenerate_events"] = events
    meta["degenerate_chars"] = sum(e["chars"] for e in events)
    chunks: list[Chunk] = []
    cur: list[Turn] = []
    size = 0

    def flush():
        nonlocal cur, size
        if cur:
            chunks.append(Chunk(len(chunks), cur, meta={}))
            cur, size = [], 0

    for t in turns:
        if len(t.text) > budget:
            flush()
            for i, piece in enumerate(_split_long(t.text, budget)):
                chunks.append(Chunk(len(chunks), [Turn(t.speaker, piece, t.idx, t.ts)],
                                    meta={"split_turn": True, "part": i}))
            continue
        if size + len(t.text) > budget and cur:
            flush()
        cur.append(t)
        size += len(t.text)
    flush()

    if overlap_chars:
        for i in range(1, len(chunks)):
            prev = chunks[i - 1].body
            chunks[i].overlap_prefix = prev[-overlap_chars:]
    meta["n_turns"] = len(turns)
    meta["n_chunks"] = len(chunks)
    meta["speakers"] = sorted({t.speaker for t in turns if t.speaker})
    return chunks, meta


if __name__ == "__main__":
    import json
    import sys
    from pathlib import Path

    for p in sys.argv[1:]:
        raw = Path(p).read_text(errors="ignore")
        for budget in (900, 1800, 3600):
            cs, meta = chunk(raw, budget)
            sizes = [c.chars for c in cs]
            print(f"{Path(p).name[:40]:<42} budget={budget:<5} chunks={len(cs):<4} "
                  f"turns={meta['n_turns']:<5} spk={len(meta['speakers']):<3} "
                  f"backend={meta['backend']:<12} "
                  f"size min/avg/max={min(sizes)}/{sum(sizes)//len(sizes)}/{max(sizes)}")