"""Scugnizz Llama-PCS interactive chat (custom PCS decoder, not transformers AutoModel)."""

from __future__ import annotations

import importlib.util
import gc
import json
import os
import shutil
import sys
import time
from pathlib import Path

import gradio as gr
import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

HUB_REPO = os.environ.get("HUB_REPO", "ProjectScugnizz/scugnizz-llama-pcs")
HUB_PATH = os.environ.get(
    "HUB_PATH",
    "training-runs/sft-chat-v2-ground-adhere-8b-20260807-053640",
)
SCRIPT_REPO = os.environ.get("SCRIPT_REPO", "ProjectScugnizz/scugnizz-llama-training")
MODEL_MOUNT = Path(os.environ.get("MODEL_MOUNT", "/models"))
CODE_MOUNT = Path(os.environ.get("CODE_MOUNT", "/code"))
SYSTEM = os.environ.get(
    "SYSTEM_PROMPT",
    "You are a helpful assistant. Answer clearly and correctly.",
)

_CODE_FALLBACK = Path("/tmp/scugnizz-code")
_MODEL_FALLBACK = Path("/tmp/scugnizz-model")
_WEIGHT_CACHE = Path("/tmp/scugnizz-weights")
_state = {
    "model": None,
    "tok": None,
    "dev": None,
    "error": None,
    "status": "cold",
    "hub_path": HUB_PATH.strip("/"),
    "weights_file": "model_final.pt",
}


def _active_hub_path() -> str:
    return (_state.get("hub_path") or HUB_PATH).strip("/")


def _active_weights_file() -> str:
    return (_state.get("weights_file") or "model_final.pt").strip() or "model_final.pt"


def _resolve_code_dir() -> Path:
    # Hub volume mounts truncate some text files (saw `import time` → `import tim`).
    # Trainer is tiny — always pull via Hub API.
    snapshot_download(
        SCRIPT_REPO,
        local_dir=str(_CODE_FALLBACK),
        allow_patterns=["scugnizz-llama.py", "sft_data.py"],
    )
    print(f"code from Hub → {_CODE_FALLBACK}", flush=True)
    return _CODE_FALLBACK


def _resolve_model_dir(hub_path: str | None = None) -> Path:
    hub_path = (hub_path or _active_hub_path()).strip("/")
    nested = MODEL_MOUNT / hub_path
    # Prefer path-specific mount (volume pointed at a run tree).
    if (nested / "args.json").is_file() or (nested / "model_final.pt").is_file():
        print(f"model mount nested: {nested}", flush=True)
        return nested
    # Legacy: Space volume mounts one run at /models root — only if it matches active path.
    env_path = os.environ.get("HUB_PATH", "").strip("/")
    if hub_path == env_path and (MODEL_MOUNT / "model_final.pt").is_file():
        print(f"model mount: {MODEL_MOUNT}", flush=True)
        return MODEL_MOUNT
    local = _MODEL_FALLBACK / hub_path
    # Re-download if missing weights or args
    need = not (local / "args.json").is_file()
    wname = _active_weights_file()
    if not (local / wname).is_file() and not (local / "model_final.pt").is_file():
        need = True
    if need:
        print(f"snapshot_download {HUB_REPO}/{hub_path} (~7GB)…", flush=True)
        snapshot_download(
            HUB_REPO,
            local_dir=str(_MODEL_FALLBACK),
            allow_patterns=[f"{hub_path}/*", f"{hub_path}/tokenizer/*"],
        )
    return local


def _load_train_module():
    code = _resolve_code_dir()
    path = code / "scugnizz-llama.py"
    spec = importlib.util.spec_from_file_location("scugnizz_train", path)
    mod = importlib.util.module_from_spec(spec)
    sys.path.insert(0, str(code))
    sys.modules["scugnizz_train"] = mod
    spec.loader.exec_module(mod)
    return mod


