#!/usr/bin/env python3
"""exp15 — byte-faithful Python PORT of the backend memory-bank hot path (memory.go + memnorm.go).

Reconstructs the injection the Go pipeline actually feeds per chunk: normalize -> Aho-Corasick
(here: exact multi-substring, same result set) -> spaced-phonetic boundary drop -> trust-gated
longest-match suppression -> spoiler window -> three-way disposition -> sticky scene-inertia (depth 2)
-> priority token budget -> render (translator glossary block / editor constraint block).

WHY a port (orchestrator verdict 2026-07-17): retrieval_state.injected_ids carries EXACT hits only
(sticky-blind by construction), so faithful arm injection needs the full Select, VALIDATED against Go's
actually-rendered glossary block (LOG_LLM_BODIES byte-diff), NOT against injected_ids. On divergence:
trust Go (fix the port). This module is the injection engine shared by A0-validation and all arms.

Source of truth mirrored (backend/internal/pipeline/):
  memnorm.go  normalizeSourceKey / significantLen / trad2simp
  memory.go   Select / suppressContained / matchTrust / trustRank / dispositionFor / spoilerBlocked
              / priorityRank / glossaryLineTokens / renderGlossaryBlock / renderEditorConstraintBlock
              / stickyDepth=2 / minKeyLenHan=2 / minKeyLenPhonetic=3
  memseed.go  loadGlossarySeed (status ""->approved; surfaces = src + aliases when dst non-empty)

Self-review by EXECUTION: run `python mem_select.py --selftest` for invariants; the real gate is the
byte-diff vs Go's logged block in validate_injection.py.
"""
from __future__ import annotations

import sys
import unicodedata
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import regex  # \p{Han} etc. — exact Unicode script tables, matching Go's unicode.Han/Hiragana/...
import yaml

BACKEND = Path("/home/ubuntu/projects/textmachine/backend")
TRAD2SIMP_FILE = BACKEND / "internal/pipeline/data/trad2simp.txt"

MIN_KEY_LEN_HAN = 2
MIN_KEY_LEN_PHONETIC = 3
STICKY_DEPTH = 2

# disposition constants (memory.go:89-93)
CONFIRMED = "confirmed"
AMBIGUOUS = "ambiguous"
REJECT = "reject"

GLOSSARY_HEADER = ("ГЛОССАРИЙ (используй эти утверждённые переводы имён и терминов последовательно; "
                   "строки с пометкой ⟨проверить⟩ — неподтверждённые кандидаты):")
EDITOR_HEADER = ("КАНОНИЧЕСКИЕ ПЕРЕВОДЫ имён и терминов (в черновике эти сущности должны быть переданы "
                 "ИМЕННО этими формами — приводи к ним любые расхождения, склоняя по контексту; не вводи "
                 "иных вариантов и не меняй ничего другого):")

# --- Unicode script predicates (match Go unicode.Han|Hiragana|Katakana|Hangul + IsLetter/Cf) --------
_RE_CJK = regex.compile(r"[\p{Han}\p{Hiragana}\p{Katakana}\p{Hangul}]")
_RE_LETTER = regex.compile(r"\p{L}")
_RE_HAN = regex.compile(r"\p{Han}")
_RE_LATIN = regex.compile(r"\p{Latin}")
_RE_CYRILLIC = regex.compile(r"\p{Cyrillic}")


def _is_cjk(ch: str) -> bool:
    return _RE_CJK.match(ch) is not None


def _is_han(ch: str) -> bool:
    return _RE_HAN.match(ch) is not None


def _is_letter(ch: str) -> bool:
    return _RE_LETTER.match(ch) is not None


def _is_apostrophe_variant(ch: str) -> bool:
    return ch in ("‘", "’", "ʼ", "＇")  # ' ' ʼ ＇  (memnorm.go:189-195)


