"""rsi/selfprove.py — self-proof + repair for generated output (session v8 proofer).

Drives the denoising expert (trained on <<<NOISY ...>>> / <<<CLEAN>>> records from
build_proof_u8.py) with an iterative masked-span refinement loop — the discrete
diffusion inference half: per round, the worst-CE spans are masked out and
re-sampled from the model's reconstruction; repeated over a decaying top-k schedule.

Causal ctx=256 trunk ⇒ span-local edits with rolling context (chunking mirrors the
bundle brain's continuation_logprob): the prompt is always truncated to the last
cfg.ctx bytes, so the NOISY record is the local window being repaired.

CLI:
  python3 rsi/selfprove.py --model experts/cand_v8/proof.pt --in out.html \
      --report report.json --repaired fixed.html [--rounds 3]

Exit 0 always — this is a measurement tool, not a gate.
"""
import argparse
import json
import os
import sys
from html.parser import HTMLParser

import numpy as np
import torch

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

from gwlm_v6 import GoldWormV6, GoldWormV6Config              # noqa: E402
from gwlm_v7 import GoldWormV7, GoldWormV7Config              # noqa: E402

LN2 = float(np.log(2.0))
NOISY_H = lambda n: b"<<<NOISY " + str(n).encode() + b">>>\n"  # noqa: E731
CLEAN_H = b"\n<<<CLEAN>>>\n"
END_MARK = b"<<<END>>>"

VOID = {"AREA", "BASE", "BR", "COL", "EMBED", "HR", "IMG", "INPUT", "LINK",
        "META", "PARAM", "SOURCE", "TRACK", "WBR"}


def load(path, dev="cuda"):
    """Family dispatch mirrors moe_gate.upcycle_experts (uor_tau -> v7 else v6)."""
    ck = torch.load(path, map_location=dev, weights_only=False)
    cfg = ck["cfg"]
    if "uor_tau" in cfg:
        m = GoldWormV7(GoldWormV7Config(**cfg)).to(dev)
    else:
        m = GoldWormV6(GoldWormV6Config(**cfg)).to(dev)
    m.load_state_dict(ck["model_state_dict"])
    m.eval()
    return m


@torch.no_grad()
def score(model, text, dev="cuda", W=256, stride=128):
    """Rolling-window mean next-byte CE (bits). Returns (bpc_mean, spans) with
    spans = [(offset, bpc)] sorted worst-first. Windows shorter than 8 B skipped."""
    if len(text) < 16:
        return 0.0, []
    offs = list(range(0, max(1, len(text) - W), stride))
    if not offs or offs[-1] + W < len(text):
        offs.append(max(0, len(text) - W))
    spans = []
    tot = n = 0.0
    for off in offs:
        w = text[off:off + W]
        if len(w) < 8:
            continue
        b = np.frombuffer(w, dtype=np.uint8).astype(np.int64)
        x = torch.from_numpy(b[:-1]).long().unsqueeze(0).to(dev)
        y = torch.from_numpy(b[1:]).long().unsqueeze(0).to(dev)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(dev == "cuda")):
            _, loss = model(x, targets=y)
        bpc = float(loss.item()) / LN2
        spans.append((off, bpc))
        tot += float(loss.item()) * (len(w) - 1)
        n += len(w) - 1
    spans.sort(key=lambda t: -t[1])
    return (tot / n / LN2 if n else 0.0), spans


@torch.no_grad()
def repair_span(model, ctx_prefix, span, dev="cuda", ctx=256):
    """One denoising shot: NOISY record prompt -> greedy reconstruction."""
    prompt = ctx_prefix[-128:] + NOISY_H(len(span)) + span + CLEAN_H
    toks = list(prompt)
    gen = bytearray()
    budget = int(len(span) * 1.2) + 16
    for _ in range(budget):
        x = torch.tensor([toks[-ctx:]], dtype=torch.long, device=dev)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(dev == "cuda")):
            logits = model(x)
        nxt = int(logits[:, -1].float().argmax(-1).item())
        gen.append(nxt)
        toks.append(nxt)
        if gen.endswith(END_MARK):
            gen = gen[:-len(END_MARK)]
            break
    return bytes(gen)


