"""alephllm-chat — talk to Beatrix.

A live window onto the AlephLLM training runs: loads the newest
checkpoint from the training repo, serves streaming completion and a
chat-format preview on ZeroGPU (with KV-cached decode and an A/B
toggle), and — with clear disclosure — logs conversations to a PUBLIC
research dataset. Code: https://github.com/AbstractEyes/alephllm
"""
import json
import os
import re
import threading
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from huggingface_hub import HfApi, hf_hub_download
from safetensors.torch import load_file

try:
    import spaces
    GPU = spaces.GPU
except Exception:  # local testing without the spaces runtime
    def GPU(fn=None, **kw):
        return fn if fn is not None else (lambda f: f)

import geolip.alephllm as _al
print(f"[boot] geolip.alephllm {_al.__version__}")
from geolip.alephllm.presets import PRESETS, AlephLMConfig
from geolip.alephllm.model.alephlm import AlephLM

REPO = "AbstractPhil/alephllm-mini-beatrix-training"
CRAFT = "mini-beatrix-1"
HISTORY_REPO = "AbstractPhil/alephllm-chat-history"
CHAT_HEADER = ("A conversation with Beatrix, a small byte-level "
               "language model still in pretraining.\n")

_lock = threading.Lock()
_state = {"model": None, "step": None, "manifest": None}

# ---------------------------------------------------------------- history log
_history_dir = Path("history")
_boot_id = uuid.uuid4().hex[:8]
_history_dir.mkdir(exist_ok=True)
_scheduler = None
if os.environ.get("HF_TOKEN"):
    try:
        from huggingface_hub import CommitScheduler
        _scheduler = CommitScheduler(
            repo_id=HISTORY_REPO, repo_type="dataset",
            folder_path=str(_history_dir), path_in_repo="data",
            every=5, token=os.environ["HF_TOKEN"], private=False)
    except Exception as e:  # noqa: BLE001
        print(f"[history] logging disabled: {e}")

LOG_NOTE = (
    f"🌐 **Public research log — read before typing.** Everything entered "
    f"here (prompts, chats, sampling settings, and Beatrix's replies) is "
    f"stored **publicly** in "
    f"[{HISTORY_REPO}](https://huggingface.co/datasets/{HISTORY_REPO}) "
    f"as JSON, tagged with an anonymous per-visit session id and the "
    f"checkpoint step. No accounts, names, IPs, or device data are "
    f"collected — but anything you type becomes public data, so **do not "
    f"enter personal information**."
    if _scheduler else
    "📴 Conversation logging is currently **disabled** (no storage token "
    "configured) — nothing you type here is recorded.")


def _log(kind: str, session: str, payload: dict):
    if _scheduler is None:
        return
    payload = dict(payload, kind=kind, session=session,
                   ts=datetime.now(timezone.utc).isoformat(),
                   craft=CRAFT, checkpoint_step=_state["step"])
    day = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    path = _history_dir / f"sessions-{day}-{_boot_id}.jsonl"
    with _scheduler.lock:
        with open(path, "a", encoding="utf-8", errors="replace") as f:
            f.write(json.dumps(payload, ensure_ascii=False) + "\n")


