"""
Space: CHRISTS20/modeloEntrenado (SDK: gradio)
Agente UML — qwen2.5-coder-7b fine-tune (Q4_K_M) con doble cara:
  - UI Gradio en / .................... vidriera para la defensa
  - POST /v1/chat/completions ......... API OpenAI-compatible (la consume el frontend)
  - GET  /v1/health ................... ping anti-sueño (UptimeRobot / cron)

Contrato con el frontend (ex1-frontend, src/lib/ia/agentLoop.ts):
  - Acepta messages con roles system/user/assistant/tool (+ tool_calls en
    assistant) y el array `tools` en formato OpenAI function-calling.
  - Devuelve {"choices": [{"message": {...}, "finish_reason": ...}]}.
  - NO necesita tool_calls nativos: si el modelo emite <tool_call>{"name":..}
    en el texto (es lo que aprendió en el fine-tune), el frontend lo rescata
    con su extractor de fallback. Por eso el template del GGUF se deja en
    auto (None) salvo override por env CHAT_FORMAT.

Pesos horneados: el .gguf viaja en el repo del Space (git-lfs), se carga por
ruta local. Cero descargas en runtime -> el despertar solo recarga a RAM.
"""
import os
import time
import uuid
from typing import Any, Dict, List, Optional

import gradio as gr
import uvicorn
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Pesos: los repos de Space tienen tope de 1 GB, así que el .gguf (4.68 GB)
# vive en el model repo y se descarga una vez por arranque al disco efímero
# (50 GB disponibles, sobra). Si el archivo ya existe local (dev), se usa.
MODEL_REPO = os.getenv("MODEL_REPO", "ByCrilux-XMode/miModeloUML")
MODEL_FILE = os.getenv("MODEL_FILE", "qwen2.5-coder-7b-instruct.Q4_K_M-001.gguf")
MODEL_PATH = os.getenv("MODEL_PATH", MODEL_FILE)
if not os.path.exists(MODEL_PATH):
    print(f"[boot] Descargando {MODEL_FILE} desde {MODEL_REPO} ...", flush=True)
    MODEL_PATH = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
N_CTX = int(os.getenv("N_CTX", "8192"))
CHAT_FORMAT: Optional[str] = os.getenv("CHAT_FORMAT", "").strip() or None
DEFAULT_TEMPERATURE = float(os.getenv("TEMPERATURE", "0.2"))
DEFAULT_MAX_TOKENS = int(os.getenv("MAX_TOKENS", "2048"))

print(f"[boot] Cargando {MODEL_PATH} (n_ctx={N_CTX}, chat_format={CHAT_FORMAT}) ...", flush=True)
llm = Llama(model_path=MODEL_PATH, n_ctx=N_CTX, chat_format=CHAT_FORMAT, verbose=False)
print("[boot] Modelo listo.", flush=True)


def sanitize_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """Normaliza roles al dialecto que llama-cpp-python acepta siempre.

    Espeja el TEMPLATE con el que se entrenó el modelo (ver Modelfile):
    los mensajes role=tool van como user envueltos en <tool_response>.
    Los tool_calls del assistant se conservan tal cual.
    """
    out: List[Dict[str, Any]] = []
    for m in messages or []:
        role = m.get("role", "user")
        if role == "tool":
            out.append({
                "role": "user",
                "content": f"<tool_response>\n{m.get('content', '')}\n</tool_response>",
            })
        elif role == "assistant" and m.get("tool_calls"):
            out.append({
                "role": "assistant",
                "content": m.get("content") or "",
                "tool_calls": m.get("tool_calls"),
            })
        else:
            out.append({"role": role, "content": m.get("content", "")})
    return out


def complete(
    messages: List[Dict[str, Any]],
    tools: Optional[List[Dict[str, Any]]] = None,
    tool_choice: Any = None,
    temperature: float = DEFAULT_TEMPERATURE,
    max_tokens: int = DEFAULT_MAX_TOKENS,
) -> Dict[str, Any]:
    kwargs: Dict[str, Any] = {
        "messages": sanitize_messages(messages),
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    # tool_choice sin tools rompe en algunas versiones: solo si hay tools.
    if tools:
        kwargs["tools"] = tools
        kwargs["tool_choice"] = tool_choice or "auto"
    out = llm.create_chat_completion(**kwargs)
    msg = out["choices"][0]["message"]
    return {
        "role": "assistant",
        "content": msg.get("content") or "",
        **({"tool_calls": msg["tool_calls"]} if msg.get("tool_calls") else {}),
    }


app = FastAPI(title="modeloEntrenado — agente UML")


@app.get("/health")
@app.get("/v1/health")
def health() -> Dict[str, Any]:
    return {"ok": True, "model": os.path.basename(MODEL_PATH)}


@app.post("/v1/chat/completions")
def chat_completions(body: Dict[str, Any]) -> JSONResponse:
    try:
        message = complete(
            messages=body.get("messages", []),
            tools=body.get("tools"),
            tool_choice=body.get("tool_choice"),
            temperature=float(body.get("temperature", DEFAULT_TEMPERATURE)),
            max_tokens=int(body.get("max_tokens", DEFAULT_MAX_TOKENS)),
        )
        return JSONResponse({
            "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": body.get("model", os.path.basename(MODEL_PATH)),
            "choices": [{"index": 0, "message": message, "finish_reason": "stop"}],
        })
    except Exception as err:  # error claro para el trace del frontend, no stack
        return JSONResponse({"error": f"Space modeloEntrenado: {err}"}, status_code=500)


def ui_fn(prompt: str) -> str:
    """UI Gradio: chat simple de un turno (la API de arriba es lo serio)."""
    if not (prompt or "").strip():
        return "Escribí un mensaje."
    msg = complete([{"role": "user", "content": prompt.strip()}])
    return msg.get("content", "")


demo = gr.Interface(
    fn=ui_fn,
    inputs=gr.Textbox(label="Mensaje"),
    outputs=gr.Textbox(label="Respuesta UML"),
    title="Agente UML — modeloEntrenado",
    description="UI de demostración. El frontend usa POST /v1/chat/completions con tools.",
)
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)