"""
app.py — Pragya by Arush Kumar · Hugging Face Spaces (PyTorch backend)

Expects alongside app.py:
    model.py       – GPT / GPTConfig (training codebase, unchanged)
    inference.py   – generate(), load_model(), sampling helpers
    tokenizer.py   – TokenizerWrapper
    tokenizer/tokenizer.pkl – pickled tiktoken.Encoding
    checkpoint.pt  – trained checkpoint (from train.py)

Environment variables (Space Settings → Variables):
    CHECKPOINT_PATH   path to local checkpoint  (default ./checkpoint.pt)
    CHECKPOINT_REPO   HF Hub repo id            (e.g. yourname/pragya)
    CHECKPOINT_FILE   filename inside Hub repo  (default checkpoint.pt)
    TOKENIZER_PATH    path to tokenizer directory (default ./tokenizer)
    MAX_TOKENS_CAP    hard cap on slider         (default 512)
    DEFAULT_MAX_TOKENS default slider value      (default 128)
"""

from __future__ import annotations

import contextlib
import os
import queue
import sys
import time
import threading
from pathlib import Path

import torch
import gradio as gr

_HERE = Path(__file__).resolve().parent
if str(_HERE) not in sys.path:
    sys.path.insert(0, str(_HERE))

from inference import (         # noqa: E402
    generate, load_model, pick_device, _resolve_dtype,
)
from tokenizer import TokenizerWrapper  # noqa: E402

# ── env config ───────────────────────────────────────────────────────────────
_CHECKPOINT_LOCAL   = os.environ.get("CHECKPOINT_PATH",    str(_HERE / "pragya.pt"))
_HF_REPO            = os.environ.get("CHECKPOINT_REPO",    "")
_HF_FILE            = os.environ.get("CHECKPOINT_FILE",    "pragya.pt")
_TOKENIZER          = os.environ.get("TOKENIZER_PATH",     str(_HERE / "tokenizer.pkl"))
_MAX_TOKENS_CAP     = int(os.environ.get("MAX_TOKENS_CAP",      "512"))
_DEFAULT_TOKENS     = int(os.environ.get("DEFAULT_MAX_TOKENS",  "128"))


def _resolve_checkpoint() -> str:
    if _HF_REPO and _HF_FILE:
        try:
            from huggingface_hub import hf_hub_download
            print(f"[App] Pulling {_HF_REPO}/{_HF_FILE} from HF Hub …", flush=True)
            return hf_hub_download(repo_id=_HF_REPO, filename=_HF_FILE)
        except Exception as exc:
            print(f"[App] Hub download failed ({exc}), trying local.", flush=True)
    if Path(_CHECKPOINT_LOCAL).exists():
        return _CHECKPOINT_LOCAL
    raise FileNotFoundError(
        f"No checkpoint at {_CHECKPOINT_LOCAL}. "
        "Set CHECKPOINT_REPO + CHECKPOINT_FILE to pull from HF Hub."
    )


# ── one-time model init ───────────────────────────────────────────────────────
print("[App] Initialising …", flush=True)

_device         = pick_device(None)
_cast_dtype     = _resolve_dtype("auto", _device)
_ckpt_path      = _resolve_checkpoint()
_tokenizer      = TokenizerWrapper(_TOKENIZER)
_model, _cfg    = load_model(_ckpt_path, _device, _tokenizer)
_model.eval()

_n_params_str   = f"{_model.get_num_params() / 1e6:.1f}M"
_ctx_len        = _cfg.block_size
_pattern        = getattr(_cfg, "pattern", "dense").upper()
_n_layer        = getattr(_cfg, "n_layer",  "?")
_n_head         = getattr(_cfg, "n_head",   "?")
_kv_heads       = getattr(_cfg, "num_kv_heads", _n_head)
_gen_lock       = threading.Lock()

print(f"[App] Ready — {_n_params_str} params  ctx={_ctx_len}  device={_device}", flush=True)


