#!/usr/bin/env python3
"""exp15 — kimi enrichment of T-ana/T-ell straddles AT flat-A0-30 boundaries (guard 1: kimi ONLY as
targeted enrichment on the boundaries, deterministic-first). For each greedy boundary we give kimi the
TAIL of chunk_i and the HEAD of chunk_{i+1} and ask ONLY: is there an anaphora (esp. zh zero-subject) or
ellipsis in the HEAD whose antecedent/elided element is in the TAIL (i.e. it straddles THIS cut)? Batched
(several boundaries per call) to bound kimi cost. Family-isolated miner (Moonshot), temp 1, 24k + retry.
Run: python enrich_ana.py --run
"""
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))
import exp15_llm as L
from chunker import greedy_chunks, split_source_sentences

SD = Path("/home/ubuntu/books/gu-zhenren/exp15")
FLAT = SD / "S2prime_big_flat.txt"
OUT = SD / "ana_enrichment.json"
LEDGER = SD / "enrich_costs.jsonl"
MINER = "kimi-k2.6"
BATCH = 6          # boundaries per kimi call
TAIL_N, HEAD_N = 3, 3

SYS = (
    "Ты — эксперт по китайско-русскому переводу и связности. Тебе дают несколько ГРАНИЦ нарезки текста "
    "(ранобэ 蛊真人): для каждой — ХВОСТ предыдущего чанка и ГОЛОВА следующего. Для КАЖДОЙ границы найди, "
    "есть ли в ГОЛОВЕ анафора (особенно китайское НУЛЕВОЕ подлежащее — предложение без явного субъекта, "
    "субъект тянется из ХВОСТА) ИЛИ эллипсис, чей антецедент/опущенный элемент — в ХВОСТЕ (т.е. зависимость "
    "ПЕРЕСЕКАЕТ именно эту границу). Если да — верни объект; если нет — пропусти границу.\n"
    'Формат (JSON-массив): [{"boundary_idx":ЦЕЛОЕ,"type":"T-ana|T-ell",'
    '"antecedent_span":"цитата из ХВОСТА","dependent_span":"цитата из ГОЛОВЫ",'
    '"phenomenon":"какая зависимость и почему разрез её ломает",'
    '"expected_ru":"что должен сделать верный перевод (референт/род/восстановленное сказуемое)"}]\n'
    "Только реальные пересекающие границу случаи; цитаты дословные. Без текста вне JSON."
)


def boundaries():
    text = FLAT.read_text(encoding="utf-8")
    ch = greedy_chunks(text)
    out = []
    for i in range(len(ch) - 1):
        tail = "".join(split_source_sentences(ch[i])[-TAIL_N:]).strip()
        head = "".join(split_source_sentences(ch[i + 1])[:HEAD_N]).strip()
        out.append({"boundary_idx": i, "tail": tail, "head": head})
    return out


def parse_array(t):
    m = re.search(r"\[.*\]", t or "", re.S)
    if not m:
        return None
    try:
        return json.loads(m.group(0))
    except Exception:
        return None


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--run", action="store_true"); a = ap.parse_args()
    b = boundaries()
    print(f"{len(b)} boundaries, batch {BATCH} -> {(len(b)+BATCH-1)//BATCH} kimi calls")
    if not a.run:
        return
    sp = L.Spender(str(LEDGER), per_call_cap=0.28, hard_cap=2.5)
    found = []
    NUDGE = {"role": "user", "content": "Не рассуждай долго. Сразу выведи ТОЛЬКО компактный валидный JSON-массив, без markdown-ограды."}
    for k in range(0, len(b), BATCH):
        batch = b[k:k + BATCH]
        user = "\n\n".join(f"ГРАНИЦА {x['boundary_idx']}:\n[ХВОСТ]: {x['tail']}\n[ГОЛОВА]: {x['head']}" for x in batch)
        msgs = [{"role": "system", "content": SYS}, {"role": "user", "content": user}]
        in_est = (len(SYS) + len(user)) // 2
        ok, pred, why = sp.gate(MINER, in_est, L.MODELS[MINER]["max_tokens"])
        print(f"[gate] batch {k//BATCH} (bnd {batch[0]['boundary_idx']}-{batch[-1]['boundary_idx']}) pred=${pred:.4f} {'OK' if ok else why}")
        if not ok:
            continue
        r = L.call(MINER, msgs, temp=1.0, max_tokens=24000, timeout=600)
        sp.record({"batch": k // BATCH, "usage": r.get("usage"), "cost": r.get("cost", 0.0), "err": r.get("err"), "latency_s": r.get("latency_s")})
        cands = parse_array(r.get("text", ""))
        if cands is None:
            r2 = L.call(MINER, msgs + [NUDGE], temp=1.0, max_tokens=24000, timeout=600)
            sp.record({"batch": k // BATCH, "retry": True, "usage": r2.get("usage"), "cost": r2.get("cost", 0.0), "err": r2.get("err")})
            cands = parse_array(r2.get("text", "")) if not r2.get("err") else None
        cands = cands or []
        found.extend(cands)
        with open(str(OUT) + ".partial.jsonl", "a", encoding="utf-8") as pf:
            pf.write(json.dumps({"batch": k // BATCH, "n": len(cands), "cands": cands}, ensure_ascii=False) + "\n")
        print(f"  batch {k//BATCH}: {len(cands)} straddles  cost=${r.get('cost',0):.4f}")
    from collections import Counter
    by = Counter(c.get("type") for c in found)
    OUT.write_text(json.dumps({"n": len(found), "by_type": dict(by), "straddles": found}, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\nenriched straddles: {dict(by)} (total {len(found)})  spend ${sp.total:.4f} -> {OUT}")


if __name__ == "__main__":
    main()