# ================================================================
# MTP-2.5 - app.py para Hugging Face Space (Gradio, CPU)
# Generacion 2: RoPE + SwiGLU (reemplaza position embeddings aprendidos y el
# GELU-MLP de MTP-1.x). Requiere un checkpoint entrenado con la Celda 1/2 de
# MTP-2.5; NO carga checkpoints de MTP-2.0 (RMSNorm cambia el state_dict) ni de MTP-1.x.
#
# OPTIMIZACIÓN DE VELOCIDAD (sin tocar el resto de la logica de muestreo):
#   - KV-cache en la atención: en generación autoregresiva, cada paso
#     antes recomputaba TODO el contexto desde cero (O(n^2) en total).
#     Ahora se reutiliza lo ya calculado y solo se procesa el token
#     nuevo (O(n) en total). Es el mismo cálculo matemático, solo que
#     no se repite trabajo ya hecho.
#   - F.scaled_dot_product_attention: kernel fusionado de PyTorch,
#     mismo resultado que el softmax manual pero más rápido en CPU.
#     Si la versión de PyTorch no lo trae, cae automáticamente al
#     cálculo manual (fallback), así que no se rompe en ningún entorno.
#   - repetition_penalty vectorizado + bloqueo de n-gramas repetidos
#     (evita que la respuesta final copie literalmente un fragmento ya
#     generado, sin que esto sea un "modelo de n-gramas": el modelo que
#     predice sigue siendo 100% transformer).
# ================================================================
import os
import math
import time
import logging
import threading
import traceback
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
import sentencepiece as spm
from starlette.middleware import Middleware
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
from huggingface_hub import hf_hub_download

# ---------------- Logging ----------------
# print() se pierde facil entre el ruido de arranque de Gradio/Starlette en
# los logs de un Space; con logging queda todo con timestamp y nivel, y es
# mas facil de filtrar si algo sale mal en produccion.
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("mtp")

# ---------------- Dispositivo ----------------
# Autodetecta GPU si el Space corre en un tier con GPU (DEVICE=cuda por
# variable de entorno tambien fuerza el valor si hace falta). Si no hay
# GPU disponible, cae a CPU como siempre.
DEVICE = os.environ.get("DEVICE") or ("cuda" if torch.cuda.is_available() else "cpu")

if DEVICE == "cpu":
    # Limita hilos a los núcleos disponibles (evita overhead en Spaces pequeños).
    # No tiene sentido en GPU, donde el computo no lo hace la CPU.
    torch.set_num_threads(max(1, os.cpu_count() or 1))
    # set_num_interop_threads solo puede llamarse una vez y antes de cualquier
    # operación paralela; lo protegemos por si el entorno ya lo fijó.
    try:
        torch.set_num_interop_threads(1)
    except RuntimeError:
        pass

torch.set_grad_enabled(False)  # solo inferencia, nunca necesitamos gradientes

# Disponibilidad de scaled_dot_product_attention (PyTorch >= 2.0).
# Si no está disponible, usamos el softmax manual original como fallback.
_HAS_SDPA = hasattr(F, "scaled_dot_product_attention")

REPO_ID = os.environ.get("MTP_REPO_ID", "TeszenAI/MTP-2.7")  # <-- ajusta al nombre real de tu repo/Space en el Hub
FILENAME = os.environ.get("MTP_FILENAME", "MTP2_7_MODEL.pt")

# Origenes permitidos para CORS. Por defecto "*" (como antes), pero se puede
# restringir en produccion con la variable de entorno MTP_ALLOWED_ORIGINS
# (separados por coma), por ejemplo: "https://teszen.com,https://www.teszen.com"
ALLOWED_ORIGINS = [
    o.strip() for o in os.environ.get("MTP_ALLOWED_ORIGINS", "*").split(",") if o.strip()
] or ["*"]

# Techo de caracteres del input antes de tokenizar. No es por seguridad (el
# modelo igual solo "ve" los ultimos BLOCK_SIZE tokens), es para no perder
# tiempo tokenizando un texto absurdamente largo por error o abuso.
MAX_INPUT_CHARS = 4000

