import os
import sys

import torch
import gradio as gr
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download, snapshot_download

# ── Repo IDs ───────────────────────────────────────────────────────────────────
REPO_V1 = "IvmeLabs/Ivme-Conversate-v1-Base"
REPO_V2 = "IvmeLabs/Ivme-Conversate-v2-Base"

device = "cuda" if torch.cuda.is_available() else "cpu"


# ── Load v1 ────────────────────────────────────────────────────────────────────
def load_v1():
    tokenizer_path = hf_hub_download(repo_id=REPO_V1, filename="ivme_tokenizer.json")
    model_path = hf_hub_download(repo_id=REPO_V1, filename="ivme_base_ema.pt")
    model_py_path = hf_hub_download(repo_id=REPO_V1, filename="model.py")

    model_dir = os.path.dirname(model_py_path)
    if model_dir not in sys.path:
        sys.path.insert(0, model_dir)
    from model import IvmeConversate  # noqa: E402

    tok = Tokenizer.from_file(tokenizer_path)
    ckpt = torch.load(model_path, map_location=device, weights_only=False)
    cfg = ckpt["cfg"]
    cfg.attn_backend = "sdpa"

    model = IvmeConversate(cfg).to(device)
    model.load_state_dict(ckpt["model"])
    model.eval()

    max_ctx = (
        getattr(cfg, "block_size", None)
        or getattr(cfg, "n_ctx", None)
        or getattr(cfg, "max_seq_len", None)
        or getattr(cfg, "context_length", None)
        or 1024
    )
    eos_id = tok.token_to_id("<|eos|>")
    return {"tokenizer": tok, "model": model, "max_ctx": max_ctx, "eos_id": eos_id}


# ── Load v2 ────────────────────────────────────────────────────────────────────
def load_v2():
    # v2's architecture code lives under a `model/` package in the repo, which
    # collides by name with v1's already-imported top-level `model` module, so
    # we must remove any cached `model` module before (re)importing v2's package.
    for mod_name in list(sys.modules):
        if mod_name == "model" or mod_name.startswith("model."):
            del sys.modules[mod_name]

    repo_local_dir = snapshot_download(REPO_V2, allow_patterns=["model/*"])
    if repo_local_dir not in sys.path:
        sys.path.insert(0, repo_local_dir)
    from model import IvmeConfig, IvmeConversateV2  # noqa: E402

    tokenizer_path = hf_hub_download(repo_id=REPO_V2, filename="tokenizer.json")
    ckpt_path = hf_hub_download(repo_id=REPO_V2, filename="ckpt_final.pt")

    tok = Tokenizer.from_file(tokenizer_path)

    torch.serialization.add_safe_globals([IvmeConfig])
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    cfg = ckpt["config"]

    model = IvmeConversateV2(cfg)
    state_dict = ckpt["ema_state_dict"]
    state_dict = {k.removeprefix("_orig_mod."): v for k, v in state_dict.items()}
    model.load_state_dict(state_dict)
    model.to(device).eval()

    max_ctx = getattr(cfg, "context_len", 1024)
    eos_id = tok.token_to_id("<|endoftext|>")
    return {"tokenizer": tok, "model": model, "max_ctx": max_ctx, "eos_id": eos_id}


print("Loading İvme-Conversate-v1-Base...")
V1 = load_v1()
print("Loading İvme-Conversate-v2-Base...")
V2 = load_v2()

REGISTRY = {
    "İvme-Conversate-v2-Base (recommended)": V2,
    "İvme-Conversate-v1-Base": V1,
}

BENCH = {
    # name: (v1, v2, higher_is_better)
    "WikiText-2 byte perplexity": (2.96, 2.2250, False),
    "BLiMP (macro-avg)": (61.40, 75.09, True),
    "ARC-Easy (acc_norm)": (30.85, 39.98, True),
}


