#!/usr/bin/env python3
"""exp16 — A6: local 9B entity spotter over the webnovel source (research/20 §D2 A6, §B4 Z3). $0 (stand).

Re-measures the exp06 caveat (exp06 spot recall 0.95-1.0 was on PD-classic; the webnovel-slice re-measure
is the eval/README-flagged open question). Prompts a local 9B to spot named entities/terms per source chunk,
aggregates across the slice, and measures spot-recall against GT + hallucination. Health-checked per exp06
(empty/echo/unparseable -> retry with bigger budget -> HEALTH=failed, NOT counted as recall 0). Deterministic.

localhost ollama needs no-proxy transport (stand quirk: 403 on localhost under the corp proxy).
"""
from __future__ import annotations

import json
import re
import sys
import urllib.request
from pathlib import Path

import exp16_common as X

MODEL = "huihui_ai/qwen3.5-abliterated:9b"   # exp06 best local spot recall (1.0)
OLLAMA = "http://localhost:11434/api/chat"
OUT = Path("/home/ubuntu/books/gu-zhenren/exp16")

# no-proxy opener (stand quirk)
_opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))

SPOT_PROMPT = (
    "Ты — экстрактор именованных сущностей из китайского текста веб-новеллы. "
    "Найди в СЛЕДУЮЩЕМ фрагменте все имена персонажей, топонимы, титулы/ранги и уникальные термины мира "
    "книги (артефакты, техники, существа). Верни ТОЛЬКО JSON-массив строк — сами китайские подстроки как "
    "они в тексте, без перевода и пояснений. Пример: [\"方源\",\"蛊师\",\"青茅山\"].\n\nФрагмент:\n")


def call_local(prompt: str, num_predict=1024, timeout=120):
    # think:false + /no_think — spotting is extraction, not reasoning; thinking mode is slow + noisy here
    body = json.dumps({"model": MODEL, "think": False,
                       "messages": [{"role": "user", "content": prompt + "\n/no_think"}],
                       "stream": False, "options": {"temperature": 0.0, "num_predict": num_predict}}).encode()
    req = urllib.request.Request(OLLAMA, data=body, headers={"Content-Type": "application/json"})
    try:
        with _opener.open(req, timeout=timeout) as r:
            d = json.loads(r.read())
        return (d.get("message", {}) or {}).get("content", ""), None
    except Exception as e:
        return "", f"{type(e).__name__}: {str(e)[:150]}"


def parse_entities(text: str):
    """Extract the JSON array of zh strings; tolerant of <think> blocks and code fences."""
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
    m = re.search(r"\[.*\]", text, re.S)
    if not m:
        # fallback: any Han runs of length 2-6 in quotes
        return [q for q in re.findall(r'"([㐀-鿿]{1,8})"', text)]
    try:
        arr = json.loads(m.group(0))
        return [str(x).strip() for x in arr if isinstance(x, (str,)) and re.search(r"[㐀-鿿]", str(x))]
    except json.JSONDecodeError:
        return [q for q in re.findall(r'"([㐀-鿿]{1,8})"', text)]


def main():
    chunks = X.load_chunks()
    gt = X.load_gt()
    limit = int(sys.argv[1]) if len(sys.argv) > 1 else len(chunks)
    spotted = set()           # normalized spotted substrings
    spotted_raw = []
    health = {"ok": 0, "failed": 0, "retried": 0}
    per_chunk = []
    for c in chunks[:limit]:
        prompt = SPOT_PROMPT + c.source
        ents, err, npred = [], None, 1024
        for attempt in range(2):
            out, err = call_local(prompt, num_predict=npred)
            ents = parse_entities(out)
            if ents and not err:
                if attempt > 0:
                    health["retried"] += 1
                break
            npred = 2048   # bigger budget on retry
        if err or not ents:
            health["failed"] += 1
            per_chunk.append({"id": f"{c.chapter}.{c.chunk_idx}", "health": "failed", "err": err})
            print(f"  ch{c.chapter}.{c.chunk_idx}: HEALTH=failed ({err})")
            continue
        health["ok"] += 1
        for e in ents:
            en = X.norm(e)
            if en:
                spotted.add(en)
                spotted_raw.append(e)
        per_chunk.append({"id": f"{c.chapter}.{c.chunk_idx}", "health": "ok", "n": len(ents)})
        print(f"  ch{c.chapter}.{c.chunk_idx}: {len(ents)} ents")

    # recall vs GT: a GT entity is spotted if any of its normalized surfaces is a spotted substring
    # OR a spotted string contains/is contained in the surface (surface-level match, spot-lenient)
    def gt_spotted(ent):
        for ns in ent.norm_surfaces:
            if ns in spotted:
                return True
            for sp in spotted:
                if ns in sp or sp in ns:
                    return True
        return False

    hits = [e for e in gt if gt_spotted(e)]
    by_stratum = {"f>=10": [], "f3-9": [], "f<3": []}
    for e in gt:
        by_stratum[X.freq_stratum(X.gt_occurrences(e, chunks))].append(gt_spotted(e))
    recall = len(hits) / len(gt)
    # hallucination proxy: spotted normalized strings that do NOT occur in ANY source (as substring)
    allsrc = "".join(c.nsource for c in chunks)
    halluc = [s for s in spotted if s not in allsrc]

    print(f"\n=== A6 local 9B spot ({MODEL}) — {limit} chunks ===")
    print(f"health: {health}")
    print(f"spot recall vs GT: {len(hits)}/{len(gt)} = {recall:.3f}")
    for st, v in by_stratum.items():
        print(f"   {st}: {sum(v)}/{len(v)} = {sum(v)/len(v):.2f}" if v else f"   {st}: n/a")
    print(f"distinct spotted: {len(spotted)}; hallucinated (not in source): {len(halluc)} {list(halluc)[:10]}")
    miss = [e.src for e in gt if not gt_spotted(e)]
    print(f"GT misses: {miss}")
    json.dump({"model": MODEL, "limit": limit, "health": health, "recall": recall,
               "by_stratum": {k: (sum(v), len(v)) for k, v in by_stratum.items()},
               "n_spotted": len(spotted), "halluc": list(halluc), "misses": miss,
               "spotted": sorted(spotted)},
              open(OUT / "a6_local_spot.json", "w"), ensure_ascii=False, indent=1)
    print(f"[written] {OUT/'a6_local_spot.json'}")


if __name__ == "__main__":
    main()