def _text(content) -> str:
    """gradio 6 Chatbot delivers turn content as a list of content-block
    dicts — normalize to plain text before the model or the public log
    ever sees it."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(b.get("text", "") for b in content
                       if isinstance(b, dict) and b.get("type") == "text")
    return str(content)


def _session(request: gr.Request | None) -> str:
    import hashlib
    try:
        raw = str(request.session_hash)
    except Exception:
        raw = uuid.uuid4().hex
    return hashlib.sha256(raw.encode()).hexdigest()[:16]


# ---------------------------------------------------------------- model load
def _latest_checkpoint():
    files = HfApi().list_repo_files(REPO)
    cks = sorted(f for f in files
                 if f.startswith(f"{CRAFT}/checkpoints/step_")
                 and f.endswith(".safetensors") and "/fp8/" not in f)
    return cks[-1] if cks else None


def _manifest():
    try:
        p = hf_hub_download(REPO, f"{CRAFT}/manifest.json",
                            force_download=True)
        return json.load(open(p, encoding="utf-8"))
    except Exception:
        return None


def load_model():
    """Never raises (a transient hub error at boot must not crash-loop the
    space); downloads happen OUTSIDE the lock (a reload mid-generation must
    not block streams); the device move happens HERE once — under ZeroGPU
    the module-level .to('cuda') is the documented supported pattern."""
    try:
        ck = _latest_checkpoint()
        if ck is None:
            return "No checkpoint on the hub yet — the run has not reached its first save."
        step = int(re.search(r"step_(\d+)", ck).group(1))
        if _state["step"] == step and _state["model"] is not None:
            _state["manifest"] = _manifest() or _state["manifest"]
            return _status_md()
        man = _manifest()
        cfg = (AlephLMConfig.from_dict(man["model_config"])
               if man and man.get("model_config") else PRESETS[CRAFT].model)
        model = AlephLM(cfg)
        sd = load_file(hf_hub_download(REPO, ck))
        model.load_state_dict(sd)
        model.eval()
        if torch.cuda.is_available():
            model.to("cuda")
        with _lock:
            _state.update(model=model, step=step, manifest=man)
        return _status_md()
    except Exception as e:  # noqa: BLE001
        return (f"**Status:** hub unreachable ({type(e).__name__}) — "
                "press *Reload newest checkpoint* to retry.")


def _status_md():
    man, step = _state["manifest"], _state["step"]
    if _state["model"] is None:
        return "**Status:** no checkpoint loaded — press *Reload newest checkpoint*."
    lines = [f"**Checkpoint:** step {step:,}"]
    if man:
        lines.append(f"**Tokens trained:** {man['tokens_seen']/1e9:.3f}B")
        ph = next((p for p in man["phases"] if p["status"] == "active"), None)
        if ph:
            lines.append(f"**Phase:** {ph['name']} ({ph['dataset']}), "
                         f"{ph.get('tokens_done', 0)/1e9:.2f}"
                         f"/{ph['planned_tokens']/1e9:.1f}B")
        bpbs = [c.get("val_bpb") for c in man.get("checkpoints", [])
                if c.get("val_bpb")]
        if bpbs:
            lines.append(f"**Val bits-per-byte:** {bpbs[-1]:.3f}")
    lines.append(f"**Params:** {_state['model'].param_count()/1e6:.1f}M · "
                 f"context {_state['model'].cfg.context} bytes")
    return "  \n".join(lines)


# ---------------------------------------------------------------- generation
@GPU
def _stream(prompt: str, max_new: int, temperature: float, top_p: float,
            use_cache: bool, stop: str | None = None):
    """Yields (text_so_far, stats_line). KV-cached by default; the
    uncached path recomputes the full context every byte for A/B timing."""
    model = _state["model"]
    if model is None:
        yield "(no checkpoint loaded — press Reload)", ""
        return
    dev = next(model.parameters()).device
    raw = np.frombuffer(prompt.encode("utf-8", errors="replace"),
                        dtype=np.uint8).astype(np.int64)
    keep = model.cfg.context - int(max_new) - 1
    ids = torch.tensor([raw[-keep:].tolist()], device=dev)
    out, t0 = [], time.time()
    with torch.no_grad():
        cache = None
        if use_cache:
            with torch.autocast("cuda", dtype=torch.bfloat16,
                                enabled=dev.type == "cuda"):
                logits, cache = model.prefill(ids)
            prefill_s = time.time() - t0
        t1 = time.time()
        for i in range(int(max_new)):
            if use_cache:
                nxt = model._sample(logits, temperature, top_p)
            else:
                with torch.autocast("cuda", dtype=torch.bfloat16,
                                    enabled=dev.type == "cuda"):
                    fl, _ = model(ids[:, -model.cfg.context:])
                nxt = model._sample(fl, temperature, top_p)
            ids = torch.cat([ids, nxt], dim=1)
            out.append(int(nxt.item()))
            if use_cache:
                with torch.autocast("cuda", dtype=torch.bfloat16,
                                    enabled=dev.type == "cuda"):
                    logits = model.decode_step(nxt, cache)
            text = bytes(out).decode("utf-8", errors="replace")
            bps = len(out) / max(time.time() - t1, 1e-6)
            stats = (f"kv-cache **on** · prefill {prefill_s*1000:.0f} ms · "
                     f"decode {bps:.0f} bytes/s" if use_cache else
                     f"kv-cache **off** (full recompute) · decode {bps:.0f} bytes/s")
            if stop and stop in text:
                yield text.split(stop)[0], stats
                return
            yield text, stats


def complete(prompt, max_new, temperature, top_p, use_cache,
             request: gr.Request):
    if not prompt or not prompt.strip():
        yield "(enter a prompt)", ""
        return
    text, done = "", False
    try:
        for text, stats in _stream(prompt, max_new, temperature, top_p,
                                   use_cache):
            yield prompt + text, stats
        done = True
    finally:  # disconnects/cancels still reach the public log, flagged
        _log("completion", _session(request),
             {"prompt": prompt, "output": text, "completed": done,
              "params": {"max_new": int(max_new), "temperature": temperature,
                         "top_p": top_p, "kv_cache": bool(use_cache)}})


def chat(message, history, max_new, temperature, top_p, use_cache,
         request: gr.Request):
    if not message or not message.strip():
        yield list(history or []), "", ""
        return
    history = [{"role": turn["role"], "content": _text(turn["content"])}
               for turn in (history or [])]
    transcript = CHAT_HEADER
    for turn in history:
        who = "User" if turn["role"] == "user" else "Beatrix"
        transcript += f"{who}: {turn['content']}\n"
    transcript += f"User: {message}\nBeatrix:"
    history = history + [{"role": "user", "content": message},
                         {"role": "assistant", "content": ""}]
    done = False
    try:
        for text, stats in _stream(transcript, max_new, temperature, top_p,
                                   use_cache, stop="\nUser:"):
            history[-1]["content"] = text.strip() or "…"
            yield history, "", stats
        done = True
    finally:
        _log("chat", _session(request),
             {"messages": history, "completed": done,
              "params": {"max_new": int(max_new), "temperature": temperature,
                         "top_p": top_p, "kv_cache": bool(use_cache)}})


DESCRIPTION = """
# Beatrix — a live AlephLLM training run