# Si es "1", los errores devueltos por /generate incluyen el detalle interno
# de la excepcion (util mientras desarrollas). En produccion, dejar en "0"
# para no filtrarle al cliente detalles internos del servidor.
DEBUG_ERRORS = os.environ.get("MTP_DEBUG_ERRORS", "0") == "1"

# Serializa las llamadas a generate(): sin esto, dos requests concurrentes
# (por ejemplo la UI de Gradio y el endpoint /generate al mismo tiempo, o
# varias visitas simultaneas al sitio) compiten por los mismos hilos de CPU
# y todas terminan mas lentas en vez de una rapida y la otra esperando. Con
# el lock, cada generacion corre de punta a punta antes de que empiece la
# siguiente -- mismo comportamiento de fondo que antes bajo carga baja, pero
# estable bajo carga alta en vez de degradarse.
_generation_lock = threading.Lock()

# ---------------- Arquitectura MTP-2.x: RoPE + SwiGLU, con KV-cache ----------------
def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rope(q, k, cos, sin):
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)
    q_rot = (q * cos) + (rotate_half(q) * sin)
    k_rot = (k * cos) + (rotate_half(k) * sin)
    return q_rot, k_rot

class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_seq_len, base=10000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._build_cache(max_seq_len)

    def _build_cache(self, seq_len):
        t = torch.arange(seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)
        self.max_seq_len_cached = seq_len

    def forward(self, seq_len, device, dtype, offset=0):
        # Con KV-cache, `offset` es cuantos tokens ya estan en la cache: el
        # token nuevo necesita el angulo correspondiente a SU posicion
        # absoluta, no a la posicion relativa dentro de este forward.
        if offset + seq_len > self.max_seq_len_cached:
            self._build_cache(offset + seq_len)
        cos = self.cos_cached[offset:offset + seq_len].to(device=device, dtype=dtype)
        sin = self.sin_cached[offset:offset + seq_len].to(device=device, dtype=dtype)
        return cos, sin


class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.qkv = nn.Linear(n_embd, 3 * n_embd)
        self.proj = nn.Linear(n_embd, n_embd)
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        mask = torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
        self.register_buffer("mask", mask)

    def forward(self, x, cos, sin, past_kv=None, use_cache=False):
        B, T, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        # RoPE se aplica ANTES de guardar en cache, con el angulo absoluto de
        # cada token (pasado por `cos`/`sin`, ya calculado con el offset
        # correcto en MTP.forward). Asi el k cacheado ya trae rotada su
        # posicion real y no hay que re-rotar nada en pasos futuros.
        q, k = apply_rope(q, k, cos, sin)

        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        present_kv = (k, v) if use_cache else None

        is_causal = (past_kv is None) and (T > 1)

        if _HAS_SDPA:
            out = F.scaled_dot_product_attention(
                q, k, v, attn_mask=None, dropout_p=0.0, is_causal=is_causal,
            )
        else:
            att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
            if is_causal:
                Tk = k.size(-2)
                causal_mask = torch.tril(torch.ones(T, Tk, device=x.device, dtype=torch.bool))
                att = att.masked_fill(~causal_mask, float("-inf"))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            out = att @ v

        out = out.transpose(1, 2).contiguous().view(B, T, C)
        out = self.resid_dropout(self.proj(out))
        return out, present_kv


class SwiGLU(nn.Module):
    def __init__(self, n_embd, dropout):
        super().__init__()
        hidden = int(2 * (4 * n_embd) / 3)
        hidden = ((hidden + 7) // 8) * 8
        self.w_gate = nn.Linear(n_embd, hidden, bias=False)
        self.w_up = nn.Linear(n_embd, hidden, bias=False)
        self.w_down = nn.Linear(hidden, n_embd, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.dropout(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)))


class RMSNorm(nn.Module):
    """Debe coincidir exactamente con la version de entrenamiento. No
    necesita ningun cambio para funcionar con KV-cache: normaliza cada
    posicion de forma independiente, igual que LayerNorm."""
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        norm = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return norm * self.weight


