"""Cairn's compute layer for the browser build (Pyodide, no server, no gradio).

Why this exists instead of Gradio-lite: gradio-lite imports ``gradio`` *before*
it installs the page's requirements, and gradio 5.x currently cannot be resolved
against ``huggingface-hub`` 1.x inside Pyodide -- the capped builds fail to
install, the uncapped ones import and then die on a missing ``httpcore``. Since
the boot order is not ours to change there, we skip the framework and drive
Pyodide directly. The page ends up lighter too: no pandas, pydantic or orjson,
just numpy, scipy, pillow and the ``cairn`` wheel.

Every function here returns plain JSON-able data (HTML fragments and base64 PNG
data URIs) which ``index.html`` drops into the DOM. All of the science is
imported unchanged from the ``cairn`` package; nothing is reimplemented for the
browser.

Claim: R/E -- lets anyone check the two headline claims themselves, for free,
with no GPU, no install and no account.
"""

from __future__ import annotations

import base64
import io
from typing import Any, Dict, List, Tuple

import numpy as np
from PIL import Image

from cairn.runner import CONDITION_NAMES, RunConfig, run_condition
from cairn.world import (
    make_departure_return_trajectory,
    make_scene,
    render,
    schedule_edit,
    usable_targets,
)

GAP = 6
BASELINES = {"A": "Vanilla", "B": "Context-window", "C": "Compressed-memory"}


# --------------------------------------------------------------------------
# rendering helpers
# --------------------------------------------------------------------------


def _png(arr: np.ndarray) -> str:
    """``(H, W, 3)`` float image -> base64 PNG data URI for an ``<img>`` tag."""
    a = (np.clip(np.asarray(arr), 0.0, 1.0) * 255).astype(np.uint8)
    buf = io.BytesIO()
    Image.fromarray(a).save(buf, format="PNG", optimize=True)
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii")


def _filmstrip(frames: np.ndarray, idx: List[int]) -> str:
    """Selected frames side by side, separated by pale gutters.

    Pyodide has no ffmpeg, so there is no video. A strip is arguably the better
    medium for this claim anyway: "before leaving" and "after returning" sit in
    one glance instead of several seconds apart in a loop.

    Claim: R -- the experiment, in one picture.
    """
    idx = [i for i in idx if 0 <= i < len(frames)]
    if not idx:
        return _png(np.zeros((8, 8, 3)))
    h = frames[0].shape[0]
    gutter = np.full((h, GAP, 3), 0.93, dtype=np.float32)
    panels: List[np.ndarray] = []
    for k, i in enumerate(idx):
        if k:
            panels.append(gutter)
        panels.append(frames[i])
    return _png(np.concatenate(panels, axis=1))


def _windows(traj) -> List[int]:
    return [
        0,
        max(0, traj.observe_frames[1] - 1),
        (traj.departure_frame + traj.return_frame) // 2,
        traj.return_frame + 1,
        min(len(traj) - 1, traj.return_frame + 7),
    ]


PANEL_CAPTION = (
    "opening shot &middot; last frame before leaving &middot; "
    "<b>looking away</b> &middot; just back &middot; settled after return"
)


def _ledger_html(res) -> str:
    rows = []
    for e in res.ledger.entries(include_absent=True):
        p = e.pose.position
        state = "present" if e.present else '<b style="color:#b3261e">REMOVED</b>'
        rows.append(
            f"<tr><td>{e.object_id}</td><td>{p[0]:.2f}</td><td>{p[2]:.2f}</td>"
            f"<td>{e.pose.yaw:+.2f}</td>"
            f"<td><span class='sw' style='background:rgb("
            f"{int(e.appearance[0]*255)},{int(e.appearance[1]*255)},{int(e.appearance[2]*255)})'></span>"
            f"{e.appearance[0]:.2f}, {e.appearance[1]:.2f}, {e.appearance[2]:.2f}</td>"
            f"<td>{e.n_observations}</td><td>{e.confidence:.2f}</td><td>{state}</td></tr>"
        )
    return (
        "<table><thead><tr><th>id</th><th>x</th><th>z</th><th>yaw</th><th>rgb</th>"
        "<th>seen</th><th>conf</th><th>state</th></tr></thead><tbody>"
        + "".join(rows)
        + "</tbody></table>"
    )


def _cfg(cond: str, seed: int) -> RunConfig:
    # reference_video=False skips rendering a second full clip that only the
    # FVD-proxy consumes, and this page never shows that number. Roughly halves
    # the work per click in the browser.
    return RunConfig(condition=cond, seed=1000 + int(seed), reference_video=False)


