import os

# Redirigir caché de Hugging Face a un directorio temporal para evitar errores de permisos
os.environ["HF_HOME"] = "/tmp/hfcache"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hfcache"
os.makedirs("/tmp/hfcache", exist_ok=True)

from fastapi import FastAPI, Request
from pydantic import BaseModel
from typing import List, Optional, Dict
from transformers import pipeline
import json
import re

app = FastAPI(title="Trivia y Consultas IA generativa ligera")

# Cambia aquí el modelo y tokens para pruebas rápidas
generator = pipeline(
    "text-generation",
    model="distilgpt2",        # Modelo más pequeño y rápido para pruebas
    max_new_tokens=50,         # Menos tokens para responder más rápido
    do_sample=True,
    temperature=0.9
)

# Modelo de moderación pequeño (toxicidad)
moderator = pipeline(
    "text-classification",
    model="unitary/toxic-bert"
)

class TriviaResponse(BaseModel):
    pregunta: str
    opciones: List[str]
    respuesta_correcta: int
    explicacion: str

class CompecomRequest(BaseModel):
    pregunta: str
    historial: Optional[List[Dict[str, str]]] = None  # Para contexto conversacional

class CompecomResponse(BaseModel):
    respuesta: str
    moderado: Optional[bool] = True
    moderacion: Optional[str] = None

@app.post("/trivia", response_model=TriviaResponse)
async def trivia(request: Request):
    print("Solicitud recibida en /trivia")  # Log para depuración
    try:
        body = await request.json()
        tema = body.get("tema", "competencia económica")
    except Exception:
        tema = "competencia económica"

    prompt = (
        f"Genera una trivia de opción múltiple de 4 opciones sobre el tema '{tema}'. "
        "Devuelve SOLO el JSON, SIN ningún texto extra. Ejemplo:\n"
        '{"pregunta": "¿Qué es un monopolio?", "opciones": ["Una sola empresa controla el mercado", "Muchas empresas compiten", "El Estado regula todo", "No hay empresas"], "respuesta_correcta": 0, "explicacion": "Un monopolio ocurre cuando una sola empresa domina el mercado."}\n'
        "La pregunta debe ser clara, las opciones concisas y la explicación breve. Recuerda: SOLO el JSON."
    )

    for _ in range(3):
        print("Enviando prompt al modelo generativo...")  # Log para depuración
        output = generator(prompt)[0]['generated_text']
        print("Respuesta recibida del modelo generativo.")  # Log para depuración
        match = re.search(r"\{.*\}", output, re.DOTALL)
        if not match:
            continue
        try:
            trivia_data = json.loads(match.group(0))
            # Validaciones robustas
            if not all(k in trivia_data for k in ("pregunta", "opciones", "respuesta_correcta", "explicacion")):
                continue
            if not isinstance(trivia_data["opciones"], list) or len(trivia_data["opciones"]) != 4:
                continue
            if not isinstance(trivia_data["respuesta_correcta"], int) or not (0 <= trivia_data["respuesta_correcta"] < 4):
                continue
            mod_result = moderator(trivia_data["pregunta"] + " " + trivia_data["explicacion"])
            if any(r["label"] == "TOXIC" and r["score"] > 0.5 for r in mod_result):
                continue
            return TriviaResponse(**trivia_data)
        except Exception as e:
            print("Error procesando la respuesta del modelo:", e)
            continue

    # Fallback robusto si falla todo
    print("Usando trivia de respaldo")
    return TriviaResponse(
        pregunta="¿Qué es un monopolio?",
        opciones=[
            "Una sola empresa controla el mercado.",
            "Muchas empresas compiten.",
            "El Estado regula todo.",
            "No hay empresas."
        ],
        respuesta_correcta=0,
        explicacion="Un monopolio ocurre cuando una sola empresa domina el mercado."
    )

@app.post("/compecom", response_model=CompecomResponse)
async def compecom(req: CompecomRequest):
    print("Solicitud recibida en /compecom")  # Log para depuración
    historial = req.historial or []
    historial_texto = ""
    if historial:
        for turno in historial:
            preg = turno.get("pregunta", "")
            resp = turno.get("respuesta", "")
            if preg and resp:
                historial_texto += f"Pregunta: {preg}\nRespuesta: {resp}\n"

    prompt = (
        "Eres un experto en competencia económica, monopolios y el Plan de la Patria en Venezuela. "
        "Responde de forma breve y clara.\n"
        f"{historial_texto}Pregunta: {req.pregunta}\nRespuesta:"
    )

    for _ in range(2):
        print("Enviando prompt a modelo generativo para consulta...")  # Log
        output = generator(prompt)[0]['generated_text']
        print("Respuesta recibida del modelo generativo para consulta.")  # Log
        respuesta = output.strip()
        if not respuesta or len(respuesta) < 5:
            continue
        mod_result = moderator(respuesta)
        if any(r["label"] == "TOXIC" and r["score"] > 0.5 for r in mod_result):
            continue
        return CompecomResponse(
            respuesta=respuesta,
            moderado=True,
            moderacion="SAFE"
        )
    # Fallback claro
    print("Usando respuesta de respaldo para consulta")
    return CompecomResponse(
        respuesta="No se pudo generar una respuesta en este momento. Por favor, intenta de nuevo.",
        moderado=False,
        moderacion="TOXIC"
    )

@app.get("/")
def root():
    return {"status": "ok", "message": "Trivia y COMPECOM IA generativa lista"}