"""CiteNet demo server: live phantom-citation checking on Core ML.

  python demo/server.py            -> http://localhost:8025

Every request runs the hand-written model through Core ML: b1 on the CPU for a
single citation, b32 / b512 on the Neural Engine for more. Timings in the
responses are measured around MLModel.predict in this process. The firehose
runs the native Swift binary bench/firehose (ANE b512, GPU b8192, or both, two
threads per unit).
"""
import json
import os
import random
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import numpy as np
import coremltools as ct

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
sys.path.insert(0, os.path.join(ROOT, "scripts"))
from citenet import spec, text            # noqa: E402
from citenet.net import knot_readout, K_REP, K_COURT, ANCHOR_W  # noqa: E402
from lepc import load                     # noqa: E402

BUILD = os.path.join(ROOT, "build")
PORT = int(os.environ.get("PORT", 8025))

META = json.load(open(os.path.join(BUILD, "meta.json")))
W = np.load(os.path.join(BUILD, "weights.npz"))
CONCEPTS = META["concepts"]
EVIDENCE = META["evidence"]
EV_INFO = {e[0]: dict(weight=e[1], cap=e[2], text=e[3]) for e in spec.EVIDENCE}

MODELS = {
    1: ("CPU", ct.models.MLModel(os.path.join(BUILD, "CiteNet_b1.mlpackage"), compute_units=ct.ComputeUnit.CPU_ONLY)),
    32: ("ANE", ct.models.MLModel(os.path.join(BUILD, "CiteNet_b32.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)),
    512: ("ANE", ct.models.MLModel(os.path.join(BUILD, "CiteNet_b512.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)),
}
LOCK = threading.Lock()
LAST_ACTIVE = [0.0]       # time of the last wall request; the keep-alive runs while it is recent


def keep_warm():
    """While a wall is open, run a tiny batch every 100 ms so the Neural Engine stays awake.
    Idle for a few seconds, the ANE powers down and the next call pays ~15-20 ms to wake up;
    warm, a 64 KB chunk takes ~1.5 ms. About 1% of the ANE's time; stops 8 s after the last request."""
    pings = {b: np.zeros((b, 1, 1, spec.SEQ), np.float16) for b in (32, 512)}
    while True:
        time.sleep(0.1)
        if time.time() - LAST_ACTIVE[0] < 8:
            with LOCK:
                for b, x in pings.items():
                    MODELS[b][1].predict({"codes": x})
TOTAL = sum(int(np.prod(v.shape)) for v in W.values())
NONZERO = sum(int((v != 0).sum()) for v in W.values())



def run(codes):
    """codes [N,1,1,SEQ] float -> (p, evidence, concepts, model label, ms)."""
    n = len(codes)
    b = 1 if n == 1 else (32 if n <= 32 else 512)
    unit, m = MODELS[b]
    ps, evs, cs = [], [], []
    t0 = time.perf_counter()
    with LOCK:
        for i in range(0, n, b):
            chunk = codes[i:i + b]
            x = np.zeros((b, 1, 1, spec.SEQ), np.float16)
            x[:len(chunk)] = chunk
            out = m.predict({"codes": x})
            k = len(chunk)
            ps.append(out["p_fake"][:k, 0]); evs.append(out["evidence"][:k]); cs.append(out["concepts"][:k])
    ms = (time.perf_counter() - t0) * 1e3
    if n == 0:
        return np.zeros(0), np.zeros((0, len(EVIDENCE))), np.zeros((0, len(CONCEPTS))), f"b{b} {unit}", 0.0
    return (np.concatenate(ps).astype(float), np.concatenate(evs).astype(float),
            np.concatenate(cs).astype(float), f"b{b} · {unit}", ms)


def describe(p, ev, c):
    act = {CONCEPTS[j]: float(c[j]) for j in np.nonzero(c > 0.5e-2)[0]}
    status = "phantom" if p >= 0.5 else ("unknown" if p > 0.1 else "ok")
    reasons = [dict(name=e, value=round(float(v), 2), text=EV_INFO[e]["text"],
                    logit=round(float(min(v, EV_INFO[e]["cap"]) * EV_INFO[e]["weight"]), 2))
               for e, v in zip(EVIDENCE, ev) if v > 0.01]
    vol = act.get("volume", 0.0)
    year = act.get("year", 0.0) + 1700 if act.get("has_year", 0) > 0.5 else None
    curve = next((n[3:] for n in act if n.startswith("is:") and n[3:] in spec.CURVES), None)
    fired = [n for n in act if not n.startswith(("vol>", "is:", "has_", "known", "volume", "year", "vol4y", "future"))
             and not n.endswith("^hi")]
    fired += [n for n in act if n.startswith("region:")]
    return dict(p=round(float(p), 4), status=status, reasons=reasons, volume=int(round(vol)) or None,
                year=int(round(year)) if year else None, curve=curve, fired=fired,
                knots=sum(1 for n in act if n.startswith("vol>")), logit_bias=spec.BIAS)


def filters_fired(cite):
    """Which first-layer filters fired (recomputed on the host from the same weights, for display)."""
    codes = text.encode([cite])[0, 0, 0].astype(int)
    onehot = np.zeros((spec.NCHAR, spec.SEQ), np.float32)
    onehot[codes, np.arange(spec.SEQ)] = 1
    out = []
    for conv, names, k, width in (("conv_rep", META["anchored"], K_REP, ANCHOR_W),
                                  ("conv_float", META["floating"], K_REP, spec.SEQ),
                                  ("conv_court", [n for n in META["features"] if n.startswith("court:")], K_COURT, spec.SEQ)):
        w, b = W[conv + ".weight"][:, :, 0, :], W[conv + ".bias"]
        for i, name in enumerate(names):
            for pos in range(width - k + 1):
                if (w[i] * onehot[:, pos:pos + k]).sum() + b[i] > 0.5:
                    out.append(dict(name=name, pos=pos - 1, len=int(-b[i] + 1)))
                    break
    return out


CLASS_DESC = {
    "SCOTUS": "a Supreme Court reporter", "FEDAPP": "a federal appeals reporter",
    "FEDAPP_OLD": "a federal appeals reporter", "FEDDIST": "a district-court reporter",
    "BANKR": "a bankruptcy reporter", "REGION": "a state-court reporter", "STATE": "a state-court reporter",
}
STATE_NAME = {"cal": "California", "n.y": "New York", "mass": "Massachusetts", "ill": "Illinois", "ind": "Indiana",
              "va": "Virginia", "pa": "Pennsylvania", "n.j": "New Jersey", "wash": "Washington", "ohio": "Ohio",
              "md": "Maryland", "n.c": "North Carolina", "ga": "Georgia", "kan": "Kansas", "mich": "Michigan",
              "wis": "Wisconsin", "conn": "Connecticut", "tex": "Texas", "s.c": "South Carolina",
              "r.i": "Rhode Island", "colo": "Colorado", "ariz": "Arizona", "or": "Oregon", "okla": "Oklahoma"}
REGION_NAME = {"A": "Atlantic", "NE": "North Eastern", "NW": "North Western", "P": "Pacific",
               "SE": "South Eastern", "SW": "South Western", "SO": "Southern"}


def explain(d, cite):
    """One plain sentence per evidence neuron that fired, filled in from the citation."""
    import re
    from citenet.net import curve_at
    m = text.CITE_RE.match(cite.strip())
    rep = m.group("rep") if m else "this reporter"
    paren = (m.group("paren") or "").strip(" ()") if m else ""
    court = re.sub(r"\s*(?:[A-Z][a-z]{2,4}\.? \d{1,2}, )?\d{4}$", "", paren).strip(" ,") or None
    act = set(d["fired"])
    cls = next((n[6:] for n in act if n.startswith("class:")), None)
    out = []
    for r in d["reasons"]:
        n = r["name"]
        if n == "phantom_series":
            out.append(f"{rep} was never published")
        elif n == "court_mismatch":
            out.append(f"{rep} is {CLASS_DESC.get(cls, 'a reporter')}; it never prints {court or 'that court'} opinions")
        elif n == "state_mismatch":
            reg = next((n2[7:] for n2 in act if n2.startswith("region:")), None)
            home = next((n2[5:] for n2 in act if n2.startswith("home:")), None)
            if reg:
                sts = [s.title() + ("" if s in spec.UNDOTTED else ".") for s in spec.REGION_STATES[reg]]
                lst = ", ".join(sts[:-1]) + " and " + sts[-1] if len(sts) <= 5 else ", ".join(sts[:4]) + " and others"
                out.append(f"{rep} ({REGION_NAME[reg]} Reporter) covers {lst}, not {court}")
            else:
                out.append(f"{rep} only prints {STATE_NAME.get(home, 'its own state')} cases, not {court}")
        elif n == "volume_year" and any(r2["name"] == "volume_missing" for r2 in d["reasons"]):
            continue
        elif n == "volume_year" and d["curve"]:
            c = spec.CURVES[d["curve"]]
            exp = round(curve_at(c["pts"], min(d["volume"], c["vmax"])))
            out.append(f"volume {d['volume']} of {rep} is from about {exp}, not {d['year']}")
        elif n == "volume_year":
            out.append(f"the Westlaw/Lexis number is from {d['volume']}, but the case is dated {d['year']}")
        elif n == "volume_missing":
            vmax = spec.CURVES[d["curve"]]["vmax"] if d["curve"] else None
            out.append(f"{rep} has no volume {d['volume']}; it ends around volume {vmax}" if vmax else f"{rep} has no volume {d['volume']}")
        elif n == "future_year":
            out.append(f"dated {d['year']}, which is in the future")
        elif n == "unknown_reporter":
            out.append(f"“{rep}” is not a reporter I recognize")
    return out


def check_text(t):
    t0 = time.perf_counter()
    found = text.extract(t)
    codes = text.encode([c for _, _, c in found])
    t1 = time.perf_counter()
    p, ev, c, label, ms = run(codes)
    items = []
    for (s, e, cite), pi, evi, ci in zip(found, p, ev, c):
        d = describe(pi, evi, ci)
        d.update(start=s, end=e, cite=cite, norm=text.normalize(cite))
        d["why"] = explain(d, cite)
        items.append(d)
    return dict(items=items, n=len(items), model=label, model_ms=round(ms, 3),
                tokenize_ms=round((t1 - t0) * 1e3, 3))


def curve_payload(key):
    c = spec.CURVES[key]
    a, coef = knot_readout(c["pts"], c["vmax"])
    vs = np.linspace(0, max(10, c["vmax"]), 120)
    net = [a + sum(w * max(0.0, v - float(n[4:])) for n, w in coef.items()) + 1700 for v in vs]
    return dict(key=key, tol=c["tol"], vmax=c["vmax"], anchors=c["pts"],
                net=[[float(v), float(y)] for v, y in zip(vs, net)],
                reporters=[a for a, _, _, k in spec.REPORTERS if k == key])


def _label(rs, labels):
    if len(rs) == 1:
        return labels[rs[0]]
    digits = [r for r in rs if spec.DIGIT0 <= r < spec.DIGIT0 + 10]
    letters = [r for r in rs if spec.LETTER_A <= r < spec.LETTER_A + 26]
    rest = [labels[r] for r in rs if r not in digits and r not in letters]
    out = []
    if len(digits) == 10:
        out.append("0-9")
    elif len(digits) >= 6:
        out.append("0-9*")
    else:
        out += [labels[r] for r in digits]
    if len(letters) >= 20:
        out.append("a-z" if len(letters) == 26 else "a-z*")
    else:
        out += [labels[r] for r in letters]
    if len(rest) > 6:
        return "other"
    return " ".join(out + rest)


def filter_payload(name):
    conv, names, k = None, None, None
    for cv, ns, kk in (("conv_rep", META["anchored"], K_REP), ("conv_float", META["floating"], K_REP),
                       ("conv_court", [n for n in META["features"] if n.startswith("court:")], K_COURT)):
        if name in ns:
            conv, names, k = cv, ns, kk
    if conv is None:
        return None
    i = names.index(name)
    w = W[conv + ".weight"][i, :, 0, :]
    labels = ["pad", "bos", "eos"] + [str(d) for d in range(10)] + [chr(97 + j) for j in range(26)] + \
        ["␣", ".", ",", "(", ")", "'", "&", "-", "?"]
    # merge characters with identical weight rows ("0-9", "a-z") and order rows by
    # their first column, so a template reads as a diagonal
    groups = {}
    for r in range(spec.NCHAR):
        if w[r].any():
            groups.setdefault(tuple(int(x) for x in w[r]), []).append(r)
    rows = []
    for vec, rs in sorted(groups.items(), key=lambda kv: (kv[0].index(1), kv[1][0])):
        rows.append(dict(ch=_label(rs, labels), w=list(vec)))
    return dict(name=name, k=k, bias=float(W[conv + ".bias"][i]), rows=rows, layer=conv,
                params=int(w.size) + 1, nonzero=int((w != 0).sum()))


# ---------------------------------------------------------------- samples and benchmarks
SAMPLES = {
    "brief": open(os.path.join(ROOT, "demo", "sample_brief.txt")).read(),
    "mata": open(os.path.join(ROOT, "demo", "sample_mata.txt")).read(),
}
EVAL_ROWS = [json.loads(l) for l in open(os.path.join(ROOT, "data", "eval.jsonl"))]
EVAL_CITES = load("eval")
ALL_CITES = load("aux_train") + EVAL_CITES
ALL_CODES = text.encode([r["cite"] for r in ALL_CITES]).astype(np.float16)


def lepc_sample():
    r = random.choice([r for r in EVAL_ROWS if any("non_existent_citation" in v
                                                   for v in (r.get("list_hallucination_types") or {}).values())])
    return dict(text=r["text"], file=r["filename"],
                planted={k: v for k, v in (r.get("list_hallucinations") or {}).items()})


def bench_eval():
    codes = text.encode([r["cite"] for r in EVAL_CITES])
    p, _, _, label, ms = run(codes)
    lab = np.array([r["label"] for r in EVAL_CITES]); flag = p >= 0.5
    tp = int((flag & (lab == "phantom")).sum()); P = int((lab == "phantom").sum())
    fp = int((flag & (lab == "genuine")).sum()); G = int((lab == "genuine").sum())
    flagged = [r["cite"] for r, f in zip(EVAL_CITES, flag) if f and r["label"] == "genuine"]
    return dict(n=len(EVAL_CITES), model=label, ms=round(ms, 2), recall=[tp, P], false_flags=[fp, G],
                false_flag_cites=flagged)


FIREHOSE = os.path.join(ROOT, "bench", "firehose")


def firehose(n, lanes):
    """Native Swift Core ML firehose (bench/firehose.swift): Python threads cannot keep the
    ANE and the GPU busy at the same time, a native process can."""
    import subprocess
    if lanes not in ("ane", "gpu", "ane+gpu"):
        lanes = "ane+gpu"
    with LOCK:
        out = subprocess.run([FIREHOSE, BUILD, lanes, str(n)], capture_output=True, text=True, timeout=120)
    return json.loads(out.stdout)


STREAM = os.path.join(ROOT, "bench", "stream_bin")
CORPUS = os.path.join(ROOT, "data", "cold", "corpus.txt")
EV_SHORT = ["never-published series", "court mismatch", "state mismatch", "year doesn't fit volume",
            "volume doesn't exist", "future date", "unknown reporter"]


def corpus_info():
    try:
        return json.load(open(CORPUS.rsplit(".", 1)[0] + ".json"))
    except OSError:
        return None


def live_stream(handler, loop):
    """Run bench/stream_bin over the case-law corpus; relay its JSON lines as server-sent events."""
    import subprocess
    cmd = [STREAM, BUILD, CORPUS] + (["--loop"] if loop else [])
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1)
    handler.send_response(200)
    handler.send_header("Content-Type", "text/event-stream")
    handler.send_header("Cache-Control", "no-store")
    handler.end_headers()
    try:
        for line in proc.stdout:
            ev = json.loads(line)
            for c in ev.get("caught", []):
                items = check_text(c["cite"])["items"]
                c["why"] = "; ".join(items[0]["why"]) if items else EV_SHORT[c["ev"]]
                c["kind"] = EV_SHORT[c["ev"]]
            handler.wfile.write(b"data: " + json.dumps(ev).encode() + b"\n\n")
            handler.wfile.flush()
            if ev.get("done"):
                break
    except (BrokenPipeError, ConnectionResetError):
        pass
    finally:
        proc.kill()
        proc.wait()


# ---------------------------------------------------------------- text wall
import mmap
WALL_SRC = {"cold": CORPUS, "lepc": os.path.join(BUILD, "lepc_text.txt")}
_maps = {}


def _wall_map(src):
    if src not in _maps:
        f = open(WALL_SRC[src], "rb")
        _maps[src] = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    return _maps[src]


def wall_random(src):
    """Byte offset of the start of a random decision / brief excerpt."""
    src = "cold" if src == "planted" else src
    m = _wall_map(src)
    off = random.randrange(0, max(1, len(m) - 100_000))
    i = m.find(b"\n\n\n" if src == "cold" else b"\n\n", off)
    return dict(off=0 if i < 0 else i + (3 if src == "cold" else 2))


# ---- planted fakes: mutate real citations into impossible ones, then let the model find them
import re as _re
CURVE_REPS = {a for a, _, _, k in spec.REPORTERS if k and k in spec.CURVES}
FIRST_SERIES = {"F.", "A.", "P.", "So.", "N.E.", "N.W.", "S.E.", "S.W."}
FED_APP = {"f.3d", "f.2d", "f.4th", "f.app'x", "fed.appx."}
STATE_COURTS = ["Tex. App.", "Cal. Ct. App.", "N.Y. App. Div.", "Fla. Dist. Ct. App.", "Ill. App. Ct."]
_SERIES = _re.compile(r"^(?P<base>.*?)(?P<sp>\s?)(?P<ord>2d|3d|4th|5th|2nd|3rd)$")


PLANT_RATE = 0.125    # target share of planted citations (a quota: unmutable ones pass it to the next)


def plant(cite, rng):
    """One impossible variant of a real citation, or None if none of the mutations applies."""
    m = text.CITE_RE.match(cite)
    if not m:
        return None
    rep = m.group("rep")
    norm = text.normalize(rep + " 1").rsplit(" ", 1)[0]
    paren = m.group("paren") or ""
    yr = _re.search(r"(?<=[\s(.,])(\d{4})\)\s*$", paren)          # a real year, not "*7771985"
    opts = []
    sm = None if "A.L.R" in rep or "ALR" in rep else _SERIES.match(rep)      # A.L.R. really runs to 7th
    if sm:
        bump = {"2d": "4th", "2nd": "4th", "3d": rng.choice(["4th", "5th"]), "3rd": "5th", "4th": "6th", "5th": "7th"}[sm["ord"]]
        if not (norm in ("f.3d", "f.supp.3d") and bump == "4th"):        # F.4th is real
            opts += [("series", rep[:sm.start("ord")] + bump)] * 3
    elif rep in FIRST_SERIES:      # F.4th is real, so the first Federal Reporter jumps to 5th/6th
        opts += [("series", rep + (rng.choice(["5th", "6th"]) if rep == "F." else rng.choice(["4th", "5th"])))] * 3
    if yr and norm in CURVE_REPS:
        y = int(yr.group(1))
        ny = y + rng.choice([-1, 1]) * rng.randint(18, 40)
        if not 1800 <= ny <= 2024:
            ny = y - rng.randint(18, 40) if ny > 2024 else y + rng.randint(18, 40)
        if 1800 <= ny <= 2024:
            opts += [("year", None, ny)] * 2
    if paren and "Cir." in paren and norm in FED_APP:
        opts.append(("court", rng.choice(STATE_COURTS)))
    if not opts:
        return None
    kind, *val = rng.choice(opts)
    if kind == "series":
        a, b = m.span("rep")
        fake_rep = val[0]
        # never produce a series that really exists (F.4th, Cal. 4th/5th, Cal. App. 4th/5th)
        real = set(spec.VALID_4TH) | set(spec.VALID_5TH)
        if text.normalize(fake_rep + " 1").rsplit(" ", 1)[0] in real:
            fake_rep = _re.sub(r"(4th|5th)$", "6th", fake_rep)
        return cite[:a] + fake_rep + cite[b:]
    if kind == "year":
        a = m.start("paren") + yr.start(1)
        return cite[:a] + str(val[1]) + cite[a + 4:]
    if not yr:
        return None
    new_paren = paren[:paren.index("(") + 1] + val[0] + " " + yr.group(1) + ")"
    return cite[:m.start("paren")] + new_paren + cite[m.end("paren"):]


def plant_text(t, rate, seed):
    """Replace ~rate of the citations in t with planted fakes. Returns (new text, planted ranges)."""
    rng = random.Random(seed)
    out, ranges, pos = [], [], 0
    owed = rng.random()                     # quota carried across citations, so density stays near `rate`
    for s_, e_, cite in text.extract(t):
        owed = min(owed + rate, 3.0)
        if owed < 1 or rng.random() < 0.3:  # a little jitter so fakes are not evenly spaced
            continue
        fake = plant(cite, rng)
        if not fake or fake == cite:
            continue
        owed -= 1
        out.append(t[pos:s_])
        start = sum(len(x) for x in out)
        out.append(fake)
        ranges.append((start, start + len(fake)))
        pos = e_
    out.append(t[pos:])
    return "".join(out), ranges


def wall_chunk(src, off, size):
    """The next ~size bytes of text from `off` (ending on a line break), with every citation
    checked by the Core ML model. Returns text plus spans [start, end, status, why]."""
    planting = src == "planted"
    if planting:
        src = "cold"
    m = _wall_map(src)
    n = len(m)
    off %= n
    end = m.find(b"\n", min(n, off + size))
    end = n if end < 0 else end + 1
    raw = m[off:end]
    nxt = end if end < n else 0                     # loop at the end of the corpus
    t = raw.decode("utf-8", errors="replace").replace("\r\n", "\n").replace("\r", "\n").replace("\x0c", "\n")
    ranges = []
    if planting:                                   # the model below only ever sees the mutated text
        t, ranges = plant_text(t, PLANT_RATE, seed=off)
    t0 = time.perf_counter()
    found = text.extract(t)
    codes = text.encode([c for _, _, c in found])
    t1 = time.perf_counter()
    p, ev, c, label, ms = run(codes)
    spans, flagged, caught, missed = [], 0, 0, 0
    for (s_, e_, cite), pi, evi, ci in zip(found, p, ev, c):
        status = "phantom" if pi >= 0.5 else ("unknown" if pi > 0.1 else "ok")
        why = ""
        if status != "ok":
            d = describe(pi, evi, ci)
            why = "; ".join(explain(d, cite))
            flagged += status == "phantom"
        planted = any(a < e_ and s_ < b for a, b in ranges)
        caught += planted and status == "phantom"
        missed += planted and status != "phantom"
        spans.append([s_, e_, status, why, planted])
    return dict(text=t, spans=spans, next=nxt, bytes=len(raw), n=len(found), flagged=flagged,
                planted=len(ranges), caught=caught, missed=missed,
                model=label, model_ms=round(ms, 3), tokenize_ms=round((t1 - t0) * 1e3, 3))


SWIFT = open(os.path.join(ROOT, "results", "bench_swift.txt")).read()


def swift_headline():
    """Pull the header-card numbers out of the native benchmark log."""
    import re
    out = {}
    for line in SWIFT.splitlines():
        m = re.match(r"b(\d+)\s+(\w+)\s+p50\s+([\d.]+) us.*?(\d+) cites/s", line)
        if m:
            out[f"b{m[1]}_{m[2]}"] = dict(p50_us=float(m[3]), per_sec=int(m[4]))
        m = re.match(r"concurrent (.*) -> (\d+) cites/s", line)
        if m:
            out["concurrent:" + re.sub(r" [\d.]+k/s", "", m[1])] = int(m[2])
    return out


def info():
    return dict(params=TOTAL, nonzero=NONZERO, evidence=[dict(name=e[0], weight=e[1], cap=e[2], text=e[3]) for e in spec.EVIDENCE],
                bias=spec.BIAS, swift=SWIFT, headline=swift_headline(),
                layers=dict(rep=len([n for n in META["features"] if n.startswith(("rep:", "series:", "other:"))]),
                            court=len([n for n in META["features"] if n.startswith("court:")]),
                            concepts=len(CONCEPTS), rules=len(META["rules"]), evidence=len(EVIDENCE)),
                curves=sorted(spec.CURVES))


class H(BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass

    def send(self, obj, code=200, ctype="application/json"):
        body = obj if isinstance(obj, bytes) else json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        path, _, q = self.path.partition("?")
        args = dict(kv.split("=", 1) for kv in q.split("&") if "=" in kv)
        from urllib.parse import unquote
        if path in ("/", "/index.html"):
            return self.send(open(os.path.join(ROOT, "demo", "index.html"), "rb").read(), ctype="text/html; charset=utf-8")
        if path == "/wall":
            return self.send(open(os.path.join(ROOT, "demo", "wall.html"), "rb").read(), ctype="text/html; charset=utf-8")
        if path == "/wall/random":
            return self.send(wall_random(args.get("src", "cold")))
        if path == "/wall/chunk":
            LAST_ACTIVE[0] = time.time()
            return self.send(wall_chunk(args.get("src", "cold"), int(args.get("off", 0)),
                                        min(int(args.get("size", 16384)), 1 << 20)))
        if path == "/live":
            return self.send(open(os.path.join(ROOT, "demo", "live.html"), "rb").read(), ctype="text/html; charset=utf-8")
        if path == "/live/info":
            return self.send(corpus_info() or {"error": "no corpus"})
        if path == "/live/stream":
            return live_stream(self, args.get("loop") == "1")
        if path == "/inspect":          # detailed view: neurons, filters, volume/year curves
            return self.send(open(os.path.join(ROOT, "demo", "inspect.html"), "rb").read(), ctype="text/html; charset=utf-8")
        if path == "/info":
            return self.send(info())
        if path == "/sample":
            kind = args.get("kind", "brief")
            return self.send(lepc_sample() if kind == "lepc" else dict(text=SAMPLES[kind]))
        if path == "/curve":
            return self.send(curve_payload(unquote(args["key"])))
        if path == "/filter":
            f = filter_payload(unquote(args["name"]))
            return self.send(f) if f else self.send({"error": "no such filter"}, 404)
        if path == "/fired":
            return self.send(filters_fired(unquote(args["cite"])))
        if path == "/bench":
            return self.send(bench_eval())
        if path == "/firehose":
            return self.send(firehose(min(int(args.get("n", 1000000)), 20_000_000), unquote(args.get("lanes", "ane+gpu"))))
        self.send({"error": "not found"}, 404)

    def do_POST(self):
        n = int(self.headers.get("Content-Length", 0))
        body = json.loads(self.rfile.read(n) or b"{}")
        if self.path == "/check":
            return self.send(check_text(body.get("text", "")))
        self.send({"error": "not found"}, 404)


if __name__ == "__main__":
    run(text.encode(["347 U.S. 483 (1954)"]))                     # warm up all three models
    run(text.encode(["347 U.S. 483 (1954)"] * 20))
    run(text.encode(["347 U.S. 483 (1954)"] * 100))
    threading.Thread(target=keep_warm, daemon=True).start()
    print(f"CiteNet demo on http://localhost:{PORT}  ({TOTAL:,} params, {NONZERO:,} nonzero, all hand-written)")
    ThreadingHTTPServer(("127.0.0.1", PORT), H).serve_forever()