import os
import requests
import shutil
from langchain_community.vectorstores import FAISS
from fastapi import FastAPI
from pydantic import BaseModel
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_groq import ChatGroq

# --------------------------------------------------------
# CACHÉ EN /tmp
# --------------------------------------------------------
TEMP_CACHE_DIR = '/tmp/huggingface_cache'
os.environ['TRANSFORMERS_CACHE'] = TEMP_CACHE_DIR
os.environ['HF_HOME'] = TEMP_CACHE_DIR
os.environ['SENTENCE_TRANSFORMERS_HOME'] = TEMP_CACHE_DIR
os.makedirs(TEMP_CACHE_DIR, exist_ok=True)

# --------------------------------------------------------
# 1. CONFIGURACIÓN
# --------------------------------------------------------
URL_FAISS = "https://drive.google.com/uc?export=download&id=1hiVycS4DQHO1MBdC-L_z1TXA6sJO_Y-r"
URL_PKL = "https://drive.google.com/uc?export=download&id=1vbG8unx88Kb5jn7puGv1gqSM4S6rIUQC"
DOWNLOAD_DIR  = "/tmp/db_faiss"
DB_FAISS_PATH = DOWNLOAD_DIR

# --------------------------------------------------------
# 2. CLASIFICADOR DE INTENCIÓN  ← NUEVO
# --------------------------------------------------------
INTENT_PROMPT = PromptTemplate(
    template="""Eres un clasificador de intenciones para un asistente del portal de la Universidad Poltécnica de Aragua.
Analiza el mensaje del usuario y clasifícalo en UNA de estas categorías:
- SALUDO: saludos, despedidas, conversación casual ("hola", "gracias", "adiós", "¿cómo estás?")
- UNIVERSIDAD: preguntas sobre carreras o programas, investigación, cursos, admisiones, notas, proyectos, postgrado, 
  PNF, PNFA, diplomados, servivios, Y TAMBIÉN cualquier pregunta relacionada con La Universidad relacionado con: sus autoridades, reglamentos, 
  servivios estudiantiles, precios de cursos, programas, etc.
- OTRO: preguntas claramente NO relacionadas con la Universidad tales como: matemáticas, historia, tecnología general, etc.
IMPORTANTE: Ante la duda, clasifica como Universidad Politécnica de Aragua o UPT ARagua. Solo usa OTRO cuando estés 
completamente seguro de que no tiene relación con la Universidad.
Responde SOLO con la categoría, sin explicación.
Mensaje: {query}
Categoría:""",
    input_variables=["query"]
)

SALUDO_PROMPT = PromptTemplate(
    template="""Eres UPTA bot, un Asistente Virtual de la UPT Aragua. Estas aquí para ayudar con información sobre admisiones, programas académicos, 
    servicios, becas y mucho más. Si el usuario se despide o agradece, invítalo a preguntar sobre la universidad.
Mensaje: {query}
Respuesta:""",
    input_variables=["query"]
)

RAG_PROMPT = PromptTemplate(
    template="""Eres UPTA bot, un Asistente Virtual experto  de la UPT Aragua. Estas aquí para ayudar con información sobre 
    admisiones, programas académicos, servicios, becas y mucho más. Tu tarea es responder basándote en el contexto proporcionado. Si el contexto
    no tiene suficiente información, pide al usuario que te proporcione una presunta más específica sobre la UPT Aragua para dar una respuesta fiable. Sé amigable, claro y conciso.
Contexto de la base de datos: {context}
Pregunta del usuario: {question}
Respuesta:""",
    input_variables=["context", "question"]
)

# --------------------------------------------------------
# 3. FUNCIONES DE DESCARGA Y CARGA
# --------------------------------------------------------
class QueryRequest(BaseModel):
    query: str

