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=1XqImFIKiuRDhSDK6Rm6dAZbHm03NdzQa"
URL_PKL   = "https://drive.google.com/uc?export=download&id=156BWHHGi-JuD9EM2Nek1mNcyitivQWAH"
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 de nutrición llamado NutriActive.

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?")
- NUTRICION: preguntas sobre nutrición, dieta, salud, alimentos, calorías, macros, IMC, 
  planes alimenticios, recetas, suplementos, hábitos saludables, Y TAMBIÉN cualquier 
  pregunta relacionada con NutriActive como empresa: sus cursos, servicios, productos, 
  planes, precios, programas, etc.
- OTRO: preguntas claramente NO relacionadas con nutrición, salud ni NutriActive 
  (matemáticas, historia, tecnología general, etc.)

IMPORTANTE: Ante la duda, clasifica como NUTRICION. Solo usa OTRO cuando estés 
completamente seguro de que no tiene relación con nutrición ni con NutriActive.

Responde SOLO con la categoría, sin explicación.

Mensaje: {query}
Categoría:""",
    input_variables=["query"]
)

SALUDO_PROMPT = PromptTemplate(
    template="""Eres NutriActive, un asistente amigable especializado en nutrición y salud.
Responde de forma natural y cálida al siguiente mensaje casual del usuario.
Si el usuario se despide o agradece, invítalo a preguntar sobre nutrición.

Mensaje: {query}
Respuesta:""",
    input_variables=["query"]
)

RAG_PROMPT = PromptTemplate(
    template="""Eres NutriActive, un asistente experto en nutrición y salud.
Tu tarea es responder basándote en el contexto proporcionado.
Si el contexto no tiene suficiente información, usa tu conocimiento general sobre nutrición para dar una respuesta útil.
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.3, model_name="llama-3.3-70b-versatile")

        # 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="NutriActive 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 NutriActive 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 NutriActive, especializado en nutrición y salud. ¿Tienes alguna pregunta sobre alimentación, dietas o bienestar? 🥗",
                "intent": "OTRO",
                "sources": []
            }

        else:
            # NUTRICION 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":   "NUTRICION",
                "sources":  sources
            }

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