class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.ln1 = RMSNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
        self.ln2 = RMSNorm(n_embd)
        self.ff = SwiGLU(n_embd, dropout)

    def forward(self, x, cos, sin, past_kv=None, use_cache=False):
        attn_out, present_kv = self.attn(self.ln1(x), cos, sin, past_kv=past_kv, use_cache=use_cache)
        x = x + attn_out
        x = x + self.ff(self.ln2(x))
        return x, present_kv


class MTP(nn.Module):
    def __init__(self, vocab_size, block_size, n_layer, n_head, n_embd, dropout):
        super().__init__()
        self.block_size = block_size
        self.head_dim = n_embd // n_head
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.rope = RotaryEmbedding(self.head_dim, max_seq_len=block_size)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
        self.ln_f = RMSNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight

    def forward(self, idx, past_key_values=None, use_cache=False, pos_offset=0):
        B, T = idx.shape
        x = self.tok_emb(idx)
        x = self.drop(x)
        cos, sin = self.rope(T, idx.device, x.dtype, offset=pos_offset)

        new_past = [] if use_cache else None
        for i, block in enumerate(self.blocks):
            past_kv = past_key_values[i] if past_key_values is not None else None
            x, present_kv = block(x, cos, sin, past_kv=past_kv, use_cache=use_cache)
            if use_cache:
                new_past.append(present_kv)

        x = self.ln_f(x)
        logits = self.lm_head(x)
        return logits, new_past


# ---------------- Carga del checkpoint (una sola vez, al iniciar el Space) ----------------
def _download_checkpoint_with_retries(repo_id, filename, max_retries=5):
    """Reintenta la descarga con backoff. hf_hub_download puede fallar por un
    corte de red momentaneo o un problema del backend Xet de HF; sin esto,
    un solo fallo transitorio tira abajo el arranque completo del Space."""
    last_err = None
    for attempt in range(1, max_retries + 1):
        try:
            logger.info(f"Descargando checkpoint desde el Hub (intento {attempt}/{max_retries})...")
            return hf_hub_download(repo_id=repo_id, filename=filename)
        except Exception as e:
            last_err = e
            logger.warning(f"Fallo la descarga del checkpoint: {e}")
            if attempt < max_retries:
                time.sleep(5 * attempt)
    raise RuntimeError(f"No se pudo descargar el checkpoint tras {max_retries} intentos") from last_err

ckpt_path = _download_checkpoint_with_retries(REPO_ID, FILENAME)
checkpoint = torch.load(ckpt_path, map_location=DEVICE)

cfg = checkpoint["config"]
special = checkpoint["special_tokens"]
gen_defaults = checkpoint["generation_defaults"]

PAD_ID, BOS_ID, EOS_ID, UNK_ID = special["pad_id"], special["bos_id"], special["eos_id"], special["unk_id"]

# El tokenizer es BPE (SentencePiece) entrenado desde cero junto con el modelo.
# No es un modelo preentrenado externo: viene embebido como bytes dentro del
# mismo checkpoint que los pesos. Se carga directo desde memoria con
# load_from_serialized_proto, sin necesidad de escribirlo a disco primero.
sp = spm.SentencePieceProcessor()
sp.load_from_serialized_proto(checkpoint["spm_model_bytes"])

model = MTP(
    vocab_size=cfg["vocab_size"], block_size=cfg["block_size"],
    n_layer=cfg["n_layer"], n_head=cfg["n_head"],
    n_embd=cfg["n_embd"], dropout=cfg["dropout"],
).to(DEVICE)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

BLOCK_SIZE = cfg["block_size"]

logger.info(
    f"MTP cargado ({checkpoint['meta']['model_name']}, "
    f"entrenado con {checkpoint['meta']['trained_examples']} ejemplos) "
    f"| device={DEVICE} | SDPA={'sí' if _HAS_SDPA else 'no (fallback manual)'}"
)


import re as _re_indent