class TagBalance(HTMLParser):
    """Minimal strict-ish checker: every non-void open must close in LIFO order."""

    def __init__(self):
        super().__init__()
        self.stack = []
        self.mismatch_off = None

    def handle_starttag(self, tag, attrs):
        if tag.upper() not in VOID:
            self.stack.append(tag)

    def handle_startendtag(self, tag, attrs):
        pass

    def handle_endtag(self, tag):
        if tag.upper() in VOID:
            return
        line, col = self.getpos()
        if self.stack and self.stack[-1] == tag:
            self.stack.pop()
        elif tag in self.stack:
            if self.mismatch_off is None:
                self.mismatch_off = (line, col)
            while self.stack and self.stack[-1] != tag:
                self.stack.pop()
            if self.stack:
                self.stack.pop()
        else:
            if self.mismatch_off is None:
                self.mismatch_off = (line, col)


def html_validate(text):
    """(ok, mismatch_offset) — ok = balanced stack and no mismatch seen."""
    p = TagBalance()
    try:
        p.feed(text.decode("utf-8", "replace"))
        p.close()
    except Exception as e:  # HTMLParser raises on some malformed input
        return False, str(e)[:80]
    ok = (p.mismatch_off is None) and not p.stack
    return ok, (None if p.mismatch_off is None else str(p.mismatch_off))


def looks_html(text):
    head = text.lstrip()[:64].lower()
    return head.startswith(b"<!doctype") or head.startswith(b"<html")


def prove(model, text, dev="cuda", rounds=3, topk=(4, 2, 1), W=256):
    """Iterative masked-span refinement. Returns (repaired_bytes, report_dict)."""
    html_mode = looks_html(text)
    bpc0, _ = score(model, text, dev, W=W)
    html_ok0, _ = html_validate(text) if html_mode else (None, None)
    cur = text
    edits = 0
    worst = []
    rounds_applied = 0
    for r in range(rounds):
        k = topk[min(r, len(topk) - 1)]
        bpc_mean, spans = score(model, cur, dev, W=W)
        if not spans:
            break
        worst = spans[:1]
        floor = bpc_mean * 2.0      # healthy tissue is not repaired: a span must be
        chosen = []                 # clearly pathological (corruption spikes >>2x)
        for off, bpc in spans:
            if bpc <= floor:
                continue
            if len(chosen) >= k:
                break
            if all(abs(off - o) >= W for o, _ in chosen):
                chosen.append((off, bpc))
        if not chosen:
            break
        rounds_applied += 1
        for off, _bpc in sorted(chosen, key=lambda t: -t[0]):  # right-to-left
            span = cur[off:off + W]
            prefix = cur[max(0, off - 128):off]
            new = repair_span(model, prefix, span, dev, ctx=model.cfg.ctx)
            if new and len(new) <= int(len(span) * 1.2) + 16:
                cur = cur[:off] + new + cur[off + len(span):]
                edits += 1
    bpc1, _ = score(model, cur, dev, W=W)
    html_ok1, _ = html_validate(cur) if html_mode else (None, None)
    report = {"bpc_before": round(bpc0, 4), "bpc_after": round(bpc1, 4),
              "worst_spans": [{"off": o, "bpc": round(b, 4)} for o, b in worst],
              "rounds_applied": rounds_applied, "edits": edits,
              "html_ok_before": html_ok0, "html_ok_after": html_ok1}
    return cur, report


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--in", dest="inp", required=True)
    ap.add_argument("--report", default=None)
    ap.add_argument("--repaired", default=None)
    ap.add_argument("--rounds", type=int, default=3)
    a = ap.parse_args()
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    model = load(a.model, dev)
    with open(a.inp, "rb") as f:
        text = f.read()
    repaired, report = prove(model, text, dev, rounds=a.rounds)
    report = {"file": a.inp, **report}
    if a.repaired:
        with open(a.repaired, "wb") as f:
            f.write(repaired)
    if a.report:
        with open(a.report, "w") as f:
            json.dump(report, f, indent=1)
    print(json.dumps(report, indent=1))


if __name__ == "__main__":
    main()