def _is_ignorable_for_match(ch: str) -> bool:
    # unicode.Cf OR variation selectors (memnorm.go:207-212)
    if unicodedata.category(ch) == "Cf":
        return True
    o = ord(ch)
    return (0xFE00 <= o <= 0xFE0F) or (0xE0100 <= o <= 0xE01EF)


# --- trad2simp table (loaded from the Go embed file, parsed identically) ----------------------------
def _load_trad2simp() -> dict:
    m = {}
    for line in TRAD2SIMP_FILE.read_text(encoding="utf-8").split("\n"):
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        fields = line.split()
        if len(fields) != 2:
            raise ValueError(f"trad2simp corrupt line: {line!r}")
        tr, si = fields
        if len(tr) != 1 or len(si) != 1:
            raise ValueError(f"trad2simp non-single-rune: {line!r}")
        m[tr] = si
    return m


_TRAD2SIMP = _load_trad2simp()


def normalize_source_key(s: str) -> str:
    """memnorm.go:104 — NFKC -> apostrophe fold -> drop ignorables -> trad2simp -> kana->hira -> lower."""
    s = unicodedata.normalize("NFKC", s)
    out = []
    for ch in s:
        if _is_apostrophe_variant(ch):
            ch = "'"
        if _is_ignorable_for_match(ch):
            continue
        ch = _TRAD2SIMP.get(ch, ch)
        o = ord(ch)
        if 0x30A1 <= o <= 0x30F6:  # katakana -> hiragana
            ch = chr(o - 0x60)
        out.append(ch.lower())
    return "".join(out)


def significant_len(normalized: str) -> int:
    """memnorm.go:220 — count CJK/kana/Hangul ideographs + letters."""
    n = 0
    for ch in normalized:
        if _is_cjk(ch) or _is_letter(ch):
            n += 1
    return n


def runes_any_han(s: str) -> bool:
    return _RE_HAN.search(s) is not None


def min_key_len_for(norm_key: str) -> int:
    return MIN_KEY_LEN_HAN if runes_any_han(norm_key) else MIN_KEY_LEN_PHONETIC


def collision_prone_key(norm_key: str) -> bool:
    return (not runes_any_han(norm_key)) and significant_len(norm_key) <= MIN_KEY_LEN_PHONETIC


def glossary_line_tokens(src: str, dst: str) -> int:
    """memory.go:304 — cjk + other//3 over 'src → dst' (floor-free)."""
    cjk = other = 0
    for ch in f"{src} → {dst}":
        if _is_cjk(ch):
            cjk += 1
        elif ch.isspace():
            pass
        else:
            other += 1
    return cjk + other // 3


def trust_rank(d: str) -> int:
    return {CONFIRMED: 2, AMBIGUOUS: 1}.get(d, 0)


# --- data model -------------------------------------------------------------------------------------
@dataclass
class MemoryEntry:
    id: str
    src: str
    dst: str
    status: str
    sense: str
    since_ch: int
    until_ch: int
    norm_keys: list = field(default_factory=list)
    allow_short: bool = False


@dataclass
class Picked:
    entry: MemoryEntry
    via: str          # firing key, or "sticky"
    disp: str


@dataclass
class Selection:
    injected: list = field(default_factory=list)   # [Picked] in budget priority order
    rejected: list = field(default_factory=list)
    evicted: list = field(default_factory=list)
    trust_gated: list = field(default_factory=list)
    active_ids: dict = field(default_factory=dict)  # id -> disp (EXACT only; feeds next sticky)


