import spaces
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODELS = {
    "TinyBuddy-80K": "Eeppa/TinyBuddy-80K",
    "TinyBuddy-500K": "Eeppa/TinyBuddy-500K",
    "TinyBuddy-30M": "Eeppa/TinyBuddy-30M",
}

loaded = {}
load_errors = {}


def load_model(name):
    if name in loaded or name in load_errors:
        return
    repo = MODELS[name]
    try:
        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
        model.eval()
        loaded[name] = (model, tokenizer)
        print(f"[loaded] {name}")
    except Exception as e:
        load_errors[name] = f"{type(e).__name__}: {e}"
        print(f"[load failed] {name}: {e}")


for name in MODELS:
    load_model(name)


# ---- manual sampling fallback (works for ANY model, ignores custom generate) ----
def manual_sample(model, input_ids, max_new_tokens, temperature, top_k, min_new_tokens=0):
    """Generic top-k sampler. Bypasses any custom .generate() method."""
    generated = input_ids.clone()
    eos_id = getattr(model.config, "eos_token_id", None)

    for i in range(int(max_new_tokens)):
        with torch.no_grad():
            outputs = model(generated)
            logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
        next_logits = logits[0, -1, :] / max(temperature, 1e-5)

        # Suppress EOS for the first `min_new_tokens` steps
        if i < min_new_tokens and eos_id is not None:
            next_logits[eos_id] = float("-inf")

        if top_k > 0:
            k = min(int(top_k), next_logits.size(-1))
            top_vals, top_idx = torch.topk(next_logits, k)
            probs = torch.softmax(top_vals, dim=-1)
            chosen = torch.multinomial(probs, 1)
            next_token = top_idx[chosen]
        else:
            probs = torch.softmax(next_logits, dim=-1)
            next_token = torch.multinomial(probs, 1)

        generated = torch.cat([generated, next_token.view(1, 1)], dim=-1)

        # Only allow EOS to stop generation after min_new_tokens
        if i >= min_new_tokens and eos_id is not None:
            if next_token.item() == eos_id:
                break

    return generated


def encode_prompt(tokenizer, prompt):
    """Handle both fast and slow tokenizers, and flatten nested outputs."""
    if hasattr(tokenizer, "encode"):
        ids = tokenizer.encode(prompt)
        if hasattr(ids, "ids"):
            ids = ids.ids
    else:
        ids = tokenizer(prompt)["input_ids"]

    if torch.is_tensor(ids):
        ids = ids.tolist()
    if isinstance(ids, list) and len(ids) > 0 and isinstance(ids[0], list):
        ids = ids[0]

    return torch.tensor([ids], dtype=torch.long)


def decode_output(tokenizer, ids):
    """Normalize to a flat python list of ints before decoding."""
    if torch.is_tensor(ids):
        ids = ids.tolist()
    if isinstance(ids, list) and len(ids) > 0 and isinstance(ids[0], list):
        ids = ids[0]
    return tokenizer.decode(ids, skip_special_tokens=True)


@spaces.GPU(duration=30)
def run_generate(model_name, prompt, max_new_tokens, temperature, top_k):
    model, tokenizer = loaded[model_name]
    input_ids = encode_prompt(tokenizer, prompt)

    # Try the model's own generate first, fall back to manual sampling on failure
    try:
        with torch.no_grad():
            out = model.generate(
                input_ids,
                max_new_tokens=int(max_new_tokens),
                min_new_tokens=10,          # <-- force at least 10 new tokens
                temperature=float(temperature),
                top_k=int(top_k),
                do_sample=True,
            )
        return decode_output(tokenizer, out)
    except Exception as e:
        print(f"[fallback] {model_name} custom generate failed ({e}), using manual sampling")
        out = manual_sample(
            model, input_ids, max_new_tokens, temperature, top_k,
            min_new_tokens=10,              # <-- same here
        )
        return decode_output(tokenizer, out)


def compare_all(prompt, max_new_tokens, temperature, top_k):
    results = []
    for name in MODELS:
        if name in load_errors:
            results.append(f"[failed to load: {load_errors[name]}]")
            continue
        if name not in loaded:
            results.append("[not loaded]")
            continue
        try:
            text = run_generate(name, prompt, max_new_tokens, temperature, top_k)
        except Exception as e:
            text = f"[error: {type(e).__name__}: {e}]"
        results.append(text)
    return results


with gr.Blocks(title="TinyBuddy Trilogy") as demo:
    gr.Markdown(
        """
        # 🧸 TinyBuddy Trilogy
        Three tiny language models — **80K**, **500K**, and **30M** parameters.
        *Educational demos only. Not assistants.*
        """
    )

    with gr.Row():
        prompt = gr.Textbox(
            label="Prompt",
            value="Once upon a time, there was a little girl named Lisa.",
            lines=2, scale=4,
        )

    with gr.Row():
        max_new_tokens = gr.Slider(10, 200, value=60, step=10, label="Max new tokens")
        temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature")
        top_k = gr.Slider(1, 100, value=40, step=1, label="Top-k")

    with gr.Row():
        btn = gr.Button("Generate with all three", variant="primary")

    with gr.Row():
        out_80k = gr.Textbox(label="TinyBuddy-80K (83,856 params)", lines=8)
        out_500k = gr.Textbox(label="TinyBuddy-500K (547,296 params)", lines=8)
        out_30m = gr.Textbox(label="TinyBuddy-30M (30,371,840 params)", lines=8)

    btn.click(
        fn=compare_all,
        inputs=[prompt, max_new_tokens, temperature, top_k],
        outputs=[out_80k, out_500k, out_30m],
    )

demo.launch()