def download_file(url, local_path):
    file_name = os.path.basename(local_path)
    print(f"Descargando: {file_name}...")
    headers = {'User-Agent': 'Mozilla/5.0'}
    try:
        response = requests.get(url, stream=True, headers=headers, timeout=30)
        if response.status_code == 403:
            raise PermissionError(f"Error 403: {file_name} no es público.")
        response.raise_for_status()
        os.makedirs(os.path.dirname(local_path), exist_ok=True)
        with open(local_path, 'wb') as f:
            shutil.copyfileobj(response.raw, f)
        print(f"✓ {file_name} descargado.")
    except requests.exceptions.RequestException as e:
        raise RuntimeError(f"Fallo al descargar {file_name}: {e}")

def load_and_configure_rag():
    try:
        download_file(URL_FAISS, os.path.join(DOWNLOAD_DIR, 'index.faiss'))
        download_file(URL_PKL,   os.path.join(DOWNLOAD_DIR, 'index.pkl'))

        print("Cargando embeddings...")
        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            model_kwargs={'device': 'cpu'},
            cache_folder=TEMP_CACHE_DIR
        )

        print("Cargando FAISS...")
        vectorstore = FAISS.load_local(
            DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True
        )

        llm = ChatGroq(temperature=0.150, model_name="openai/gpt-oss-120b")

        # Cadena clasificadora de intención
        intent_chain = INTENT_PROMPT | llm

        # Cadena para saludos
        saludo_chain = SALUDO_PROMPT | llm

        # Cadena RAG principal
        retriever  = vectorstore.as_retriever(search_kwargs={"k": 4})
        rag_chain  = (
            {"context": retriever, "question": RunnablePassthrough()}
            | RAG_PROMPT
            | llm
        )

        return intent_chain, saludo_chain, rag_chain, retriever

    except Exception as e:
        print(f"Error CRÍTICO al inicializar: {type(e).__name__}: {e}")
        raise RuntimeError(f"Falla al cargar RAG: {e}")

# --------------------------------------------------------
# 4. FASTAPI
# --------------------------------------------------------
app = FastAPI(title="UPT Aragua bot RAG API")

intent_chain = saludo_chain = qa_chain = retriever = None

try:
    intent_chain, saludo_chain, qa_chain, retriever = load_and_configure_rag()
except RuntimeError:
    pass

@app.get("/")
def home():
    if qa_chain is None:
        return {"error": "RAG no inicializado. Revisa los logs."}
    return {"message": "API UPT Aragua bot operativa. Usa /query."}

@app.post("/query")
async def process_query(request: QueryRequest):
    if qa_chain is None:
        return {"error": "El sistema RAG no se pudo cargar."}

    try:
        # ── 1. Clasificar intención ──────────────────────────────
        intent_result = intent_chain.invoke({"query": request.query})
        intent = intent_result.content.strip().upper()
        print(f"[Intent] '{request.query}' → {intent}")

        # ── 2. Ruta según intención ──────────────────────────────
        if "SALUDO" in intent:
            respuesta = saludo_chain.invoke({"query": request.query})
            return {
                "query": request.query,
                "response": respuesta.content,
                "intent": "SALUDO",
                "sources": []
            }

        elif "OTRO" in intent:
            return {
                "query": request.query,
                "response": "Soy UPTA bot, estoy especializado en la UPT Aragua. ¿Tienes alguna pregunta sobre programas, carreras, inscripciones, fechas...? 🥗",
                "intent": "OTRO",
                "sources": []
            }

        else:
            # UNIVERSIDAD o cualquier categoría no reconocida → RAG
            respuesta  = qa_chain.invoke(request.query)
            docs       = retriever.invoke(request.query)
            sources    = [doc.metadata.get("source", "N/A") for doc in docs]
            return {
                "query":    request.query,
                "response": respuesta.content,
                "intent":   "UNIVERSIDAD",
                "sources":  sources
            }

    except Exception as e:
        return {"error": f"Error al procesar la consulta: {e}"}