import os

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

MODEL_ID = "play451/Moron"
BASE_MODEL_ID = "unsloth/qwen2.5-coder-7b-instruct-bnb-4bit"

HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# IMPORTANT (ZeroGPU): do NOT load the model or touch CUDA at module level.
# On ZeroGPU Spaces there is no GPU attached to the process until you're
# actually inside a function decorated with @spaces.GPU. The base model here
# is also a pre-quantized bitsandbytes 4-bit checkpoint, which requires a
# live CUDA context just to load — so model construction has to happen
# lazily, on first call, inside the GPU-decorated function.
model = None


def load_model():
    global model
    if model is not None:
        return model

    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        token=HF_TOKEN,
        device_map="auto",
        dtype=torch.bfloat16,
    )

    loaded = PeftModel.from_pretrained(
        base_model,
        MODEL_ID,
        token=HF_TOKEN,
    )
    loaded.eval()
    model = loaded
    return model


def format_prompt(prompt):
    return tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}],
        tokenize=False,
        add_generation_prompt=True,
    )


@spaces.GPU(duration=60)
def generate(prompt, max_tokens, temperature, top_p):
    prompt = (prompt or "").strip()
    if not prompt:
        return ""

    try:
        current_model = load_model()
        text = format_prompt(prompt)
        inputs = tokenizer(text, return_tensors="pt")
        model_device = next(current_model.parameters()).device
        inputs = {k: v.to(model_device) for k, v in inputs.items()}

        with torch.inference_mode():
            output = current_model.generate(
                **inputs,
                max_new_tokens=int(max_tokens),
                temperature=float(temperature),
                top_p=float(top_p),
                do_sample=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )

        new_tokens = output[0][inputs["input_ids"].shape[-1]:]
        return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

    except Exception as e:
        return f"{type(e).__name__}: {e}"


examples = [
    ["Explain how a Minecraft Fabric mod is structured."],
    ["Write a simple Java example for a Minecraft Fabric client module."],
    ["Help me debug a Java compilation error."],
]

with gr.Blocks() as demo:
    gr.Markdown("# 🤖 play451/Moron\n\nA Gradio demo for **play451/Moron**.")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Ask Moron something...",
                lines=8,
            )
            with gr.Row():
                generate_btn = gr.Button("Generate", variant="primary")
                clear_btn = gr.Button("Clear")

        with gr.Column():
            output = gr.Textbox(label="Response", lines=14)

    with gr.Accordion("Generation settings", open=False):
        max_tokens = gr.Slider(16, 2048, value=512, step=1, label="Max new tokens")
        temperature = gr.Slider(0.1, 2.0, value=0.7, step=0.05, label="Temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")

    gr.Examples(examples=examples, inputs=prompt, label="Examples")

    generate_btn.click(
        generate,
        inputs=[prompt, max_tokens, temperature, top_p],
        outputs=output,
    )
    prompt.submit(
        generate,
        inputs=[prompt, max_tokens, temperature, top_p],
        outputs=output,
    )
    clear_btn.click(lambda: ("", ""), outputs=[prompt, output])

demo.launch(theme=gr.themes.Soft())