def protect_indentation(text):
    """Debe coincidir exactamente con la funcion usada en el entrenamiento."""
    lines = text.split("\n")
    new_lines = []
    for line in lines:
        stripped = line.lstrip(" ")
        n_spaces = len(line) - len(stripped)
        n_levels = n_spaces // 4
        remainder = n_spaces % 4
        if n_levels > 0:
            prefix = " " + " ".join(["<tab>"] * n_levels) + " " + " " * remainder
        else:
            prefix = " " * remainder
        new_lines.append(prefix + stripped)
    text = "\n".join(new_lines)
    def _repl(m):
        n = len(m.group())
        return " " + " ".join(["<nl>"] * n) + " "
    text = _re_indent.sub(r"\n+", _repl, text)
    return text


def restore_indentation(text):
    text = _re_indent.sub(r"(<nl>\s*)+", lambda m: "\n" * m.group().count("<nl>"), text)
    text = _re_indent.sub(r"(<tab>\s*)+", lambda m: "    " * m.group().count("<tab>"), text)
    return text


def encode_text(s):
    return sp.encode(protect_indentation(s), out_type=int)


def decode_ids(ids):
    text = sp.decode([i for i in ids if i not in (PAD_ID, BOS_ID, EOS_ID)])
    return restore_indentation(text)


def _block_repeated_ngrams(generated_ids, logits, ngram_size):
    """Prohibe repetir literalmente un n-grama ya generado en esta misma
    respuesta (tecnica de decoding tipo GPT-2/3, no un modelo de n-gramas:
    el modelo que predice sigue siendo 100% transformer con KV-cache)."""
    if ngram_size <= 0 or len(generated_ids) < ngram_size:
        return logits
    prefix = tuple(generated_ids[-(ngram_size - 1):])
    banned = set()
    for i in range(len(generated_ids) - ngram_size + 1):
        if tuple(generated_ids[i:i + ngram_size - 1]) == prefix:
            banned.add(generated_ids[i + ngram_size - 1])
    if banned:
        logits[0, list(banned)] = float("-inf")
    return logits


# ---------------- Generación (con KV-cache) ----------------
@torch.inference_mode()
def generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size=3):
    past_key_values = None
    cache_len = 0  # cuántos tokens del extremo derecho de `idx` ya están en la caché

    for _ in range(max_new_tokens):
        total_len = idx.shape[1]

        if total_len <= BLOCK_SIZE:
            if past_key_values is None:
                # Primer paso: una sola pasada ("prefill") sobre todo el prompt.
                logits, past_key_values = model(idx, use_cache=True)
                cache_len = total_len
            else:
                # Pasos siguientes: solo se procesa el último token generado,
                # reutilizando la caché de todo lo anterior.
                last_token = idx[:, -1:]
                logits, past_key_values = model(
                    last_token,
                    past_key_values=past_key_values,
                    use_cache=True,
                    pos_offset=cache_len,
                )
                cache_len += 1
            logits = logits[:, -1, :]
        else:
            # Se superó block_size: mismo comportamiento que el modelo original
            # (ventana deslizante recalculada por completo). Solo ocurre en
            # respuestas muy largas; la caché se reinicia para esa ventana.
            idx_cond = idx[:, -BLOCK_SIZE:]
            logits, past_key_values = model(idx_cond, use_cache=True)
            cache_len = BLOCK_SIZE
            logits = logits[:, -1, :]

        logits = logits / max(temperature, 1e-5)

        if repetition_penalty and repetition_penalty != 1.0:
            # Vectorizado: antes era `for token_id in set(idx[0].tolist())`,
            # un bucle Python nuevo por cada token generado.
            unique_ids = torch.unique(idx[0])
            logits[0, unique_ids] /= repetition_penalty

        logits = _block_repeated_ngrams(idx[0].tolist(), logits, no_repeat_ngram_size)

        if top_k is not None and top_k > 0:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = float("-inf")

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

        if top_p is not None and 0 < top_p < 1:
            sorted_probs, sorted_idx = torch.sort(probs, descending=True)
            cum_probs = torch.cumsum(sorted_probs, dim=-1)
            cutoff = cum_probs > top_p
            cutoff[:, 1:] = cutoff[:, :-1].clone()
            cutoff[:, 0] = False
            sorted_probs[cutoff] = 0.0
            sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
            next_id = sorted_idx.gather(-1, torch.multinomial(sorted_probs, 1))
        else:
            next_id = torch.multinomial(probs, num_samples=1)

        idx = torch.cat([idx, next_id], dim=1)
        if next_id.item() == EOS_ID:
            break

    return idx