def _load_tokenizer(model_dir: Path, hub_path: str | None = None):
    hub_path = (hub_path or _active_hub_path()).strip("/")
    tok_dir = model_dir / "tokenizer"
    cfg = tok_dir / "tokenizer_config.json"
    try:
        if cfg.is_file() and cfg.stat().st_size > 1000:
            json.loads(cfg.read_text(encoding="utf-8"))
            return AutoTokenizer.from_pretrained(str(tok_dir))
    except (OSError, json.JSONDecodeError, ValueError):
        pass
    print("tokenizer via Hub (mount JSON unreliable)", flush=True)
    return AutoTokenizer.from_pretrained(HUB_REPO, subfolder=f"{hub_path}/tokenizer")


def _materialize_weights(src: Path, hub_path: str | None = None, weights_file: str | None = None) -> Path:
    """Copy once to local disk — Hub volume / FUSE is slow for torch.load."""
    hub_path = (hub_path or _active_hub_path()).strip("/")
    weights_file = weights_file or _active_weights_file()
    sz = src.stat().st_size
    if sz < 1_000_000_000:
        raise RuntimeError(f"weight file too small ({sz} B) at {src} — LFS/mount broken?")
    _WEIGHT_CACHE.mkdir(parents=True, exist_ok=True)
    safe = hub_path.replace("/", "__") + "__" + weights_file
    dest = _WEIGHT_CACHE / safe
    if dest.is_file() and dest.stat().st_size == sz:
        print(f"weights cache hit {dest} ({sz / 1e9:.2f} GB)", flush=True)
        return dest
    print(f"copying {sz / 1e9:.2f} GB → {dest} …", flush=True)
    t0 = time.perf_counter()
    shutil.copyfile(src, dest)
    print(f"copy done in {time.perf_counter() - t0:.1f}s", flush=True)
    return dest