class MemoryBank:
    """Materialized frozen glossary. Built once from seed rows; never mutated (memory.go:180)."""

    def __init__(self, rows: list):
        self.entries: list[MemoryEntry] = []
        self.key_owners: dict[str, list[int]] = {}
        seen_key = set()
        all_keys = []
        # frozen entry order = store's GlossaryForBook: ORDER BY src, sense, since_ch, until_ch, status, dst
        # (SQLite BINARY collation == UTF-8 byte order == codepoint order for BMP CJK). This fixes the
        # injection BLOCK line order (the id SET already matched; only the order differed).
        rows = sorted(rows, key=lambda r: (r["src"], r.get("sense") or "",
                                           int(r.get("since_ch", 0) or 0), int(r.get("until_ch", 0) or 0),
                                           r.get("status") or "approved", r.get("dst") or ""))
        for row in rows:
            src = row["src"]
            sense = row.get("sense", "") or ""
            since = int(row.get("since_ch", 0) or 0)
            until = int(row.get("until_ch", 0) or 0)
            status = row.get("status") or "approved"       # memseed.go:100-102
            allow_short = bool(row.get("allow_short", False))
            e = MemoryEntry(
                id=f"{src}\x1f{sense}\x1f{since}\x1f{until}",
                src=src, dst=row.get("dst", "") or "", status=status,
                sense=sense, since_ch=since, until_ch=until, allow_short=allow_short,
            )
            # surfaces = src + aliases, ONLY when dst non-empty (memory.go:202-208)
            surfaces = []
            if (e.dst or "").strip() != "":
                surfaces.append(src)
                for a in (row.get("aliases") or []):
                    al = (a.get("alias") or "").strip()
                    if al:
                        surfaces.append(al)
            for s in surfaces:
                nk = normalize_source_key(s)
                if nk == "":
                    continue
                if significant_len(nk) < min_key_len_for(nk) and not allow_short:
                    continue   # single-key ban
                e.norm_keys.append(nk)
            idx = len(self.entries)
            self.entries.append(e)
            for nk in e.norm_keys:
                self.key_owners.setdefault(nk, []).append(idx)
                if nk not in seen_key:
                    seen_key.add(nk)
                    all_keys.append(nk)
        self.keys = sorted(all_keys)                # deterministic (memory.go:241)
        self.key_index = {k: i for i, k in enumerate(self.keys)}

    # ---- matching (AC-equivalent exact multi-substring over normalized text) ----
    def _matches(self, ntext: str):
        """Return [(key_idx, start, end)] for every occurrence of every key, sorted (start,end,key_idx).
        Overlapping occurrences included (pos advances by 1), matching AC output exactly."""
        out = []
        for k in self.keys:
            ki = self.key_index[k]
            klen = len(k)
            pos = 0
            while True:
                i = ntext.find(k, pos)
                if i < 0:
                    break
                out.append((ki, i, i + klen))
                pos = i + 1
        out.sort(key=lambda m: (m[1], m[2], m[0]))
        return out

    def _spaced_phonetic_script(self, norm_key: str):
        """memory.go:781 — Latin/Cyrillic word-delimited script, or None for Han/kana/mixed."""
        script = None
        for ch in norm_key:
            if _is_cjk(ch):
                return None
            if _RE_LATIN.match(ch):
                if script == "cyr":
                    return None
                script = "lat"
            elif _RE_CYRILLIC.match(ch):
                if script == "lat":
                    return None
                script = "cyr"
        return script

    def _suppress_unbounded_phonetic(self, ms, ntext):
        kept = []
        for (ki, s, e) in ms:
            script = self._spaced_phonetic_script(self.keys[ki])
            if script is not None:
                re_s = _RE_LATIN if script == "lat" else _RE_CYRILLIC
                before_bad = s > 0 and _is_letter(ntext[s - 1]) and re_s.match(ntext[s - 1])
                after_bad = e < len(ntext) and _is_letter(ntext[e]) and re_s.match(ntext[e])
                if before_bad or after_bad:
                    continue
            kept.append((ki, s, e))
        return kept

    def _match_trust(self, m, chapter: int):
        """memory.go:656 — best disposition over spoiler-valid owners of the key; hasValid flag."""
        key = self.keys[m[0]]
        disp, has_valid = None, False
        for ei in self.key_owners[key]:
            e = self.entries[ei]
            if spoiler_blocked(e, chapter):
                continue
            d = disposition_for(e, key)
            if not has_valid or trust_rank(d) > trust_rank(disp):
                disp, has_valid = d, True
        return disp, has_valid

    def _suppress_contained(self, ms, chapter: int):
        """memory.go:684 — trust-gated longest-match. Returns (kept, trust_gate_events)."""
        kept = []
        gated = []
        for i, m in enumerate(ms):
            con_disp, con_valid = self._match_trust(m, chapter)
            contained = False
            refused = None
            for j, o in enumerate(ms):
                if i == j:
                    continue
                # o must STRICTLY contain m (memory.go:696)
                if not (o[1] <= m[1] and o[2] >= m[2] and (o[2] - o[1]) > (m[2] - m[1])):
                    continue
                sup_disp, sup_valid = self._match_trust(o, chapter)
                if not sup_valid:
                    continue
                if (not con_valid) or trust_rank(sup_disp) >= trust_rank(con_disp):
                    contained = True
                    break
                if refused is None:
                    refused = {
                        "suppressor": self.keys[o[0]], "suppressor_disp": sup_disp,
                        "protected": self.keys[m[0]], "protected_disp": con_disp,
                    }
            if contained:
                continue
            kept.append(m)
            if refused is not None:
                gated.append(refused)
        return kept, _dedupe_trust_gate(gated)

    def select(self, chunk: str, chapter: int, sticky_prev: dict, budget_tokens: int) -> Selection:
        ntext = normalize_source_key(chunk)
        occ = self._matches(ntext)
        occ = self._suppress_unbounded_phonetic(occ, ntext)
        occ, trust_gated = self._suppress_contained(occ, chapter)

        # surviving key occurrences -> entries, longest firing key per entry (memory.go:340-348)
        matched_via = {}
        for m in occ:
            k = self.keys[m[0]]
            for ei in self.key_owners[k]:
                cur = matched_via.get(ei)
                if cur is None or len(k) > len(cur):
                    matched_via[ei] = k

        sel = Selection(trust_gated=trust_gated)
        hits = {}   # id -> Picked (exact)
        for ei, e in enumerate(self.entries):
            via = matched_via.get(ei)
            if via is None:
                continue
            if spoiler_blocked(e, chapter):
                sel.rejected.append(Picked(e, via, REJECT))
                continue
            disp = disposition_for(e, via)
            hits[e.id] = Picked(e, via, disp)
            sel.active_ids[e.id] = disp

        # sticky scene-inertia (memory.go:375-388) — inherit prior disposition
        for e in self.entries:
            if e.id not in sticky_prev:
                continue
            if e.id in hits:
                continue
            if spoiler_blocked(e, chapter):
                continue
            hits[e.id] = Picked(e, "sticky", sticky_prev[e.id])

        # priority order (frozen-entry order, then stable-sort by priorityRank) (memory.go:393-401)
        order = [hits[e.id] for e in self.entries if e.id in hits]
        order.sort(key=priority_rank)  # Python sort is stable

        if budget_tokens > 0:
            used, cut = 0, len(order)
            for i, p in enumerate(order):
                used += glossary_line_tokens(p.entry.src, p.entry.dst)
                if used > budget_tokens:
                    cut = i
                    break
            sel.injected = order[:cut]
            sel.evicted = order[cut:]
        else:
            sel.injected = order
        return sel