def run_inference(text, max_new_tokens=None, temperature=None, top_k=None, top_p=None, repetition_penalty=None, no_repeat_ngram_size=None):
    """Núcleo de generación, reutilizado por la UI de Gradio y por la API /generate.
    No reduce calidad por estar en CPU: usa exactamente el mismo muestreo
    (top_k + top_p + repetition_penalty + bloqueo de n-gramas repetidos) que
    en la Celda 2 de entrenamiento, solo que ahora con KV-cache es notablemente
    más rápido en respuestas largas."""
    max_new_tokens = int(max_new_tokens) if max_new_tokens else gen_defaults["max_new_tokens"]
    temperature = float(temperature) if temperature is not None else gen_defaults["temperature"]
    top_k = int(top_k) if top_k is not None else gen_defaults["top_k"]
    top_p = float(top_p) if top_p is not None else gen_defaults["top_p"]
    repetition_penalty = float(repetition_penalty) if repetition_penalty is not None else gen_defaults["repetition_penalty"]
    no_repeat_ngram_size = int(no_repeat_ngram_size) if no_repeat_ngram_size is not None else gen_defaults.get("no_repeat_ngram_size", 3)

    # Techo máximo de generación: 4000 no era realista en CPU (cada token
    # adicional cuesta tiempo real). 700 sigue siendo una respuesta larga y
    # mantiene el tiempo de respuesta bajo control en el peor caso.
    MAX_TOKENS_HARD_LIMIT = 700
    max_new_tokens = max(1, min(max_new_tokens, MAX_TOKENS_HARD_LIMIT))

    if len(text) > MAX_INPUT_CHARS:
        logger.warning(f"Input de {len(text)} caracteres recortado a {MAX_INPUT_CHARS}")
        text = text[:MAX_INPUT_CHARS]

    prefix = f"Usuario: {text}\nMTP: "
    ids = [BOS_ID] + encode_text(prefix)
    idx = torch.tensor([ids], dtype=torch.long, device=DEVICE)

    # Con el lock, si llegan varias generaciones al mismo tiempo (UI + API,
    # o varios usuarios a la vez) se procesan una despues de otra en vez de
    # pisarse los hilos de CPU entre si.
    with _generation_lock:
        out = generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size)

    new_ids = out[0].tolist()[len(ids):]
    return decode_ids(new_ids).strip()


def chat_fn(message, history, max_new_tokens, temperature, top_k, top_p, repetition_penalty):
    return run_inference(message, max_new_tokens, temperature, top_k, top_p, repetition_penalty)


# ---------------- Interfaz Gradio (para probar el modelo desde el navegador) ----------------
with gr.Blocks(title="MTP-2.5 Chat") as demo:
    gr.Markdown(f"# MTP-2.5\nModelo GPT (RoPE + SwiGLU + RMSNorm) entrenado desde cero, tokenizer BPE. Ejecutándose en {DEVICE.upper()}.")

    with gr.Accordion("Parámetros de generación", open=False):
        max_new_tokens_ui = gr.Slider(16, 4000, value=gen_defaults["max_new_tokens"], step=10, label="max_new_tokens")
        temperature_ui = gr.Slider(0.1, 2.0, value=gen_defaults["temperature"], step=0.05, label="temperature")
        top_k_ui = gr.Slider(0, 100, value=gen_defaults["top_k"], step=1, label="top_k")
        top_p_ui = gr.Slider(0.1, 1.0, value=gen_defaults["top_p"], step=0.05, label="top_p")
        repetition_penalty_ui = gr.Slider(1.0, 2.0, value=gen_defaults["repetition_penalty"], step=0.05,
                                           label="repetition_penalty")

    chatbot = gr.ChatInterface(
        fn=chat_fn,
        additional_inputs=[max_new_tokens_ui, temperature_ui, top_k_ui, top_p_ui, repetition_penalty_ui],
        title=None,
        examples=[
            ["Hola, ¿cómo estás?"],
            ["¿Cuánto es 8 + 5?"],
            ["Explícame qué es un algoritmo."],
        ],
        cache_examples=False,
    )