# ── generation ────────────────────────────────────────────────────────────────
_ANCHOR = (
    "The old lighthouse stood at the edge of the cliff, its beam sweeping "
    "slowly across the dark water. Every night the keeper climbed the narrow "
    "stairs to check the lamp, and every night the sea answered with the steady "
    "sound of waves breaking against the rocks far below.\n\n"
)

def _cast_ctx():
    if _cast_dtype is not None and _device.type in ("cuda", "cpu"):
        return torch.autocast(device_type=_device.type, dtype=_cast_dtype)
    return contextlib.nullcontext()


def respond(
    message: str,
    history: list,
    max_new_tokens: int,
    temperature: float,
    top_k: int,
    top_p: float,
    repetition_penalty: float,
    no_repeat_ngram_size: int,
    use_anchor: bool,
):
    if not message.strip():
        yield "Please enter a prompt."
        return

    prompt = (_ANCHOR + message) if use_anchor else message
    t_start = time.perf_counter()
    ttft: list[float | None] = [None]
    n_pieces: list[int] = [0]

    def _stats() -> str:
        elapsed = max(time.perf_counter() - t_start, 1e-6)
        ttft_ms = (ttft[0] * 1000) if ttft[0] is not None else 0.0
        n = n_pieces[0]
        return (
            f"\n\n<span style='font-size:11px;font-family:\"IBM Plex Mono\",monospace;"
            f"color:#9ca3af;letter-spacing:0.03em'>"
            f"{n} tokens · TTFT {ttft_ms:.0f} ms · {n/elapsed:.1f} tok/s</span>"
        )

    q: queue.Queue = queue.Queue()
    done = object()

    def on_text(piece: str) -> None:
        if ttft[0] is None:
            ttft[0] = time.perf_counter() - t_start
        n_pieces[0] += 1
        q.put(piece)

    def runner() -> None:
        try:
            with _gen_lock, _cast_ctx():
                generate(
                    _model, _tokenizer, prompt, _device,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    top_k=top_k if top_k > 0 else None,
                    top_p=top_p if top_p < 1.0 else None,
                    repetition_penalty=repetition_penalty,
                    no_repeat_ngram_size=no_repeat_ngram_size,
                    stream=False,
                    use_cache=True,
                    on_text=on_text,
                )
        except Exception as exc:  # noqa: BLE001 -- surface generation errors in the UI
            q.put(f"\n\n[generation error: {type(exc).__name__}: {exc}]")
        finally:
            q.put(done)

    threading.Thread(target=runner, daemon=True).start()

    acc = ""
    while True:
        item = q.get()
        if item is done:
            break
        acc += item
        yield acc + _stats()

    if not acc:
        yield "(no output)"


# ── UI ────────────────────────────────────────────────────────────────────────