def _resolve_target(scene, requested: int, seed: int) -> Tuple[int, str]:
    """Snap the slider to an object that can actually host an episode.

    Claim: R -- the demo shows the same well-posed episodes the benchmark scores.
    """
    ok = usable_targets(scene, seed=int(seed))
    if not ok:
        raise ValueError("No object in this room can be left and returned to — try another seed.")
    if int(requested) in ok:
        return int(requested), ""
    chosen = min(ok, key=lambda o: abs(o - int(requested)))
    return chosen, (
        f"<p class='note'>Object #{int(requested)} is permanently hidden behind another object "
        f"in this room, so it cannot host a leave-and-return episode. Showing object "
        f"#{chosen} instead.</p>"
    )


# --------------------------------------------------------------------------
# public entry points (called from JavaScript)
# --------------------------------------------------------------------------


def scene_preview(seed: int, n_objects: int) -> str:
    """Four views of the room, so you can see what you are about to test."""
    from cairn.types import CameraPose

    scene = make_scene(int(n_objects), seed=int(seed))
    cams = [
        CameraPose(np.array([5.0, 1.55, 5.0]), a)
        for a in np.linspace(-np.pi, np.pi, 4, endpoint=False)
    ]
    frames = np.stack([render(scene.states(), c, scene.settings).rgb for c in cams])
    return _filmstrip(frames, list(range(4)))


def _verdict_html(res, label: str) -> str:
    m = res.metrics
    drawn = m["return_observed"] > 0.5
    ok = m["return_success"] > 0.5
    err = "not drawn at all" if not drawn else f"{m['return_self_trans']:.2f} m"
    badge = (
        "<span class='ok'>consistent</span>" if ok else "<span class='bad'>inconsistent</span>"
    )
    return (
        f"<div class='verdict'><h4>{label} &nbsp;{badge}</h4><ul>"
        f"<li>moved on return: <b>{err}</b></li>"
        f"<li>identity preserved: {'yes' if m['return_identity_preserved'] else '<b>no</b>'}</li>"
        f"<li>yaw error {m['return_self_yaw']:.2f} rad &middot; "
        f"appearance error {m['return_self_appearance']:.3f}</li>"
        f"<li>integrated consistency debt {m['debt_area']:.1f} (peak {m['debt_peak']:.2f})</li>"
        f"</ul></div>"
    )


def compare(seed: int, n_objects: int, absence: int, target: int, baseline: str) -> Dict[str, Any]:
    """Same scene, same trajectory, same generator seed — only the memory differs.

    Claim: R -- the interactive form of the headline experiment.
    """
    seed, n_objects, absence = int(seed), int(n_objects), int(absence)
    scene = make_scene(n_objects, seed=seed)
    target, note = _resolve_target(scene, int(np.clip(target, 0, n_objects - 1)), seed)
    traj = make_departure_return_trajectory(scene, target, absence, seed=seed)

    off = run_condition(scene, traj, _cfg(baseline, seed))
    on = run_condition(scene, traj, _cfg("D", seed))
    idx = _windows(traj)

    header = (
        f"<p>Camera left object <b>#{target}</b> at frame {traj.departure_frame} and came back "
        f"at frame {traj.return_frame} — <b>{traj.absence_frames} frames away</b>. "
        f"The return viewpoint is deliberately <i>not</i> the departure viewpoint, so neither "
        f"method can win by replaying its last frame.</p>"
        f"<p class='caption'>Panels: {PANEL_CAPTION}</p>{note}"
    )
    return {
        "header": header,
        "off_img": _filmstrip(off.frames, idx),
        "on_img": _filmstrip(on.frames, idx),
        "off_label": f"Cairn OFF — ({baseline}) {CONDITION_NAMES[baseline]}",
        "on_label": "Cairn ON — (D) explicit world ledger",
        "off_verdict": _verdict_html(off, f"Cairn OFF — ({baseline}) {BASELINES[baseline]}"),
        "on_verdict": _verdict_html(on, "Cairn ON — (D) explicit world ledger"),
        "ledger": _ledger_html(on),
    }


