#!/usr/bin/env python3
"""
Linear SPV chain · WBS ground · WBF measure · SDF_OUT boundaries
Never exit · auto hotswap · RFC → push_constants
"""
from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import hotswap_core as H

RFC_DIR = H.ROOT / "rfc"
CHAIN_STATE = H.CUR / "chain.json"


def bgf(a: int, b: int) -> int:
    return (a - b) | 1


def spv3(a: int, b: int, c: int) -> int:
    return (a ^ b ^ c) | 1


def bgs1(v: int) -> int:
    return v | 1


def ensure() -> None:
    H.ensure_dirs()
    RFC_DIR.mkdir(parents=True, exist_ok=True)


def list_rfcs() -> List[Path]:
    ensure()
    return sorted(RFC_DIR.glob("*.rfc.json"))


def load_rfc(name: str) -> Dict[str, Any]:
    ensure()
    # name may be 0001, 0001-linear-atoms, or full path
    candidates = [
        RFC_DIR / name,
        RFC_DIR / f"{name}.rfc.json",
        RFC_DIR / f"{name}.json",
    ]
    # prefix match
    for p in list_rfcs():
        if name in p.name:
            candidates.insert(0, p)
    for p in candidates:
        if p.exists():
            return json.loads(p.read_text())
    raise SystemExit(f"RFC not found: {name} · try wb rfc-list")


def parse_push_constants(rfc: Dict[str, Any], step: Dict[str, Any], state: Dict[str, Any]) -> Dict[str, Any]:
    """Merge schema defaults + step push + WBS/WBF live → push_constants block."""
    schema = rfc.get("push_constants_schema") or {}
    push = {
        "seed": 7,
        "step": int(state.get("step_i", 0)),
        "wbs": int((rfc.get("wbs") or {}).get("value", 1)) | 1,
        "wbf": int(state.get("wbf", 0)) | 1,
        "handoff_field": int(state.get("handoff_field", 1)),
        "handoff_free": int(state.get("handoff_free", 1)) | 1,
    }
    # schema keys get typed defaults
    for k, typ in schema.items():
        if k in push:
            continue
        if typ in ("u32", "i32"):
            push[k] = 0
        elif typ == "f32":
            push[k] = 0.0
        else:
            push[k] = 0
    # step overrides
    step_push = step.get("push") or {}
    for k, v in step_push.items():
        if k == "handoff":
            continue
        push[k] = v
    # force live chain fields
    push["step"] = int(state.get("step_i", 0))
    push["wbf"] = int(state.get("wbf", 0)) | 1
    push["wbs"] = bgs1(int((rfc.get("wbs") or {}).get("value", 1)))
    push["handoff_field"] = int(state.get("handoff_field", 1))
    push["handoff_free"] = int(state.get("handoff_free", 1)) | 1
    return push


def load_chain_state() -> Dict[str, Any]:
    ensure()
    if not CHAIN_STATE.exists():
        return {
            "rfc": None,
            "step_i": 0,
            "prev_step_i": 0,
            "wbs": 1,
            "wbf": 1,
            "never_exit": True,
            "linear": True,
            "handoff_field": 1,
            "handoff_free": 1,
            "ticks": 0,
            "last_kind": None,
            "running": False,
        }
    return json.loads(CHAIN_STATE.read_text())


def save_chain_state(st: Dict[str, Any]) -> None:
    ensure()
    st["updated"] = time.strftime("%Y-%m-%dT%H:%M:%S")
    st["free"] = spv3(st.get("step_i", 0), st.get("wbf", 0), st.get("ticks", 0))
    CHAIN_STATE.write_text(json.dumps(st, indent=2) + "\n")
    # also write push_constants for host/GPU consumers
    (H.CUR / "push_constants.json").write_text(
        json.dumps(st.get("push_constants") or {}, indent=2) + "\n"
    )


def start_chain(rfc_name: str) -> Dict[str, Any]:
    rfc = load_rfc(rfc_name)
    if rfc.get("chain") != "linear":
        raise SystemExit("only linear chains supported · always linear")
    st = {
        "rfc": rfc.get("rfc") or rfc_name,
        "rfc_title": rfc.get("title"),
        "step_i": 0,
        "prev_step_i": 0,
        "wbs": bgs1(int((rfc.get("wbs") or {}).get("value", 1))),
        "wbf": 1,
        "never_exit": bool(rfc.get("never_exit", True)),
        "linear": True,
        "handoff_field": 1,
        "handoff_free": 1,
        "ticks": 0,
        "last_kind": None,
        "running": True,
        "steps_n": len(rfc.get("steps") or []),
    }
    st["push_constants"] = parse_push_constants(rfc, (rfc.get("steps") or [{}])[0], st)
    save_chain_state(st)
    # apply first SPV step if needed
    return tick_chain(force_kind=None)


def _steps_map(rfc: Dict[str, Any]) -> Dict[int, Dict[str, Any]]:
    m = {}
    for s in rfc.get("steps") or []:
        m[int(s["i"])] = s
    return m


