import os
import requests
import shutil
import re
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)
def extraer_id(url):
    # Buscamos el patrón que está entre 'd/' y '/view'
    match = re.search(r'd/(.*?)/view', url)
    
    if match:
        return match.group(1)
    else:
        return "No se pudo encontrar el ID en el enlace proporcionado."

# Ejemplo de uso:

faiss_id = extraer_id("https://drive.google.com/file/d/1nCvXViXaRn_LyRxzJ4OR8Q8RPVPP5HER/view?usp=drive_link")
pkl_id = extraer_id("https://drive.google.com/file/d/1P8M8mYf-SXNAaZNV4ANYJh8oL6xOEI-z/view?usp=drive_link")


# --------------------------------------------------------
# 1. CONFIGURACIÓN
# --------------------------------------------------------
URL_FAISS = (f"https://drive.google.com/uc?export=download&id={faiss_id}")
URL_PKL   = (f"https://drive.google.com/uc?export=download&id={pkl_id}")
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?")
- mango: preguntas su veneficios, vitaminas, salud, calorías, alimentos, etc.
- OTRO: preguntas no relacionadas con nutrición ni saludos

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

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

SALUDO_PROMPT = PromptTemplate(
    template="""hola soy tu ayudante para entender todo sobre los mangos🥭.

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

RAG_PROMPT = PromptTemplate(
    template="""eres un asistente personal centrado en enseñar sobre el mango y sus veneficios entonces .

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 mango ACTIVA!!!. 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 tu ayudante para darte todo sobre los mangos y su conocimiento frutal 🥭",
                "intent": "OTRO",
                "sources": []
            }

        else:
            # NUTRICION → RAG completo
            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}"}