#!/usr/bin/env python3
"""
Adjudica los candidatos dudosos con Claude (Anthropic API).

Para cada candidato (REVISAR / REPO_TYPO / AMBIGUO):
  - se le da a Claude el CONTEXTO de la repo con el tramo disputado marcado [[...]]
  - y el texto LIMPIO de la pagina del escaneo (Surya si existe; si no tesseract psm6)
  - Claude decide si la repo se desvia del original (escaneo = verdad) y, si difiere,
    da la lectura exacta del escaneo.

Salida: out/<book>.adjudicated.jsonl  (candidato + veredicto del LLM)
        out/review_entries_claude.jsonl  (cuando se usa --entries)

Uso:
  python3 adjudicate_llama.py --book isopanisad
  python3 adjudicate_llama.py --all [--only REVISAR,REPO_TYPO,AMBIGUO]
  python3 adjudicate_llama.py --entries /path/to/review_entries.json
"""
import json, re, sys, argparse
from pathlib import Path
import anthropic

HERE = Path(__file__).resolve().parent
CFG = json.loads((HERE / "books.json").read_text())
OCR_DIR = HERE / "ocr"
SURYA_DIR = HERE / "surya"
OUT = HERE / "out"
MODEL = "claude-sonnet-4-6"

_client = None
def get_client():
    global _client
    if _client is None:
        _client = anthropic.Anthropic()
    return _client

def page_text(book, pageref):
    """pageref = '<stem>#pNNNN'. Devuelve texto limpio: surya si hay, si no psm6."""
    if not pageref or "#" not in pageref:
        return ""
    stem, pg = pageref.split("#", 1)
    sp = SURYA_DIR / book / stem / f"{pg}.surya.txt"
    if sp.exists():
        return sp.read_text(errors="ignore")
    tp = OCR_DIR / book / stem / f"{pg}.psm6.txt"
    return tp.read_text(errors="ignore") if tp.exists() else ""

PROMPT = """You are a textual scholar auditing a digital edition (REPO) of an A.C. Bhaktivedanta Swami Prabhupada book. The authority is the ORIGINAL first-edition scan text below.

REPO context — the disputed span is in [[double brackets]]:
{context}

ORIGINAL scan page OCR (may have minor OCR noise like 'l' for 'i'; ignore page numbers/headers):
\"\"\"
{page}
\"\"\"

TASK: Locate the passage in the ORIGINAL that corresponds to the [[bracketed]] span in the REPO.
1. Quote the EXACT words the ORIGINAL uses for that span.
2. Decide: does the ORIGINAL use DIFFERENT WORDS from the REPO (a genuine editorial difference, not a single-character OCR noise substitution like 'l'→'i' or 'e'→'c')?

Answer STRICT JSON only — no prose outside the JSON:
{{"found": true|false, "original_reading": "<exact words from the ORIGINAL for that span>", "is_real_diff": true|false, "confidence": "high|medium|low", "note": "<one sentence: what the original says and why it differs or not>"}}

Rules:
- is_real_diff=true when the ORIGINAL uses a genuinely different word or phrase (e.g. "each other" vs "one another", "instruction" vs "instructions").
- is_real_diff=false only when the ORIGINAL matches the REPO or the only difference is obvious single-character OCR noise (e.g. 'il' for 'il', 'pard' for 'para').
- If you cannot locate the passage, set found=false and is_real_diff=false.
JSON:"""

def call_claude(prompt):
    msg = get_client().messages.create(
        model=MODEL,
        max_tokens=512,
        temperature=0.0,
        messages=[{"role": "user", "content": prompt}],
    )
    return msg.content[0].text

def parse_json(s):
    m = re.search(r'\{.*\}', s, re.DOTALL)
    if not m:
        return None
    try:
        return json.loads(m.group(0))
    except Exception:
        return None