You are talking to **mini-beatrix-1** (112.5M parameters), an experimental
byte-level language model **currently in pretraining**. This space always
serves the newest checkpoint uploaded by the training run — she gets
smarter as you revisit.

**How she reads text.** No tokenizer: Beatrix consumes raw UTF-8 bytes.
Each position's embedding composes the byte with its two predecessors
(a trigram embedding, with a dedicated pad row for sequence starts), so
"tokens" are learned inside the network rather than fixed by a vocabulary.

**What's inside.** A 16-layer transformer where routing is done by
*signed geometric addresses* (the aleph mechanism): dispatch weights are
`sinh(u_k)/Σcosh(u_j)` over learned unit anchors — no softmax-over-choices,
no top-k, and negative (inhibitory) weights are first-class. Each layer
carries an anchored expert bank **born contributing exactly zero** (it must
earn its way in by gradient), three attention layers use a linear-cost
address read instead of softmax attention, and the output head carries a
second, address-based read behind a gate that also starts at zero.

**KV-cached decode.** Generation runs incrementally: the standard layers
keep a K/V cache, and the aleph attention layers keep a **constant-size
prefix state** regardless of context length — a structural property of the
address read. The toggle below lets you A/B cached vs full-recompute
decoding and watch the bytes/s difference live.

**What to expect right now.** Early-stage: under a billion tokens seen and
**no chat training yet** — the Chat tab wraps your words in a transcript
she continues, so expect drifty text, echoes, and format mimicry rather
than conversation. A chat-tuned stage will replace this preview once base
pretraining completes.

*Training logs, checkpoints and manifests:*
[alephllm-mini-beatrix-training](https://huggingface.co/AbstractPhil/alephllm-mini-beatrix-training)
· *code:* [github.com/AbstractEyes/alephllm](https://github.com/AbstractEyes/alephllm)
"""

with gr.Blocks(title="Beatrix — AlephLLM chat") as demo:
    gr.Markdown(DESCRIPTION)
    gr.Markdown(LOG_NOTE)
    with gr.Row():
        status = gr.Markdown(load_model())
        reload_btn = gr.Button("Reload newest checkpoint", scale=0)
    with gr.Row():
        max_new = gr.Slider(16, 1024, value=256, step=16, label="Max new bytes")
        temperature = gr.Slider(0.0, 1.5, value=0.8, step=0.05,
                                label="Temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
        use_cache = gr.Checkbox(value=True, label="KV cache")
    stats_md = gr.Markdown("")
    with gr.Tab("Completion"):
        prompt = gr.Textbox(label="Prompt", lines=4,
                            value="The history of mathematics begins")
        comp_out = gr.Textbox(label="Continuation (streams)", lines=12)
        gr.Button("Complete", variant="primary").click(
            complete, [prompt, max_new, temperature, top_p, use_cache],
            [comp_out, stats_md])
    with gr.Tab("Chat (preview — no chat training yet)"):
        gr.Markdown(LOG_NOTE)
        chatbot = gr.Chatbot(label="Beatrix", height=420)
        msg = gr.Textbox(label="Message", placeholder="Say something to Beatrix…")
        msg.submit(chat, [msg, chatbot, max_new, temperature, top_p, use_cache],
                   [chatbot, msg, stats_md])
        gr.ClearButton([chatbot, msg])
    reload_btn.click(load_model, None, status)

if __name__ == "__main__":
    demo.launch()