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

# ─── Modell ───────────────────────────────────────────────────────────────
MODEL_ID = "Qwen/Qwen3-Coder-30B-A3B-Instruct"
MAX_TOKENS = 32768
DEFAULT_TEMP = 0.6
DEFAULT_SYSTEM = "You are Qwen Coder Boosted, an expert coding assistant. Generate clean, working code. Always show complete files with [FILE:filename] markers."

print("🚀 Lade Modell...")
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)
print("✅ Modell geladen!")


# ─── GPU-Funktion (nur hier ist echte GPU) ──────────────────────────────
@spaces.GPU(duration=120)
def generate_response(messages, temperature, max_tokens):
    """Generate a response from the model. GPU is real inside this function."""
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer([text], return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.eos_token_id,
        )

    response = tokenizer.decode(
        outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True
    )
    return response


# ─── Chat-Logik ──────────────────────────────────────────────────────────
def chat(message, history, system_prompt, temperature):
    """Handle a chat turn. Builds messages, calls GPU function, returns updated history."""
    if not message.strip():
        return history

    # Build messages array from history
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    # Generate (this triggers @spaces.GPU)
    response = generate_response(messages, temperature, MAX_TOKENS)

    # Update history
    history.append((message, response))
    return history


# ─── UI ──────────────────────────────────────────────────────────────────
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo"),
    title="Qwen Coder Boosted",
) as demo:
    gr.Markdown("""
    # 🚀 Qwen Coder Boosted
    **Qwen3-Coder-30B-A3B (4-bit) auf Zero GPU** – 48 GB shared VRAM, PRO-Quota
    """)

    with gr.Row():
        with gr.Column(scale=1):
            system_prompt = gr.Textbox(
                label="⚙️ System Prompt",
                value=DEFAULT_SYSTEM,
                lines=3,
            )
            temperature = gr.Slider(
                label="🌡️ Temperature",
                minimum=0.1,
                maximum=1.5,
                value=DEFAULT_TEMP,
                step=0.05,
            )
            clear_btn = gr.Button("🗑️ Chat leeren", variant="secondary")

        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="💬 Chat",
                height=550,
                bubble_full_width=False,
                show_copy_button=True,
            )

            with gr.Row():
                msg = gr.Textbox(
                    label="Deine Nachricht",
                    placeholder="Schreib einen Prompt... z.B. 'Baue eine To-Do-App mit HTML/CSS/JS'",
                    scale=8,
                    container=True,
                )
                submit = gr.Button("📤 Senden", variant="primary", scale=1)

    gr.Markdown("""
    ---
    ### 💡 Info
    - **Modell:** Qwen3-Coder-30B-A3B (4-bit, ~15 GB VRAM)
    - **Zero GPU:** 48 GB shared, PRO = 40 Min/Tag inklusive
    - **Max Tokens:** 32K (hardcodiert)
    - **Temperature:** 0.6 (Standard für Coding)
    """)

    # Events
    submit.click(
        fn=chat,
        inputs=[msg, chatbot, system_prompt, temperature],
        outputs=[chatbot],
    ).then(
        fn=lambda: "",
        outputs=[msg],
    )

    msg.submit(
        fn=chat,
        inputs=[msg, chatbot, system_prompt, temperature],
        outputs=[chatbot],
    ).then(
        fn=lambda: "",
        outputs=[msg],
    )

    clear_btn.click(fn=lambda: None, outputs=[chatbot])


if __name__ == "__main__":
    demo.launch()