# ── Shared generation core ────────────────────────────────────────────────────
@torch.no_grad()
def _generate(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty):
    tokenizer = bundle["tokenizer"]
    model = bundle["model"]
    max_ctx = bundle["max_ctx"]
    eos_id = bundle["eos_id"]

    prompt = prompt or ""
    input_ids = tokenizer.encode(prompt).ids
    if not input_ids:
        yield prompt
        return

    generated = torch.tensor([input_ids], device=device, dtype=torch.long)
    vocab_size = None
    response_tokens: list[int] = []
    temperature = max(float(temperature), 1e-6)

    for _ in range(int(max_new_tokens)):
        window = generated[:, -max_ctx:]
        out = model(window)

        if isinstance(out, (tuple, list)):
            out = out[0]
        elif isinstance(out, dict):
            out = out.get("logits", next(iter(out.values())))

        logits = out[:, -1, :].float()

        if vocab_size is None:
            vocab_size = logits.size(-1)

        if repetition_penalty and repetition_penalty != 1.0:
            seen = torch.unique(generated[0])
            scores = logits[0, seen]
            scores = torch.where(
                scores > 0, scores / repetition_penalty, scores * repetition_penalty
            )
            logits[0, seen] = scores

        logits = logits / temperature

        k = int(top_k)
        if k > 0:
            k = min(k, vocab_size)
            topk_vals, _ = torch.topk(logits, k)
            logits[logits < topk_vals[:, -1:]] = float("-inf")

        probs = torch.softmax(logits, dim=-1)

        if not torch.isfinite(probs).all() or probs.sum() <= 0:
            next_tok = torch.argmax(logits, dim=-1, keepdim=True)
        else:
            next_tok = torch.multinomial(probs, num_samples=1)

        tok_id = next_tok.item()
        if eos_id is not None and tok_id == eos_id:
            break

        response_tokens.append(tok_id)
        generated = torch.cat([generated, next_tok], dim=1)

        yield prompt + tokenizer.decode(response_tokens)

    if not response_tokens:
        yield prompt


def continue_text(model_choice, prompt, max_new_tokens, temperature, top_k, repetition_penalty):
    bundle = REGISTRY[model_choice]
    yield from _generate(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)