def apply_consistency_rule(c):
    """Si Claude encontró una lectura CORTA diferente al repo, forzar real."""
    m = c.get("llm")
    if isinstance(m, dict) and m.get("found") and m.get("original_reading"):
        orig_clean = re.sub(r'[^a-z0-9 ]', '', m["original_reading"].lower()).strip()
        repo_clean = re.sub(r'[^a-z0-9 ]', '', c["repo"].lower()).strip()
        orig_words = set(orig_clean.split())
        repo_words = set(repo_clean.split())
        short = len(orig_words) <= 6
        repo_in_orig = repo_words.issubset(orig_words)
        if short and not repo_in_orig and orig_words != repo_words:
            m["is_real_diff"] = True

def adjudicate_one(c, book):
    page = page_text(book, c.get("page", ""))
    if not page.strip():
        c["llm"] = {"error": "no_page_text"}
    else:
        prompt = PROMPT.format(context=c["context"][:1200], repo=c["repo"],
                               scan=c.get("scan", "")[:200], page=page[:6000])
        try:
            resp = call_claude(prompt)
            c["llm"] = parse_json(resp) or {"raw": resp[:300]}
        except Exception as e:
            c["llm"] = {"error": str(e)[:200]}
    apply_consistency_rule(c)
    return c

def run_entries(entries_path):
    """Re-adjudica las entradas de review_entries.json buscando su page en candidates.jsonl."""
    entries = json.loads(Path(entries_path).read_text())

    # Build lookup: (book, file, repo, scan) -> page
    page_lookup = {}
    books_needed = {e["book"] for e in entries}
    for book in books_needed:
        cf = OUT / f"{book}.candidates.jsonl"
        if not cf.exists():
            continue
        for line in cf.open():
            c = json.loads(line)
            key = (book, c.get("file", ""), c.get("repo", ""), c.get("scan", ""))
            page_lookup[key] = c.get("page", "")

    outf = OUT / "review_entries_claude.jsonl"
    n_real = 0
    with outf.open("w") as fh:
        for i, e in enumerate(entries, 1):
            book = e["book"]
            key = (book, e.get("file", ""), e.get("repo", ""), e.get("scan", ""))
            e["page"] = page_lookup.get(key, "")
            adjudicate_one(e, book)
            if isinstance(e.get("llm"), dict) and e["llm"].get("is_real_diff"):
                n_real += 1
            fh.write(json.dumps(e, ensure_ascii=False) + "\n")
            if i % 10 == 0:
                print(f"  {i}/{len(entries)} (reales={n_real})", file=sys.stderr, flush=True)
    print(f"entries: adjudicados {len(entries)}, diffs reales={n_real} -> {outf}", file=sys.stderr, flush=True)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--book")
    ap.add_argument("--all", action="store_true")
    ap.add_argument("--only", default="REVISAR,REPO_TYPO,AMBIGUO")
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--entries", help="Path to review_entries.json for targeted re-adjudication")
    a = ap.parse_args()

    if a.entries:
        run_entries(a.entries)
        return

    classes = set(a.only.split(","))
    books = [a.book] if a.book else list(CFG["books"].keys())

    for book in books:
        cf = OUT / f"{book}.candidates.jsonl"
        if not cf.exists():
            continue
        cands = [json.loads(l) for l in cf.open()]
        todo = [c for c in cands if c["verdict"] in classes]
        if a.limit:
            todo = todo[:a.limit]
        outf = OUT / f"{book}.adjudicated.jsonl"
        n_real = 0
        with outf.open("w") as fh:
            for i, c in enumerate(todo, 1):
                adjudicate_one(c, book)
                if isinstance(c.get("llm"), dict) and c["llm"].get("is_real_diff"):
                    n_real += 1
                fh.write(json.dumps(c, ensure_ascii=False) + "\n")
                if i % 10 == 0:
                    print(f"  {book}: {i}/{len(todo)} (reales={n_real})",
                          file=sys.stderr, flush=True)
        print(f"{book}: adjudicados {len(todo)}, diffs reales={n_real} -> {outf.name}",
              file=sys.stderr, flush=True)

if __name__ == "__main__":
    main()