def _read_json(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def _load_run_args(model_dir: Path, hub_path: str | None = None) -> dict:
    hub_path = (hub_path or _active_hub_path()).strip("/")
    local = model_dir / "args.json"
    try:
        if local.is_file() and local.stat().st_size > 50:
            return _read_json(local)
    except (OSError, json.JSONDecodeError, ValueError):
        pass
    print("args.json via Hub (mount JSON unreliable)", flush=True)
    from huggingface_hub import hf_hub_download

    p = hf_hub_download(HUB_REPO, f"{hub_path}/args.json")
    return _read_json(Path(p))


def _load_weights_into_model(model: torch.nn.Module, weight_path: Path, device: str) -> None:
    """t4-small has 15GB RAM — mmap + per-tensor copy, no full fp32 duplicate."""
    t_load = time.perf_counter()
    weights = torch.load(weight_path, map_location="cpu", weights_only=True, mmap=True)
    print(f"torch.load(mmap) {time.perf_counter() - t_load:.1f}s", flush=True)
    state = weights["model"] if isinstance(weights, dict) and "model" in weights else weights
    own = model.state_dict()
    t_copy = time.perf_counter()
    with torch.no_grad():
        for k, v in state.items():
            dst = own[k]
            if torch.is_floating_point(v):
                dst.copy_(v.to(device=dst.device, dtype=dst.dtype, non_blocking=True))
            else:
                dst.copy_(v.to(device=dst.device, non_blocking=True))
    if device == "cuda":
        torch.cuda.synchronize()
    print(f"param copy {time.perf_counter() - t_copy:.1f}s", flush=True)
    del weights, state
    gc.collect()


def _build_model(mod, cfg, device: str):
    # meta init avoids allocating a 7GB fp32 empty model in 15GB RAM
    with torch.device("meta"):
        model = mod.ScugnizzDecoder(cfg)
    model = model.to_empty(device=device)
    if device == "cuda":
        model = model.half()
    # to_empty may split tied weights
    model.tok_emb.weight = model.lm_head.weight
    return model


def _stop_token_ids(tok):
    ids = set()
    if getattr(tok, "eos_token_id", None) is not None:
        ids.add(int(tok.eos_token_id))
    unk = getattr(tok, "unk_token_id", None)
    for s in ("<|eot_id|>", "<|eom_id|>"):
        tid = tok.convert_tokens_to_ids(s)
        if tid is None or tid < 0:
            continue
        if unk is not None and tid == unk:
            continue
        ids.add(int(tid))
    return ids


def unload_model():
    if _state["model"] is not None:
        print("unloading model…", flush=True)
        _state["model"] = None
        _state["tok"] = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    _state["status"] = "cold"


def list_hub_runs():
    """List training-runs/* folders on the weights repo (newest-ish last alphabetically reversed)."""
    from huggingface_hub import HfApi

    api = HfApi()
    runs = []
    for item in api.list_repo_tree(
        HUB_REPO, path_in_repo="training-runs", recursive=False, repo_type="model"
    ):
        p = getattr(item, "path", "") or ""
        if p.startswith("training-runs/"):
            runs.append(p)
    # Prefer chat SFT near top: sort sft-chat first, then reverse chrono name
    def key(p):
        name = p.rsplit("/", 1)[-1]
        pri = 0 if "sft-chat" in name else 1
        return (pri, name)

    runs.sort(key=key, reverse=True)
    return runs


def list_weight_files(hub_path: str):
    """List *.pt weight files under a run folder."""
    from huggingface_hub import HfApi

    hub_path = (hub_path or "").strip("/")
    if not hub_path:
        return ["model_final.pt"]
    api = HfApi()
    files = []
    for item in api.list_repo_tree(
        HUB_REPO, path_in_repo=hub_path, recursive=False, repo_type="model"
    ):
        p = getattr(item, "path", "") or ""
        name = p.rsplit("/", 1)[-1]
        if not name.endswith(".pt"):
            continue
        size = getattr(item, "size", None) or 0
        if size and size < 1_000_000_000:
            continue
        files.append(name)
    if not files:
        files = ["model_final.pt"]
    # final first, then pulses by step
    def wkey(n):
        if n == "model_final.pt":
            return (0, 0)
        if n.startswith("model_pulse_"):
            try:
                return (1, -int(n.replace("model_pulse_", "").replace(".pt", "")))
            except ValueError:
                return (1, 0)
        return (2, n)

    files.sort(key=wkey)
    return files


def load_model(force: bool = False):
    if _state["model"] is not None and not force:
        return _state["model"], _state["tok"], _state["dev"]
    if force:
        unload_model()
    hub_path = _active_hub_path()
    weights_file = _active_weights_file()
    _state["error"] = None
    _state["status"] = f"loading {hub_path}/{weights_file}"
    t0 = time.perf_counter()
    try:
        mod = _load_train_module()
        model_dir = _resolve_model_dir(hub_path)
        run_args = _load_run_args(model_dir, hub_path)
        dev = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"device={dev} cuda={torch.cuda.is_available()}", flush=True)
        tok = _load_tokenizer(model_dir, hub_path)
        if tok.pad_token is None:
            tok.pad_token = tok.eos_token
        cfg = mod.preset_config(
            run_args.get("model_size", "1.7b"),
            len(tok),
            # Pretrain was 4096; chat SFT now 4096 — match pretrain/Gradio window
            int(os.environ.get("BLOCK_SIZE", "4096")),
            0.0,
            run_args.get("pcs_a", 0.8309193524478643),
            run_args.get("pcs_b", 0.0),
        )
        cfg.gradient_checkpointing = False
        model = _build_model(mod, cfg, dev)
        wpath = model_dir / weights_file
        if not wpath.is_file():
            # fall back to final if pulse missing locally but listed
            alt = model_dir / "model_final.pt"
            if weights_file != "model_final.pt" and alt.is_file():
                raise FileNotFoundError(f"{wpath} missing — try Refresh weights or model_final.pt")
            raise FileNotFoundError(f"missing weights: {wpath}")
        weight_path = _materialize_weights(wpath, hub_path, weights_file)
        _load_weights_into_model(model, weight_path, dev)
        model.eval()
        _state.update(model=model, tok=tok, dev=dev, status="ready")
        print(f"model ready in {time.perf_counter() - t0:.1f}s total", flush=True)
        return model, tok, dev
    except Exception as e:
        _state["error"] = str(e)
        _state["status"] = f"error: {e}"
        print(f"load failed: {e}", flush=True)
        raise