def compare_generate(prompt, max_new_tokens, temperature, top_k, repetition_penalty):
    """Run v1 and v2 on the same prompt/settings, streaming both in parallel steps."""
    gen_v1 = _generate(V1, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
    gen_v2 = _generate(V2, prompt, max_new_tokens, temperature, top_k, repetition_penalty)

    last_v1, last_v2 = prompt, prompt
    done_v1 = done_v2 = False

    while not (done_v1 and done_v2):
        if not done_v1:
            try:
                last_v1 = next(gen_v1)
            except StopIteration:
                done_v1 = True
        if not done_v2:
            try:
                last_v2 = next(gen_v2)
            except StopIteration:
                done_v2 = True
        yield last_v1, last_v2


def benchmark_table():
    rows = []
    for name, (v1, v2, higher_better) in BENCH.items():
        delta = (v2 - v1) if higher_better else (v1 - v2)
        pct = (delta / abs(v1)) * 100 if v1 else 0
        arrow = "↑" if higher_better else "↓"
        rows.append([name + f" {arrow}", f"{v1:.2f}", f"{v2:.2f}", f"{'+' if delta >= 0 else ''}{delta:.2f} ({pct:+.0f}%)"])
    return rows


# ── UI ─────────────────────────────────────────────────────────────────────────
CSS = """
body, .gradio-container { font-family: 'Inter', system-ui, sans-serif; }
#component-0 { max-width: 900px; margin: 0 auto; padding: 16px; }
footer { display: none !important; }
.ivme-output textarea { font-size: 1.02rem; line-height: 1.6; }
"""

EXAMPLES = [
    "The theory of relativity states that",
    "In the beginning, the universe was",
    "def fibonacci(n):",
    "The most important thing to remember about cooking is",
    "Once upon a time, in a small village by the sea,",
    "Python is a programming language that",
]

with gr.Blocks(css=CSS, title="İvme-Conversate") as demo:
    gr.Markdown(
        "## İvme-Conversate — Tiny Language Models, Text Continuation\n"
        "Sub-25M-parameter decoder-only **base** models · not instruction-tuned · "
        "these models **continue** text, they do not chat or answer questions."
    )

    with gr.Tabs():
        # ── Tab 1: single-model playground with picker ──────────────────────
        with gr.Tab("Playground"):
            model_picker = gr.Dropdown(
                choices=list(REGISTRY.keys()),
                value="İvme-Conversate-v2-Base (recommended)",
                label="Model",
            )

            prompt_box = gr.Textbox(
                label="Prompt",
                placeholder="The theory of relativity states that…",
                lines=3,
                value="The theory of relativity states that",
            )

            with gr.Row():
                gen_btn = gr.Button("Generate", variant="primary", scale=3)
                clear_btn = gr.Button("Clear", scale=1)

            output_box = gr.Textbox(
                label="Continuation",
                lines=12,
                show_copy_button=True,
                elem_classes="ivme-output",
                interactive=False,
            )

            gr.Examples(examples=[[e] for e in EXAMPLES], inputs=prompt_box, label="Try a prompt")

            with gr.Accordion("Settings", open=False):
                with gr.Row():
                    max_tokens = gr.Slider(16, 512, value=200, step=8, label="Max new tokens")
                    temperature = gr.Slider(0.1, 2.0, value=0.7, step=0.05, label="Temperature")
                with gr.Row():
                    top_k = gr.Slider(0, 200, value=40, step=1, label="Top-k (0 = disabled)")
                    rep_penalty = gr.Slider(1.0, 2.0, value=1.15, step=0.05, label="Repetition penalty")

            gen_inputs = [model_picker, prompt_box, max_tokens, temperature, top_k, rep_penalty]
            gen_btn.click(continue_text, gen_inputs, output_box)
            prompt_box.submit(continue_text, gen_inputs, output_box)
            clear_btn.click(lambda: ("", ""), None, [prompt_box, output_box], queue=False)

        # ── Tab 2: side-by-side compare ──────────────────────────────────────
        with gr.Tab("Compare v1 vs v2"):
            gr.Markdown(
                "Run the **same prompt and settings** through both models at once "
                "to see the difference training data made, plus the benchmark deltas below."
            )

            cmp_prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                value="Once upon a time, there was a",
            )

            with gr.Row():
                cmp_gen_btn = gr.Button("Generate both", variant="primary", scale=3)
                cmp_clear_btn = gr.Button("Clear", scale=1)

            with gr.Row():
                cmp_out_v1 = gr.Textbox(
                    label="v1-Base",
                    lines=10,
                    show_copy_button=True,
                    elem_classes="ivme-output",
                    interactive=False,
                )
                cmp_out_v2 = gr.Textbox(
                    label="v2-Base",
                    lines=10,
                    show_copy_button=True,
                    elem_classes="ivme-output",
                    interactive=False,
                )

            gr.Examples(examples=[[e] for e in EXAMPLES], inputs=cmp_prompt, label="Try a prompt")

            with gr.Accordion("Settings", open=False):
                with gr.Row():
                    cmp_max_tokens = gr.Slider(16, 512, value=150, step=8, label="Max new tokens")
                    cmp_temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature")
                with gr.Row():
                    cmp_top_k = gr.Slider(0, 200, value=50, step=1, label="Top-k (0 = disabled)")
                    cmp_rep_penalty = gr.Slider(1.0, 2.0, value=1.0, step=0.05, label="Repetition penalty")

            gr.Markdown("### Benchmark improvement, v1 → v2")
            gr.Dataframe(
                headers=["Benchmark", "v1", "v2", "Δ (v1 → v2)"],
                value=benchmark_table(),
                interactive=False,
                row_count=(len(BENCH), "fixed"),
            )

            cmp_inputs = [cmp_prompt, cmp_max_tokens, cmp_temperature, cmp_top_k, cmp_rep_penalty]
            cmp_gen_btn.click(compare_generate, cmp_inputs, [cmp_out_v1, cmp_out_v2])
            cmp_prompt.submit(compare_generate, cmp_inputs, [cmp_out_v1, cmp_out_v2])
            cmp_clear_btn.click(lambda: ("", "", ""), None, [cmp_prompt, cmp_out_v1, cmp_out_v2], queue=False)

demo.queue().launch()