CSS = """
/* ══════════════════════════════════════════════════════
   Sarvam-style: warm white · near-black · saffron accent
   Inter (body) + IBM Plex Mono (code / telemetry)
   ══════════════════════════════════════════════════════ */

:root {
    /* surface */
    --s-bg:        #fafaf9;
    --s-white:     #ffffff;
    --s-warm-50:   #f5f5f4;
    --s-warm-100:  #e7e5e4;
    --s-warm-200:  #d6d3d1;

    /* text */
    --s-ink:       #0a0a0a;
    --s-ink-2:     #292524;
    --s-muted:     #78716c;
    --s-faint:     #a8a29e;

    /* accent — saffron */
    --s-saf:       #ea580c;
    --s-saf-hover: #c2410c;
    --s-saf-light: #fff7ed;
    --s-saf-mid:   #fed7aa;

    /* Gradio token overrides */
    --body-background-fill:          var(--s-bg);
    --background-fill-primary:       var(--s-white);
    --background-fill-secondary:     var(--s-warm-50);
    --block-background-fill:         var(--s-white);
    --block-border-color:            var(--s-warm-100);
    --block-title-text-color:        var(--s-muted);
    --block-label-text-color:        var(--s-muted);
    --body-text-color:               var(--s-ink);
    --body-text-color-subdued:       var(--s-muted);
    --border-color-primary:          var(--s-warm-100);
    --border-color-accent:           var(--s-saf);
    --color-accent:                  var(--s-saf);
    --input-background-fill:         var(--s-white);
    --input-border-color:            var(--s-warm-200);
    --input-border-color-focus:      var(--s-saf);
    --input-placeholder-color:       var(--s-faint);
    --slider-color:                  var(--s-saf);
    --checkbox-background-color-selected:   var(--s-saf);
    --checkbox-border-color-selected:       var(--s-saf);
    --checkbox-label-background-fill-selected:
        color-mix(in srgb, var(--s-saf) 8%, var(--s-white));
    --button-primary-background-fill:       var(--s-saf);
    --button-primary-background-fill-hover: var(--s-saf-hover);
    --button-primary-text-color:            #ffffff;
    --button-secondary-background-fill:     var(--s-white);
    --button-secondary-border-color:        var(--s-warm-200);
    --button-secondary-background-fill-hover: var(--s-warm-50);
    --link-text-color:                      var(--s-saf);
}

/* ── Base ──────────────────────────────────────────────── */
* { box-sizing: border-box; }

body, .gradio-container {
    font-family: 'Inter', system-ui, -apple-system, sans-serif !important;
    background: var(--s-bg) !important;
    color: var(--s-ink) !important;
}

.gradio-container {
    max-width: 820px !important;
    margin: 0 auto !important;
    padding: 0 20px !important;
}

/* ── Nav bar ──────────────────────────────────────────── */
#s-nav {
    display: flex;
    align-items: center;
    padding: 18px 0 16px;
    border-bottom: 1px solid var(--s-warm-100);
    margin-bottom: 0;
    gap: 12px;
}
#s-wordmark {
    font-family: 'Inter', sans-serif;
    font-weight: 700;
    font-size: 1.05rem;
    letter-spacing: -0.02em;
    color: var(--s-ink);
    display: flex;
    align-items: center;
    gap: 8px;
    text-decoration: none;
}
#s-wordmark-dot {
    width: 8px; height: 8px; border-radius: 50%;
    background: var(--s-saf);
    flex-shrink: 0;
}
#s-nav-tag {
    font-size: 0.72rem;
    font-weight: 500;
    color: var(--s-saf);
    background: var(--s-saf-light);
    border: 1px solid var(--s-saf-mid);
    padding: 3px 9px;
    border-radius: 999px;
    letter-spacing: 0.01em;
}
#s-nav-right {
    margin-left: auto;
    display: flex;
    align-items: center;
    gap: 20px;
}
.s-nav-link {
    font-size: 0.82rem;
    color: var(--s-muted);
    text-decoration: none;
    font-weight: 500;
    transition: color 0.12s;
}
.s-nav-link:hover { color: var(--s-ink); }

/* ── Hero ─────────────────────────────────────────────── */
#s-hero {
    padding: 56px 0 40px;
    border-bottom: 1px solid var(--s-warm-100);
}
#s-hero-eyebrow {
    font-size: 0.78rem;
    font-weight: 600;
    letter-spacing: 0.06em;
    text-transform: uppercase;
    color: var(--s-saf);
    margin-bottom: 16px;
}
#s-hero h1 {
    font-size: clamp(2rem, 5vw, 2.8rem);
    font-weight: 700;
    letter-spacing: -0.03em;
    line-height: 1.12;
    color: var(--s-ink);
    margin: 0 0 16px;
    max-width: 560px;
}
#s-hero-desc {
    font-size: 1rem;
    color: var(--s-muted);
    line-height: 1.65;
    max-width: 480px;
    margin-bottom: 28px;
    font-weight: 400;
}
#s-stat-row {
    display: flex;
    gap: 32px;
    flex-wrap: wrap;
}
.s-stat {
    display: flex;
    flex-direction: column;
    gap: 2px;
}
.s-stat-val {
    font-size: 1.1rem;
    font-weight: 700;
    letter-spacing: -0.02em;
    color: var(--s-ink);
    font-variant-numeric: tabular-nums;
}
.s-stat-lbl {
    font-size: 0.72rem;
    color: var(--s-faint);
    font-weight: 500;
    text-transform: uppercase;
    letter-spacing: 0.05em;
}

/* ── Section label ────────────────────────────────────── */
.s-section-label {
    font-size: 0.72rem;
    font-weight: 600;
    letter-spacing: 0.08em;
    text-transform: uppercase;
    color: var(--s-faint);
    padding: 28px 0 12px;
    border-bottom: none;
}

/* ── Chat bubbles ─────────────────────────────────────── */
/* User: saffron-tinted */
.user-row {
    background: var(--s-saf) !important;
    border-radius: 12px 12px 2px 12px !important;
}
.user-row * { color: #fff !important; }

/* Bot: white card with warm border */
.bot-row {
    background: var(--s-white) !important;
    border: 1px solid var(--s-warm-100) !important;
    border-radius: 12px 12px 12px 2px !important;
}

.message-bubble-border {
    box-shadow: 0 1px 3px rgba(0,0,0,0.06) !important;
}

/* Code in bot replies */
.bot-row code, .bot-row pre {
    font-family: 'IBM Plex Mono', monospace !important;
    background: var(--s-warm-50) !important;
    border: 1px solid var(--s-warm-100) !important;
    border-radius: 6px !important;
    font-size: 0.82rem !important;
}

/* ── Textbox & send button ────────────────────────────── */
textarea {
    font-family: 'Inter', sans-serif !important;
    font-size: 0.95rem !important;
    border-radius: 8px !important;
    border: 1.5px solid var(--s-warm-200) !important;
    transition: border-color 0.15s, box-shadow 0.15s !important;
    resize: none !important;
    background: var(--s-white) !important;
    color: var(--s-ink) !important;
}
textarea:focus {
    border-color: var(--s-saf) !important;
    box-shadow: 0 0 0 3px color-mix(in srgb, var(--s-saf) 12%, transparent) !important;
    outline: none !important;
}
textarea::placeholder { color: var(--s-faint) !important; }

button.primary {
    background: #ea580c !important;
    background: var(--s-saf, #ea580c) !important;
    border: none !important;
    border-radius: 8px !important;
    font-weight: 600 !important;
    font-size: 0.88rem !important;
    letter-spacing: 0.01em !important;
    color: #ffffff !important;
    transition: background 0.12s, transform 0.1s !important;
}
button.primary:hover {
    background: #c2410c !important;
    background: var(--s-saf-hover, #c2410c) !important;
    transform: translateY(-1px) !important;
}

/* ── Settings accordion ───────────────────────────────── */
.accordion {
    border: 1px solid var(--s-warm-100) !important;
    border-radius: 10px !important;
    background: var(--s-white) !important;
    overflow: hidden;
}
.accordion > .label-wrap {
    padding: 12px 16px !important;
    font-size: 0.85rem !important;
    font-weight: 600 !important;
    color: var(--s-ink-2) !important;
    background: var(--s-warm-50) !important;
    border-bottom: 1px solid var(--s-warm-100) !important;
}

/* Sliders */
input[type="range"] { accent-color: var(--s-saf) !important; }

/* ── Example chips ────────────────────────────────────── */
.example {
    font-size: 0.82rem !important;
    font-weight: 500 !important;
    border: 1px solid var(--s-warm-200) !important;
    border-radius: 6px !important;
    color: var(--s-ink-2) !important;
    background: var(--s-white) !important;
    padding: 6px 12px !important;
    transition: border-color 0.12s, color 0.12s, background 0.12s !important;
}
.example:hover {
    border-color: var(--s-saf) !important;
    color: var(--s-saf) !important;
    background: var(--s-saf-light) !important;
}

/* ── Footer ───────────────────────────────────────────── */
#s-footer {
    border-top: 1px solid var(--s-warm-100);
    padding: 28px 0 20px;
    display: flex;
    justify-content: space-between;
    align-items: flex-start;
    flex-wrap: wrap;
    gap: 16px;
    margin-top: 16px;
}
#s-footer-left {
    display: flex;
    flex-direction: column;
    gap: 5px;
}
#s-footer-brand {
    font-weight: 700;
    font-size: 0.88rem;
    color: var(--s-ink);
    letter-spacing: -0.01em;
}
#s-footer-copy {
    font-size: 0.72rem;
    color: var(--s-faint);
}
#s-footer-chips {
    display: flex;
    gap: 6px;
    flex-wrap: wrap;
    align-items: center;
}
.s-badge {
    font-size: 0.68rem;
    font-weight: 500;
    font-family: 'IBM Plex Mono', monospace;
    color: var(--s-muted);
    background: var(--s-warm-50);
    border: 1px solid var(--s-warm-100);
    padding: 3px 8px;
    border-radius: 4px;
    letter-spacing: 0.02em;
    white-space: nowrap;
}
.s-badge-accent {
    color: var(--s-saf);
    background: var(--s-saf-light);
    border-color: var(--s-saf-mid);
}

/* ── Misc fixes ───────────────────────────────────────── */
* { scrollbar-width: thin; scrollbar-color: var(--s-warm-200) transparent; }
::-webkit-scrollbar { width: 5px; }
::-webkit-scrollbar-thumb { background: var(--s-warm-200); border-radius: 4px; }
"""