def edit(seed: int, n_objects: int, absence: int, target: int, kind: str) -> Dict[str, Any]:
    """Issue an edit while the object is off screen, then score what came back.

    Claim: E -- the operation conditions A–C cannot express at all.
    """
    seed, n_objects, absence = int(seed), int(n_objects), int(absence)
    scene = make_scene(n_objects, seed=seed)
    target, note = _resolve_target(scene, int(np.clip(target, 0, n_objects - 1)), seed)
    traj = make_departure_return_trajectory(scene, target, absence, seed=seed)
    ev = schedule_edit(traj, scene, kind, seed=seed)
    traj.edits = [ev]

    on = run_condition(scene, traj, _cfg("D", seed))
    off = run_condition(scene, traj, _cfg("A", seed))
    idx = _windows(traj)

    if kind == "move":
        d = float(np.linalg.norm(np.asarray(ev.payload["position"]) - scene.get(target).pose.position))
        what = f"move object #{target} {d:.1f} m and turn it"
    elif kind == "remove":
        what = f"delete object #{target} from the world"
    else:
        v = np.asarray(ev.payload["value"], dtype=float)
        what = (
            f"recolour object #{target} to "
            f"<span class='sw' style='background:rgb({int(v[0]*255)},{int(v[1]*255)},{int(v[2]*255)})'></span>"
            f"RGB {np.round(v, 2).tolist()}"
        )

    rows = []
    for res, cond in ((off, "A"), (on, "D")):
        for s in res.edit_scores:
            expressible = (
                "yes" if cond == "D" else "<b>no</b> — no addressable world state"
            )
            if s.complied:
                v = "<span class='ok'>yes</span>"
            elif s.ledger_correct:
                v = "<span class='warn'>written, not confirmable</span>"
            else:
                v = "<span class='bad'>no</span>"
            rows.append(
                f"<tr><td>({cond}) {CONDITION_NAMES[cond]}</td><td>{expressible}</td>"
                f"<td>{v}</td><td><code>{s.detail}</code></td></tr>"
            )

    header = (
        f"<p><b>Command:</b> {what}</p>"
        f"<p>Issued at frame <b>{ev.frame}</b>, while the object is off screen "
        f"(frames {traj.departure_frame}–{traj.return_frame}).</p>"
        f"<p class='caption'>Panels: {PANEL_CAPTION}</p>{note}"
    )
    table = (
        "<table><thead><tr><th>condition</th><th>can express it?</th><th>obeyed?</th>"
        "<th>evidence</th></tr></thead><tbody>" + "".join(rows) + "</tbody></table>"
        f"<p class='note'>Conditions A–C hold the world implicitly, in activations. There is no "
        f'row named "object #{target}" to write to, so <code>move</code>/<code>remove</code>/'
        f"<code>set_attr</code> are not merely hard for them — they are undefined.</p>"
    )
    return {
        "header": header,
        "off_img": _filmstrip(off.frames, idx),
        "on_img": _filmstrip(on.frames, idx),
        "table": table,
        "ledger": _ledger_html(on),
    }


def ledger_view(seed: int, n_objects: int, absence: int, target: int, rewind_to: int) -> Dict[str, Any]:
    """Show the ledger, its transaction log, and the effect of a rewind.

    Claim: E -- the world is a table with an audit trail; rewinding is one call.
    """
    seed, n_objects, absence = int(seed), int(n_objects), int(absence)
    scene = make_scene(n_objects, seed=seed)
    tgt, _ = _resolve_target(scene, int(np.clip(target, 0, n_objects - 1)), seed)
    traj = make_departure_return_trajectory(scene, tgt, absence, seed=seed)
    res = run_condition(scene, traj, _cfg("D", seed))
    led = res.ledger

    before = _ledger_html(res)
    log_rows = "".join(
        f"<tr><td>{t.index}</td><td>{t.t}</td><td><code>{t.op}</code></td>"
        f"<td>{t.object_id}</td><td>{t.source}</td><td>{t.note or ''}</td></tr>"
        for t in led.log[-40:]
    )
    log = (
        "<table><thead><tr><th>#</th><th>frame</th><th>op</th><th>object</th>"
        "<th>source</th><th>note</th></tr></thead><tbody>" + log_rows + "</tbody></table>"
    )
    v0 = led.version
    undone = led.rollback_to_time(int(rewind_to))
    note = (
        f"<p><b>Rewound to frame {int(rewind_to)}</b>: undid {undone} of {v0} transactions. "
        f"The ledger is now exactly as it stood at that instant, object for object. No learned "
        f"memory offers this operation — its state is entangled across every object and every "
        f"timestep at once.</p>"
    )
    return {"before": before, "log": log, "after": _ledger_html(res), "note": note}