def switch_checkpoint(hub_path: str, weights_file: str):
    hub_path = (hub_path or "").strip("/")
    weights_file = (weights_file or "model_final.pt").strip() or "model_final.pt"
    if not hub_path:
        return f"**Status:** pick a run · current `{_active_hub_path()}/{_active_weights_file()}`"
    same = hub_path == _active_hub_path() and weights_file == _active_weights_file() and _state["model"] is not None
    if same:
        return f"**Status:** already loaded `{hub_path}/{weights_file}`"
    _state["hub_path"] = hub_path
    _state["weights_file"] = weights_file
    try:
        load_model(force=True)
        return f"**Status:** ready · `{hub_path}/{weights_file}`"
    except Exception as e:
        return f"**Status:** error loading `{hub_path}/{weights_file}`: {e}"

def _as_text(content) -> str:
    """Normalize Gradio message content (str | list blocks | ChatMessage-like)."""
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, str):
                parts.append(block)
            elif isinstance(block, dict):
                parts.append(str(block.get("text") or block.get("content") or ""))
            else:
                parts.append(str(getattr(block, "text", "") or ""))
        return "".join(parts)
    return str(content)


def _strip_ctx_footer(text: str) -> str:
    marker = "\n\n—\n*ctx "
    i = text.rfind(marker)
    return text[:i].rstrip() if i >= 0 else text


def history_to_messages(history):
    """Gradio 5 history: dicts, ChatMessage objects, or legacy tuples."""
    msgs = [{"role": "system", "content": SYSTEM}]
    if not history:
        return msgs
    for m in history:
        if isinstance(m, dict):
            role = m.get("role")
            content = _as_text(m.get("content"))
        elif hasattr(m, "role") and hasattr(m, "content"):
            role = getattr(m, "role", None)
            content = _as_text(getattr(m, "content", None))
        elif isinstance(m, (list, tuple)) and len(m) == 2:
            # legacy [user, assistant] pair
            u, a = _as_text(m[0]), _as_text(m[1])
            if u:
                msgs.append({"role": "user", "content": u})
            if a:
                msgs.append({"role": "assistant", "content": _strip_ctx_footer(a)})
            continue
        else:
            continue
        if role in ("user", "assistant") and content:
            if role == "assistant":
                content = _strip_ctx_footer(content)
            msgs.append({"role": role, "content": content})
    return msgs


def _prompt_ids(tok, msgs, budget: int):
    """Fit chat into block_size: drop oldest turns first, then left-trim latest user text."""
    system, rest = msgs[0], msgs[1:]
    if not rest:
        prompt = tok.apply_chat_template([system], tokenize=False, add_generation_prompt=True)
        ids = tok.encode(prompt, add_special_tokens=False)
        return prompt, ids[:budget], False

    dropped = 0
    for keep_from in range(0, len(rest)):
        candidate = [system] + rest[keep_from:]
        prompt = tok.apply_chat_template(
            candidate, tokenize=False, add_generation_prompt=True
        )
        ids = tok.encode(prompt, add_special_tokens=False)
        if len(ids) <= budget:
            if keep_from:
                print(f"context: dropped {keep_from} older turns → {len(ids)} tok", flush=True)
            return prompt, ids, bool(keep_from)
        dropped = keep_from + 1

    # Latest user message alone still too long — keep its tail (question usually at end)
    last = dict(rest[-1])
    content = last.get("content") or ""
    lo, hi = 0, len(content)
    best_prompt, best_ids = None, None
    while lo < hi:
        mid = (lo + hi) // 2
        last["content"] = content[mid:]
        prompt = tok.apply_chat_template(
            [system, last], tokenize=False, add_generation_prompt=True
        )
        ids = tok.encode(prompt, add_special_tokens=False)
        if len(ids) <= budget:
            best_prompt, best_ids = prompt, ids
            hi = mid
        else:
            lo = mid + 1
    if best_ids is None:
        last["content"] = content[-500:]
        best_prompt = tok.apply_chat_template(
            [system, last], tokenize=False, add_generation_prompt=True
        )
        best_ids = tok.encode(best_prompt, add_special_tokens=False)[-budget:]
    print(
        f"context: trimmed latest user (+ dropped {dropped} turns) → {len(best_ids)} tok",
        flush=True,
    )
    return best_prompt, best_ids, True


