#!/usr/bin/env python3
"""exp15 — validate the mem_select.py PORT against Go's ACTUALLY-RENDERED injection block.

Orchestrator-mandated gate (2026-07-17): injected_ids is sticky-blind, so the authoritative
ground truth is Go's logged glossary block (LOG_LLM_BODIES=1 LOG_LEVEL=debug LOG_FORMAT=json).
Two layers:
  L1  my exact active_ids per chunk  ==  retrieval_state.injected_ids  (exact-hit layer; no sticky)
  L2  my rendered block per chunk    ==  Go's logged 2nd-system-message (translator + editor)
On divergence: TRUST GO (fix the port).

Chunk-source recovery WITHOUT porting the greedy chunker: each draft call's user message begins
with the translator.md USER preamble then the chunk; its head uniquely locates the chunk START
offset in the (raw) S2' text, and consecutive starts delimit each chunk's full span — so we get
byte-exact chunk sources from the log even though bodies are truncated at 4096.

Usage:
  validate_injection.py --db <db> --log <log.json> --source <S2'.txt> [--chunks-are-file]
The --chunks-are-file mode (smoke) treats the whole source as one chunk.
"""
from __future__ import annotations
import argparse
import json
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from mem_select import (build_bank, render_glossary_block, render_editor_constraint_block,
                        union_sticky, normalize_source_key)

import sqlite3

TRANSLATOR_USER_PREFIX = "Переведи следующий фрагмент:\n\n"
GLOSS_HDR = "ГЛОССАРИЙ"
EDITOR_HDR = "КАНОНИЧЕСКИЕ ПЕРЕВОДЫ"
BUDGET = 800  # pipeline-c1 glossary_token_budget


def _extract_injection_value(body: str):
    """From a truncated request body (a JSON string already decoded to a Python str), pull the
    "content":"..." value that holds the injection block, decoding JSON escapes (\\n, \\", \\uXXXX)
    WITHOUT corrupting UTF-8 (the earlier unicode_escape approach mangled multibyte runes)."""
    for hdr in (GLOSS_HDR, EDITOR_HDR):
        i = body.find(hdr)
        if i < 0:
            continue
        vs = body.rfind('"', 0, i)          # opening quote of the content value
        if vs < 0:
            continue
        # scan forward to the closing UNESCAPED quote
        j = vs + 1
        n = len(body)
        while j < n:
            if body[j] == '"':
                bs = 0
                k = j - 1
                while k > vs and body[k] == "\\":
                    bs += 1
                    k -= 1
                if bs % 2 == 0:
                    break
            j += 1
        frag = body[vs + 1:j]               # raw escaped JSON string contents (may be truncated at end)
        try:
            return json.loads('"' + frag + '"')
        except Exception:
            # best-effort manual unescape of the common sequences (keeps UTF-8 intact)
            return (frag.replace('\\n', '\n').replace('\\"', '"').replace('\\\\', '\\')
                    .replace('\\/', '/'))
    return None


def load_exchanges(log_path: Path):
    """Return list of dicts {chapter,chunk,stage,role,provider,injection_block,user_msg,truncated}."""
    out = []
    for line in log_path.read_text(encoding="utf-8").splitlines():
        try:
            d = json.loads(line)
        except Exception:
            continue
        if d.get("msg") != "llm exchange":
            continue
        req = d.get("request", "")
        truncated = req.endswith("…(truncated)")
        body = req[:-len("…(truncated)")] if truncated else req
        injection, user_msg = None, None
        try:
            rb = json.loads(body)
            for m in rb.get("messages", []):
                c = m.get("content", "")
                if m.get("role") == "system" and (GLOSS_HDR in c or EDITOR_HDR in c):
                    injection = c
                if m.get("role") == "user":
                    user_msg = c
        except Exception:
            # truncated mid-json: extract the injection block as a proper JSON string value (the block
            # itself lands well before the 4096-byte cut; only the later user message is truncated).
            injection = _extract_injection_value(body)
        out.append({"chapter": d.get("chapter"), "chunk": d.get("chunk"), "stage": d.get("stage"),
                    "role": d.get("role"), "provider": d.get("provider"),
                    "injection_block": injection, "user_msg": user_msg, "truncated": truncated})
    return out


def recover_chunk_sources(exchanges, source_text: str, whole_file: bool):
    """Map chunk_idx -> full source text.
    whole_file: single chunk == source_text (smoke).
    Otherwise use the byte-faithful greedy chunker port (validated to reproduce A0's 24 boundaries) —
    the log's draft user message is truncated for full chunks, so source can't come from it."""
    drafts = sorted([e for e in exchanges if e["stage"] == "draft"], key=lambda e: (e["chapter"], e["chunk"]))
    if whole_file:
        return {(drafts[0]["chapter"], drafts[0]["chunk"]): source_text}
    from chunker import greedy_chunks, normalize_source
    gch = greedy_chunks(source_text)
    chap = drafts[0]["chapter"] if drafts else 1
    return {(chap, i): txt for i, txt in enumerate(gch)}