def main():
    # Build runtime stats for the hero section
    device_label  = str(_device).upper()
    dtype_label   = str(_cast_dtype).replace("torch.", "").upper() if _cast_dtype else "FP32"
    layer_str     = str(_n_layer)
    head_str      = f"{_n_head}/{_kv_heads}"

    header_html = f"""
    <div id="s-nav">
        <span id="s-wordmark">
            <span id="s-wordmark-dot"></span>
            Pragya
        </span>
        <span id="s-nav-tag">by Arush Kumar</span>
        <div id="s-nav-right">
            <span class="s-nav-link">PyTorch</span>
            <span class="s-nav-link">{device_label}</span>
        </div>
    </div>

    <div id="s-hero">
        <div id="s-hero-eyebrow">Arush Kumar · From-scratch language model</div>
        <h1>Pragya — a model<br>built from the ground up.</h1>
        <p id="s-hero-desc">
            A custom GPT-style architecture with Mixture-of-Depths routing,
            Multi-Token Prediction heads, and YaRN RoPE context scaling.
            Every weight trained from scratch in PyTorch. Type a prompt and let it continue.
        </p>
        <div id="s-stat-row">
            <div class="s-stat">
                <span class="s-stat-val">{_n_params_str}</span>
                <span class="s-stat-lbl">Parameters</span>
            </div>
            <div class="s-stat">
                <span class="s-stat-val">{_ctx_len:,}</span>
                <span class="s-stat-lbl">Context length</span>
            </div>
            <div class="s-stat">
                <span class="s-stat-val">{layer_str}L · {head_str}H</span>
                <span class="s-stat-lbl">Depth · GQA heads</span>
            </div>
            <div class="s-stat">
                <span class="s-stat-val">{_pattern}</span>
                <span class="s-stat-lbl">Attention pattern</span>
            </div>
            <div class="s-stat">
                <span class="s-stat-val">{dtype_label}</span>
                <span class="s-stat-lbl">Inference dtype</span>
            </div>
        </div>
    </div>
    """

    footer_html = f"""
    <div id="s-footer">
        <div id="s-footer-left">
            <span id="s-footer-brand">Pragya &nbsp;·&nbsp; <span style="font-weight:400;color:var(--s-muted)">Arush Kumar</span></span>
            <span id="s-footer-copy">
                Continuation model · conversation history is shown in UI but not fed back to the model ·
                TTFT and tok/s are live-measured each generation
            </span>
        </div>
        <div id="s-footer-chips">
            <span class="s-badge s-badge-accent">{_n_params_str} params</span>
            <span class="s-badge">{device_label}</span>
            <span class="s-badge">{dtype_label}</span>
            <span class="s-badge">ctx {_ctx_len}</span>
            <span class="s-badge">{_pattern}</span>
        </div>
    </div>
    """

    with gr.Blocks(title="Pragya by Arush Kumar") as demo:
        gr.HTML("""
        <link rel="preconnect" href="https://fonts.googleapis.com">
        <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
        <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=IBM+Plex+Mono:wght@400;500&display=swap" rel="stylesheet">
        <style>
          :root {
            --s-bg:#fafaf9; --s-white:#ffffff; --s-warm-50:#f5f5f4; --s-warm-100:#e7e5e4; --s-warm-200:#d6d3d1;
            --s-ink:#0a0a0a; --s-ink-2:#292524; --s-muted:#78716c; --s-faint:#a8a29e;
            --s-saf:#ea580c; --s-saf-hover:#c2410c; --s-saf-light:#fff7ed; --s-saf-mid:#fed7aa;
          }
        </style>
        """)
        gr.HTML(header_html)

        gr.HTML('<div class="s-section-label">Try it</div>')

        chatbot = gr.Chatbot(height=440, show_label=False, render_markdown=True)

        with gr.Row():
            textbox = gr.Textbox(
                placeholder="Start a sentence and let Pragya complete it …",
                show_label=False,
                lines=2,
                scale=8,
            )
            submit = gr.Button("Generate", variant="primary", scale=1)

        with gr.Accordion("Generation settings", open=False):
            max_new_tokens = gr.Slider(16, _MAX_TOKENS_CAP, value=_DEFAULT_TOKENS, step=8,
                                       label="Max new tokens")
            temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature")
            top_k = gr.Slider(0, 200, value=50, step=5, label="Top-k",
                              info="0 disables top-k filtering.")
            top_p = gr.Slider(0.5, 1.0, value=1.0, step=0.01, label="Top-p (nucleus)",
                              info="1.0 disables nucleus sampling.")
            rep_pen = gr.Slider(1.0, 2.0, value=1.3, step=0.05, label="Repetition penalty",
                                info="Count-scaled: penalty^occurrences. 1.0 = off.")
            ngram = gr.Slider(0, 6, value=3, step=1, label="No-repeat n-gram",
                              info="Hard-bans any n-gram of this length from repeating. 0 = off.")
            use_anchor = gr.Checkbox(value=True, label="Narrative anchor",
                                     info="Prepends a short coherent prose passage to prime the context.")

        gr.Examples(
            examples=[["Hello"], ["Tell me a story."], ["Create a simple story."],
                      ["Write a bedtime story."]],
            inputs=[textbox],
            cache_examples=False,
        )

        settings = [max_new_tokens, temperature, top_k, top_p, rep_pen, ngram, use_anchor]

        def _chat(message, history, *cfg):
            message = (message or "").strip()
            if not message:
                yield history, ""
                return

            history = list(history) + [[message, ""]]
            yield history, ""
            try:
                for partial in respond(message, history, *cfg):
                    history[-1][1] = partial
                    yield history, ""
            except Exception as exc:  # noqa: BLE001
                history[-1][1] = f"[generation error: {type(exc).__name__}: {exc}]"
                yield history, ""

        submit.click(_chat, [textbox, chatbot, *settings], [chatbot, textbox])
        textbox.submit(_chat, [textbox, chatbot, *settings], [chatbot, textbox])

        gr.HTML(footer_html)

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        css=CSS,
        theme=gr.themes.Base(
            primary_hue=gr.themes.colors.orange,
            neutral_hue=gr.themes.colors.stone,
            font=gr.themes.GoogleFont("Inter"),
        ),
    )


if __name__ == "__main__":
    main()