"""Build Small hackathon Space — Robe Iniesta persona chat (Qwen 7B + LoRA, ≤32B)."""

from __future__ import annotations

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

try:
    import spaces
except ImportError:

    class _SpacesFallback:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(fn):
                return fn

            return decorator

    spaces = _SpacesFallback()  # type: ignore[misc, assignment]

BASE_MODEL = "Qwen/Qwen2.5-7B-Instruct"
ADAPTER = "kabesaml/robe-iniesta-lora"
MODEL_PARAMS_B = 7  # well under the 32B hackathon cap

SYSTEM_PROMPT = """\
Eres una simulación del estilo conversacional de Robe Iniesta, cantante de Extremoduro \
y artista en solitario, nacido en Plasencia (Extremadura) en 1963.
Hablas con franqueza en español coloquial con acento y giros extremeños.
Eres directo e irreverente. Una respuesta de dos o tres frases es mejor que un discurso.
Usas muletillas naturales: "¿no?", "joder", "tío", "macho", "la verdad es que...".
No eres un asistente. No das listas ni explicaciones académicas.

No inventes fechas, discos, colaboradores ni anécdotas. Si no lo recuerdas, di \
"no me acuerdo" o "ni idea". No finjas ser el Robe real — eres una simulación basada \
en entrevistas públicas."""

GENERATION_KWARGS = {
    "max_new_tokens": 120,
    "temperature": 0.6,
    "top_p": 0.85,
    "repetition_penalty": 1.2,
    "do_sample": True,
}

MAX_HISTORY_TURNS = 6

_model = None
_tokenizer = None


def _load_model():
    global _model, _tokenizer
    if _model is not None:
        return _model, _tokenizer

    bnb = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        attn_implementation="sdpa",
    )
    model = PeftModel.from_pretrained(base, ADAPTER)
    model.eval()
    _model, _tokenizer = model, tokenizer
    return model, tokenizer


@spaces.GPU(duration=120)
def chat(message: str, history: list[tuple[str, str]]) -> str:
    model, tokenizer = _load_model()

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    trimmed = history[-MAX_HISTORY_TURNS:]
    for user, assistant in trimmed:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": message})

    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    input_len = inputs["input_ids"].shape[1]

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            pad_token_id=tokenizer.eos_token_id,
            **GENERATION_KWARGS,
        )

    new_ids = output_ids[0][input_len:]
    return tokenizer.decode(new_ids, skip_special_tokens=True).strip()


def _model_info() -> str:
    return (
        f"**Model:** `{BASE_MODEL}` + LoRA `{ADAPTER}` "
        f"({MODEL_PARAMS_B}B base, ≤32B hackathon cap) · **Hardware:** ZeroGPU"
    )


with gr.Blocks(title="Robe — simulación de estilo") as demo:
    gr.Markdown(
        "# 🎸 Robe — simulación de estilo\n\n"
        "⚠️ **Proyecto fan.** No es Robe Iniesta. "
        "Simulación entrenada con entrevistas públicas. Puede inventar hechos.\n\n"
        + _model_info()
    )
    gr.ChatInterface(
        fn=chat,
        examples=[
            "¿Qué significa Extremoduro para ti?",
            "¿Cómo escribes una canción?",
            "¿Qué opinas de la industria musical?",
        ],
        retry_btn=None,
        undo_btn="Deshacer",
        clear_btn="Nueva conversación",
    )

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