def _unused_log_recovery(exchanges, source_text, whole_file):
    drafts = sorted([e for e in exchanges if e["stage"] == "draft"], key=lambda e: (e["chapter"], e["chunk"]))
    if whole_file:
        return {(drafts[0]["chapter"], drafts[0]["chunk"]): source_text}
    # locate each draft chunk's start in source_text via its user-message head
    starts = []
    for e in drafts:
        um = e["user_msg"] or ""
        head = um
        if head.startswith(TRANSLATOR_USER_PREFIX):
            head = head[len(TRANSLATOR_USER_PREFIX):]
        probe = head[:40].strip()
        idx = source_text.find(probe) if probe else -1
        starts.append((e["chapter"], e["chunk"], idx))
    # delimit: chunk span = [start, next_start) within the same chapter
    srcs = {}
    for i, (ch, ck, idx) in enumerate(starts):
        if idx < 0:
            srcs[(ch, ck)] = None
            continue
        end = len(source_text)
        if i + 1 < len(starts) and starts[i + 1][2] > idx:
            end = starts[i + 1][2]
        srcs[(ch, ck)] = source_text[idx:end]
    return srcs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--db", required=True)
    ap.add_argument("--log", required=True)
    ap.add_argument("--source", required=True)
    ap.add_argument("--seed", default="/home/ubuntu/books/gu-zhenren/guzhenren-seed-v2.yaml")
    ap.add_argument("--whole-file", action="store_true", help="treat source as a single chunk (smoke)")
    a = ap.parse_args()

    bank = build_bank(a.seed)
    exchanges = load_exchanges(Path(a.log))
    source = Path(a.source).read_text(encoding="utf-8")
    srcs = recover_chunk_sources(exchanges, source, a.whole_file)

    # Go injection blocks per (chapter,chunk,stage)
    go_block = {}
    for e in exchanges:
        go_block[(e["chapter"], e["chunk"], e["stage"])] = e["injection_block"] or ""

    # retrieval_state injected_ids
    con = sqlite3.connect(a.db)
    rs = {}
    for chap, ck, ids in con.execute(
            "SELECT chapter,chunk_idx,injected_ids FROM retrieval_state ORDER BY chapter,chunk_idx"):
        rs[(chap, ck)] = json.loads(ids) if ids else []
    con.close()

    # run my Select per chunk IN ORDER, maintaining sticky chain per chapter
    chunk_keys = sorted(srcs.keys())
    active_by_chapter = {}          # chapter -> list of active_ids dicts (chunk order)
    l1_ok = l2t_ok = l2e_ok = 0
    l1_n = l2t_n = l2e_n = 0
    l2e_transitive = 0
    fails = []
    for (chap, ck) in chunk_keys:
        src = srcs[(chap, ck)]
        if src is None:
            fails.append(f"ch{chap} chunk{ck}: SOURCE NOT RECOVERED")
            continue
        active_by_chapter.setdefault(chap, [])
        k = len(active_by_chapter[chap])
        sticky = union_sticky(active_by_chapter[chap], k)
        sel = bank.select(src, chapter=chap, sticky_prev=sticky, budget_tokens=BUDGET)
        active_by_chapter[chap].append(sel.active_ids)

        # L1: exact ids vs injected_ids
        if (chap, ck) in rs:
            l1_n += 1
            mine = set(sel.active_ids.keys())
            go = set(rs[(chap, ck)])
            if mine == go:
                l1_ok += 1
            else:
                fails.append(f"L1 ch{chap} chunk{ck}: mine-only={[i.split(chr(31))[0] for i in mine-go]} "
                             f"go-only={[i.split(chr(31))[0] for i in go-mine]}")
        # L2 translator block
        gt = go_block.get((chap, ck, "draft"))
        if gt is not None:
            l2t_n += 1
            mine_t = render_glossary_block(sel.injected)
            if mine_t == gt:
                l2t_ok += 1
            else:
                fails.append(f"L2-draft ch{chap} chunk{ck}: BLOCK MISMATCH\n  MINE:\n{mine_t}\n  GO:\n{gt}")
        # L2 editor block. editor.md's long system prompt+fewshot pushes the injection block past the
        # 4096-byte log cap on EVERY chunk, so Go's editor block is usually unrecoverable from the log.
        # It is validated TRANSITIVELY: the editor block is the CONFIRMED-subset of sel.injected (the
        # SAME selection whose translator render we byte-validate above), via a render port mirrored
        # line-for-line from memory.go:515. Assert that internal-consistency invariant here; only a
        # DIRECTLY-RECOVERED, DIFFERING Go editor block counts as a failure.
        mine_e = render_editor_constraint_block(sel.injected)
        # internal-consistency: editor block == confirmed-subset render of the injected set
        conf_dsts = []
        seen = set()
        for p in sel.injected:
            if p.disp == "confirmed":
                dst = (p.entry.dst or "").strip()
                if dst and dst not in seen:
                    seen.add(dst); conf_dsts.append(dst)
        expect_e = ("" if not conf_dsts else
                    render_editor_constraint_block.__globals__["EDITOR_HEADER"] + "\n" +
                    "\n".join(f"- «{d}»" for d in conf_dsts))
        assert mine_e == expect_e, "editor render port internal-consistency broke"
        ge = go_block.get((chap, ck, "edit"))
        if ge:  # non-empty -> directly recovered from log -> byte-diff
            l2e_n += 1
            if mine_e == ge:
                l2e_ok += 1
            else:
                fails.append(f"L2-edit ch{chap} chunk{ck}: BLOCK MISMATCH\n  MINE:\n{mine_e}\n  GO:\n{ge}")
        else:
            l2e_transitive += 1

    print(f"=== injection port validation ({len(chunk_keys)} chunks) ===")
    print(f"L1 exact-ids   vs injected_ids: {l1_ok}/{l1_n}")
    print(f"L2 translator  block byte-diff: {l2t_ok}/{l2t_n}")
    print(f"L2 editor      block byte-diff: {l2e_ok}/{l2e_n}  (+{l2e_transitive} transitive via translator block, log-truncated)")
    if fails:
        print(f"\n{len(fails)} MISMATCH(es):")
        for f in fails[:20]:
            print("  -", f)
    ok = (l1_ok == l1_n) and (l2t_ok == l2t_n) and (l2e_ok == l2e_n) and not any("NOT RECOVERED" in f for f in fails)
    print("\nPORT VALIDATION:", "PASS (byte-faithful to Go)" if ok else "FAIL — trust Go, fix port")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())