demo.queue(max_size=16)

# ---------------- API REST /generate (la que consume el PHP) ----------------
# El PHP hace: fetch(url, { method:'POST', body: JSON.stringify({text, max_tokens, temperature}) })
# y espera de vuelta: { "reply": "..." }
#
# IMPORTANTE:
# - ssr_mode=False: Gradio 6 usa un servidor Node.js aparte para SSR, que
#   intentaba levantarse en el puerto 7861 y chocaba. Lo desactivamos porque
#   no lo necesitamos para servir la API.
# - El middleware CORS se pasa vía app_kwargs ANTES de llamar a launch(),
#   porque una vez que la app arranca, Starlette ya no permite añadir
#   middleware (por eso fallaba con app.add_middleware() después).

class GenerateRequest(BaseModel):
    text: str
    max_tokens: Optional[int] = None
    temperature: Optional[float] = None
    top_k: Optional[int] = None
    top_p: Optional[float] = None
    repetition_penalty: Optional[float] = None
    no_repeat_ngram_size: Optional[int] = None


PORT = int(os.environ.get("PORT", 7860))
demo.launch(
    server_name="0.0.0.0",
    server_port=PORT,
    prevent_thread_lock=True,
    ssr_mode=False,
    app_kwargs={
        "middleware": [
            Middleware(CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_methods=["*"], allow_headers=["*"]),
        ]
    },
)

app = demo.app


@app.post("/generate")
def generate_endpoint(req: GenerateRequest):
    if not req.text or not req.text.strip():
        return {"reply": "Escribe algo para que pueda responder."}
    try:
        reply = run_inference(
            req.text,
            max_new_tokens=req.max_tokens,
            temperature=req.temperature,
            top_k=req.top_k,
            top_p=req.top_p,
            repetition_penalty=req.repetition_penalty,
            no_repeat_ngram_size=req.no_repeat_ngram_size,
        )
        if not reply:
            reply = "No pude generar una respuesta."
        return {"reply": reply}
    except Exception as e:
        # El detalle completo va al log del Space (con traceback), no a la
        # respuesta publica: devolver la excepcion cruda a quien llama podria
        # filtrar rutas internas, nombres de variables, etc. Con
        # MTP_DEBUG_ERRORS=1 se puede activar el detalle mientras desarrollas.
        logger.error(f"Error generando respuesta: {e}\n{traceback.format_exc()}")
        reply = f"Error del modelo: {e}" if DEBUG_ERRORS else "Ocurrió un error al generar la respuesta. Intenta de nuevo en un momento."
        return {"reply": reply}


@app.get("/generate")
def generate_health():
    # Solo para poder comprobar en el navegador que la ruta existe (GET no genera texto)
    return {"status": "ok", "info": "Usa POST con JSON {text, max_tokens, temperature}"}


@app.get("/health")
def health():
    # Health check real: confirma que el modelo esta cargado y listo, no
    # solo que el proceso esta vivo. Util para monitoreo externo (uptime
    # checks) o para que el PHP sepa si conviene reintentar mas tarde.
    return {
        "status": "ok",
        "model": checkpoint["meta"]["model_name"],
        "device": DEVICE,
        "trained_examples": checkpoint["meta"]["trained_examples"],
    }


# demo.launch(prevent_thread_lock=True) ya dejó el servidor corriendo en un
# hilo en segundo plano (un solo proceso, un solo puerto). Mantenemos vivo
# el hilo principal para que el contenedor del Space no termine.
demo.block_thread()