def _sdf_out(st: Dict[str, Any], rfc: Dict[str, Any], step: Dict[str, Any]) -> Dict[str, Any]:
    """SDF boundary · Host normal between SPVs · measure WBF · handoff free."""
    prev = int(st.get("prev_step_i", 0))
    cur = int(st.get("step_i", 0))
    # WBF = measure linear step distance · always |1 live (non-neg span)
    dist = cur - prev if cur >= prev else (prev - cur)
    st["wbf"] = dist | 1
    st["prev_step_i"] = cur
    # handoff field = BGF(cur+1, prev) style measure · free separate
    st["handoff_field"] = bgf(cur + 1, prev)
    st["handoff_free"] = spv3(st.get("handoff_field", 1), st.get("ticks", 0), st.get("wbs", 1))
    st["last_kind"] = "SDF_OUT"
    # linear next
    nxt = step.get("next")
    if nxt is None:
        # default linear +1, wrap if never_exit
        nxt = cur + 1
        steps_n = len(rfc.get("steps") or [])
        if nxt >= steps_n:
            if st.get("never_exit", True):
                nxt = 0
            else:
                nxt = cur  # would exit — we refuse: stay
    st["step_i"] = int(nxt)
    st["push_constants"] = parse_push_constants(
        rfc, _steps_map(rfc).get(int(nxt), {}), st
    )
    return st


def _run_spv(st: Dict[str, Any], rfc: Dict[str, Any], step: Dict[str, Any]) -> Dict[str, Any]:
    """Auto hotswap plug into socket · never exit host · then linear advance."""
    plug = step.get("plug")
    socket = step.get("socket") or "ui.stage"
    if not plug:
        raise SystemExit("SPV step missing plug")
    if socket == "ui.stage":
        if plug not in H.plugs_by_id():
            raise SystemExit(f"UI plug missing: {plug}")
        H.seat_ui(plug, record_history=True)
    elif socket == "gpu.compute":
        H.seat_gpu(plug, record_history=True)
    else:
        raise SystemExit(f"unknown socket {socket}")
    st["last_kind"] = "SPV"
    st["last_plug"] = plug
    st["last_socket"] = socket
    st["push_constants"] = parse_push_constants(rfc, step, st)
    (H.CUR / "step.json").write_text(
        json.dumps(
            {
                "kind": "SPV",
                "plug": plug,
                "socket": socket,
                "push_constants": st["push_constants"],
                "end": step.get("end", "SDF_OUT"),
                "wbs": st.get("wbs"),
                "wbf": st.get("wbf"),
            },
            indent=2,
        )
        + "\n"
    )
    # always linear: after SPV, next index (usually SDF_OUT step)
    st["step_i"] = int(step["i"]) + 1
    sm = _steps_map(rfc)
    if st["step_i"] not in sm and st.get("never_exit", True):
        st["step_i"] = 0
    return st


def tick_chain(force_kind: Optional[str] = None) -> Dict[str, Any]:
    """
    One linear step of the chain.
    Host never exits · SPV runs · SDF_OUT chains to next SPV.
    """
    st = load_chain_state()
    if not st.get("rfc"):
        raise SystemExit("no chain · wb chain-start <rfc>")
    if not st.get("running", True) and st.get("never_exit", True):
        st["running"] = True  # refuse permanent exit
    rfc = load_rfc(st["rfc"])
    sm = _steps_map(rfc)
    step = sm.get(int(st.get("step_i", 0)))
    if not step:
        # wrap linear
        st["step_i"] = 0
        step = sm.get(0)
        if not step:
            raise SystemExit("empty RFC steps")
    st["ticks"] = int(st.get("ticks", 0)) + 1
    kind = force_kind or step.get("kind")
    if kind == "SPV":
        st = _run_spv(st, rfc, step)
    elif kind == "SDF_OUT":
        st = _sdf_out(st, rfc, step)
    else:
        raise SystemExit(f"unknown step kind {kind}")
    # WBS always live ground
    st["wbs"] = bgs1(int((rfc.get("wbs") or {}).get("value", 1)))
    st["never_exit"] = True
    st["linear"] = True
    st["running"] = True
    save_chain_state(st)
    return st


def auto_run(ticks: int = 1, sleep_s: float = 0.0) -> List[Dict[str, Any]]:
    """Run N linear ticks · still never sets exit · host remains."""
    out = []
    for _ in range(max(1, ticks)):
        st = tick_chain()
        out.append(
            {
                "step_i": st.get("step_i"),
                "last_kind": st.get("last_kind"),
                "last_plug": st.get("last_plug"),
                "wbf": st.get("wbf"),
                "wbs": st.get("wbs"),
                "push": st.get("push_constants"),
            }
        )
        if sleep_s > 0:
            time.sleep(sleep_s)
    return out


def chain_status() -> Dict[str, Any]:
    st = load_chain_state()
    seats = H.load_seats()
    return {
        "chain": st,
        "seats": seats.get("seats"),
        "push_constants": st.get("push_constants"),
        "law": "linear · never_exit · SDF_OUT between SPV · WBS/WBF · ride BGF",
    }