def spoiler_blocked(e: MemoryEntry, chapter: int) -> bool:
    if e.since_ch > 0 and chapter < e.since_ch:
        return True
    if e.until_ch > 0 and chapter > e.until_ch:
        return True
    return False


def disposition_for(e: MemoryEntry, via: str) -> str:
    if via != "sticky" and not e.allow_short and collision_prone_key(via):
        return AMBIGUOUS
    if e.status == "approved":
        return CONFIRMED
    return AMBIGUOUS


def priority_rank(p: Picked) -> int:
    rank = 0
    if p.disp != CONFIRMED:
        rank |= 1 << 2
    if p.via == "sticky":
        rank |= 1 << 1
    if p.entry.status != "approved":
        rank |= 1 << 0
    return rank


def _dedupe_trust_gate(evs):
    if len(evs) < 2:
        return evs
    seen = set()
    out = []
    for e in evs:
        k = (e["suppressor"], e["protected"])
        if k in seen:
            continue
        seen.add(k)
        out.append(e)
    return out


def render_glossary_block(injected) -> str:
    """memory.go:477 — translator glossary block."""
    lines = []
    for p in injected:
        if (p.entry.dst or "").strip() == "":
            continue
        line = f"{p.entry.src} → {p.entry.dst}"
        if p.disp != CONFIRMED:
            line += "  ⟨проверить⟩"
        lines.append(line)
    if not lines:
        return ""
    return GLOSSARY_HEADER + "\n" + "\n".join(lines)