@torch.no_grad()
def _generate(model, tok, ids, max_new_tokens, temperature, top_k, device):
    model.eval()
    x = torch.tensor([ids], dtype=torch.long, device=device)
    greedy = temperature <= 0.0
    stop_ids = _stop_token_ids(tok)
    out_ids = []
    t0 = time.perf_counter()
    logits, _, kvs = model(x, use_cache=True)
    for _ in range(max_new_tokens):
        step_logits = logits[:, -1, :]
        if greedy:
            next_id = step_logits.argmax(dim=-1, keepdim=True)
        else:
            step_logits = step_logits / max(temperature, 1e-5)
            if top_k > 0:
                v, _ = torch.topk(step_logits, min(top_k, step_logits.size(-1)))
                step_logits = step_logits.masked_fill(step_logits < v[:, [-1]], float("-inf"))
            probs = F.softmax(step_logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
        tid = int(next_id.item())
        out_ids.append(tid)
        if tid in stop_ids:
            break
        if len(ids) + len(out_ids) >= model.cfg.block_size:
            break
        logits, _, kvs = model(next_id, past_kvs=kvs, use_cache=True)
    dt = max(time.perf_counter() - t0, 1e-6)
    print(f"gen {len(out_ids)} tok in {dt:.2f}s ({len(out_ids) / dt:.1f} tok/s)", flush=True)
    return tok.decode(out_ids, skip_special_tokens=True).strip()


def respond(message, history, max_new_tokens, temperature, top_k):
    model, tok, dev = load_model()
    msgs = history_to_messages(history)
    msgs.append({"role": "user", "content": _as_text(message)})
    lengths = [len(m["content"]) for m in msgs]
    print(
        f"msgs={len(msgs)} roles={[m['role'] for m in msgs]} chars={lengths} "
        f"temp={temperature} top_k={top_k} max_new={max_new_tokens}",
        flush=True,
    )
    for i, m in enumerate(msgs):
        print(f"--- msg[{i}] {m['role']} ({len(m['content'])} chars) ---", flush=True)
        print(m["content"], flush=True)
    budget = model.cfg.block_size - 1
    prompt, ids, truncated = _prompt_ids(tok, msgs, budget)
    print(
        f"prompt_tokens={len(ids)}/{budget} truncated={truncated}",
        flush=True,
    )
    print("--- prompt begin ---", flush=True)
    print(prompt, flush=True)
    print("--- prompt end ---", flush=True)
    text = _generate(
        model,
        tok,
        ids,
        int(max_new_tokens),
        float(temperature),
        int(top_k),
        dev,
    )
    print("--- assistant begin ---", flush=True)
    print(text, flush=True)
    print("--- assistant end ---", flush=True)
    note = f"\n\n—\n*ctx {len(ids)}/{budget} tok*" + (" *(trimmed)*" if truncated else "")
    return text + note


def build_ui():
    try:
        run_choices = list_hub_runs()
    except Exception as e:
        print(f"list_hub_runs failed: {e}", flush=True)
        run_choices = [_active_hub_path()]
    default_run = _active_hub_path()
    if default_run not in run_choices:
        run_choices = [default_run] + run_choices
    try:
        weight_choices = list_weight_files(default_run)
    except Exception:
        weight_choices = ["model_final.pt"]

    with gr.Blocks(title="Scugnizz Llama-PCS chat") as demo:
        status = gr.Markdown(
            f"**Status:** {_state.get('status', 'cold')} · `{_active_hub_path()}/{_active_weights_file()}`"
        )
        gr.Markdown(
            f"""# Scugnizz Llama-PCS chat
Custom ~1.7B PCS decoder · pick any run under `{HUB_REPO}`  
**Context window: 4096 tokens**. Each reply shows `ctx N/4095`.  
Switching checkpoints downloads ~7GB the first time — wait for **ready**.  
For grounded/open-book answers use **temperature 0–0.3** (default 0.2).
"""
        )
        with gr.Accordion("Checkpoint", open=True):
            run_dd = gr.Dropdown(
                choices=run_choices,
                value=default_run,
                label=f"Hub run ({HUB_REPO})",
                allow_custom_value=True,
            )
            weight_dd = gr.Dropdown(
                choices=weight_choices,
                value=_active_weights_file()
                if _active_weights_file() in weight_choices
                else weight_choices[0],
                label="Weights file",
                allow_custom_value=True,
            )
            with gr.Row():
                refresh_runs = gr.Button("Refresh list", scale=1)
                load_btn = gr.Button("Load checkpoint", variant="primary", scale=2)

            def on_run_change(hub_path):
                try:
                    ws = list_weight_files(hub_path)
                except Exception as e:
                    return gr.update(choices=["model_final.pt"], value="model_final.pt"), f"**Status:** weight list error: {e}"
                val = "model_final.pt" if "model_final.pt" in ws else ws[0]
                return gr.update(choices=ws, value=val), f"**Status:** select weights · `{hub_path}`"

            def on_refresh():
                try:
                    runs = list_hub_runs()
                except Exception as e:
                    return gr.update(), f"**Status:** refresh error: {e}"
                cur = _active_hub_path()
                if cur not in runs:
                    runs = [cur] + runs
                return gr.update(choices=runs, value=cur), f"**Status:** listed {len(runs)} runs"

            run_dd.change(on_run_change, [run_dd], [weight_dd, status])
            refresh_runs.click(on_refresh, None, [run_dd, status])
            load_btn.click(switch_checkpoint, [run_dd, weight_dd], status)

        chatbot = gr.Chatbot(height=480, type="messages")
        msg = gr.Textbox(
            placeholder="Ask something… (for open-book: paste passage, then question at the end)",
            scale=1,
        )
        with gr.Accordion("Generation", open=False):
            max_new = gr.Slider(16, 512, value=128, step=8, label="max new tokens")
            temp = gr.Slider(0.0, 1.5, value=0.2, step=0.05, label="temperature (0 = greedy)")
            top_k = gr.Slider(0, 200, value=50, step=1, label="top-k (0 = off)")
        clear = gr.Button("Clear")

        # Single handler — avoids Gradio queue race where .then() saw empty/stale history
        # (logs showed prompt_tokens≈48 while user pasted long context).
        def chat(user_msg, history, max_new_tokens, temperature, top_k):
            history = list(history or [])
            user_msg = _as_text(user_msg)
            if not user_msg.strip():
                return "", history
            prior = history
            try:
                answer = respond(user_msg, prior, max_new_tokens, temperature, top_k)
            except Exception as e:
                answer = f"(error loading/generating: {e})"
            history = prior + [
                {"role": "user", "content": user_msg},
                {"role": "assistant", "content": answer},
            ]
            return "", history

        def refresh_status():
            return (
                f"**Status:** {_state.get('status', 'cold')} · "
                f"`{_active_hub_path()}/{_active_weights_file()}`"
            )

        msg.submit(
            chat, [msg, chatbot, max_new, temp, top_k], [msg, chatbot]
        ).then(refresh_status, None, status)
        clear.click(lambda: [], None, chatbot, queue=False)
        demo.load(refresh_status, None, status)
    return demo


if __name__ == "__main__":
    print("preloading model before serving…", flush=True)
    try:
        load_model()
    except Exception as e:
        print(f"preload failed (will retry on first request): {e}", flush=True)
    demo = build_ui()
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        ssr_mode=False,
    )