def render_editor_constraint_block(injected) -> str:
    """memory.go:515 — editor constraint block (CONFIRMED-only, dedup by dst)."""
    lines = []
    seen = set()
    for p in injected:
        if p.disp != CONFIRMED:
            continue
        dst = (p.entry.dst or "").strip()
        if dst == "" or dst in seen:
            continue
        seen.add(dst)
        lines.append(f"- «{dst}»")
    if not lines:
        return ""
    return EDITOR_HEADER + "\n" + "\n".join(lines)


# --- seed loading (mirror loadGlossarySeed enough for the hot path) ---------------------------------
def load_seed_rows(seed_path: str) -> list:
    d = yaml.safe_load(Path(seed_path).read_text(encoding="utf-8"))
    return d["terms"]


def union_sticky(active_ids_per_chunk: list, k: int, depth: int = STICKY_DEPTH) -> dict:
    """sticky_prev for chunk k = union of the previous `depth` chunks' active_ids (same chapter).
    active_ids_per_chunk[i] is the dict from chunk i's Selection.active_ids. (bookrun.go unionSticky)."""
    merged = {}
    for i in range(max(0, k - depth), k):
        for cid, disp in active_ids_per_chunk[i].items():
            merged[cid] = disp    # later chunk wins on the disposition it carried
    return merged


def build_bank(seed_path: str) -> MemoryBank:
    return MemoryBank(load_seed_rows(seed_path))


if __name__ == "__main__":
    if "--selftest" in sys.argv:
        seed = "/home/ubuntu/books/gu-zhenren/guzhenren-seed-v2.yaml"
        bank = build_bank(seed)
        print(f"entries={len(bank.entries)} unique_keys={len(bank.keys)}")
        # normalize sanity: NFKC folds fullwidth space U+3000 -> U+0020 (kept; Go does NOT trim)
        assert normalize_source_key("　方源　") == " 方源 ", repr(normalize_source_key("　方源　"))
        assert normalize_source_key("Ｑ版") == "q版", repr(normalize_source_key("Ｑ版"))
        # key "方源" still substring-matches inside the space-padded normalized text
        assert "方源" in normalize_source_key("　方源　")
        # a chunk from S2' head
        chunk = ("方源笑了，握紧拳头。古月方源！蛊师有九转。四代族长古月博。"
                 "沈翠端着酒菜。方正是甲等资质，方源只是丙等。")
        sel = bank.select(chunk, chapter=1, sticky_prev={}, budget_tokens=800)
        print("\n-- exact active_ids (no sticky) --")
        for cid, disp in sel.active_ids.items():
            print(f"  {disp:9} {cid.split(chr(31))[0]}")
        print("\n-- trust_gated events --")
        for g in sel.trust_gated:
            print(f"  {g['suppressor']}({g['suppressor_disp']}) X-> {g['protected']}({g['protected_disp']})")
        print("\n-- TRANSLATOR block --")
        print(render_glossary_block(sel.injected))
        print("\n-- EDITOR block --")
        print(render_editor_constraint_block(sel.injected))
        # longest-match check: 古月方源 should suppress nested 方源 and 古月 as firing keys,
        # but the ENTITY 方源 still fires via its own occurrences elsewhere in the chunk.
        print("\nselftest OK")
