import os
from pathlib import Path
from time import perf_counter
from typing import Any

import pandas as pd
from dotenv import load_dotenv
from groq import Groq
from langchain_community.vectorstores import FAISS

from app.core.logging import InstrumentedGroqClient, get_audit_logger, get_logger
from app.llm import LLMTask
from app.llm.task_text import generate_task_text
from modules.processing.embeddings import get_hf_embeddings


load_dotenv()
EMBEDDINGS_MODEL = os.getenv("EMBEDDINGS_MODEL", "intfloat/e5-small-v2")
GROQ_MODEL_CUPS = os.getenv("GROQ_MODEL_CUPS", "openai/gpt-oss-120b")
logger = get_logger(__name__)
audit_logger = get_audit_logger()


class CUPSRetriever:
    def __init__(self, faiss_index_path, groq_api_key=None, groq_model=None, llm_router: Any | None = None):
        self.embeddings = get_hf_embeddings(EMBEDDINGS_MODEL)
        self.db = None
        self._index_load_error = None
        self.index_path = faiss_index_path
        self.db = self._cargar_indice(faiss_index_path)
        self.llm_router = llm_router
        self.groq = None
        if self.llm_router is None:
            api_key = groq_api_key or os.getenv("GROQ_API_KEY")
            if not api_key:
                raise RuntimeError("Se requiere GROQ_API_KEY")
            self.groq = InstrumentedGroqClient(
                Groq(api_key=api_key),
                provider="groq",
                service="cups_retriever",
            )
        self.groq_model = groq_model or GROQ_MODEL_CUPS

    def _cargar_indice(self, path: str):
        try:
            index_dir = Path(path)
            if not index_dir.exists():
                raise FileNotFoundError(f"Directorio no encontrado: {index_dir}")
            if not (index_dir / "index.faiss").exists():
                raise FileNotFoundError(f"Archivo index.faiss no encontrado en: {index_dir}")
            return FAISS.load_local(str(index_dir), self.embeddings, allow_dangerous_deserialization=True)
        except Exception as e:
            self._index_load_error = e
            logger.error("Error cargando indice FAISS CUPS: %s", str(e))
            return None

    def _require_db(self):
        if self.db is None:
            detalle = f"{self._index_load_error}" if self._index_load_error else "desconocido"
            raise RuntimeError(
                "Indice FAISS CUPS no disponible. "
                f"Ruta: {self.index_path}. "
                f"Detalle: {detalle}. "
                "Genera el indice con: python modules/processing/cie10/carga_cups.py"
            )

    def buscar(self, consulta, k=30):
        started_at = perf_counter()
        self._require_db()
        try:
            docs = self.db.max_marginal_relevance_search(consulta, k=k, lambda_param=0.7)
            audit_logger.rag_event(
                action="cups.buscar",
                outcome="success",
                service="cups_retriever",
                resource={"index_path": self.index_path, "query_type": "procedimiento", "index_name": "cups"},
                metrics={
                    "duration_ms": round((perf_counter() - started_at) * 1000, 3),
                    "k": k,
                    "result_count": len(docs),
                },
            )
            return docs
        except Exception as exc:
            audit_logger.rag_event(
                action="cups.buscar",
                outcome="error",
                service="cups_retriever",
                resource={"index_path": self.index_path, "query_type": "procedimiento", "index_name": "cups"},
                metrics={"duration_ms": round((perf_counter() - started_at) * 1000, 3), "k": k},
                error={"class": exc.__class__.__name__, "message": str(exc)},
            )
            raise

    def generar_groq(self, consulta, contexto):
        prompt = (
            "Eres un asistente experto en codificación CUPS. "
            "Para la siguiente descripción médica, entrega el código y nombre:\n\n"
            f"Consulta: {consulta}\n\nContexto:\n{contexto}\n\n"
            "Formato: <codigo> - <nombre>\n"
            "Si no encuentras coincidencia no muestres NADA, y omite el codigo"  # , responde: 'No se pudo determinar el código CUPS.'"
        )
        if self.llm_router is not None:
            response = generate_task_text(
                self.llm_router,
                task=LLMTask.CUPS_RESOLUTION,
                prompt=prompt,
                system_prompt="Eres experto en codificación CUPS.",
                metadata={"operation": "cups_retriever_generate", "consulta": str(consulta)[:200]},
                generation_config={"temperature": 0.1, "top_p": 0.1},
            )
            return response.content.strip()

        resp = self.groq.chat.completions.create(
            model=self.groq_model,
            temperature=0.1,
            top_p=0.1,
            messages=[
                {"role": "system", "content": "Eres experto en codificación CUPS."},
                {"role": "user", "content": prompt},
            ],
        )
        return resp.choices[0].message.content.strip()

    def asignar_codigos(self, items, temperature=0.0):
        resultados = []
        for item in items:
            # Retrieve similar documents
            docs = self.buscar(item, k=30)

            # FIXED: Access metadata from each document
            contexto = "\n".join(f"{doc.metadata['codigo']} - {doc.metadata['nombre']}" for doc in docs)

            codigo = self.generar_groq(item, contexto)
            resultados.append({"procedimiento": item, "codigo_cups": codigo})

        return pd.DataFrame(resultados)


def main():
    # Ajusta este path al índice FAISS que construiste
    indice_path = "cups_faiss"

    retriever = CUPSRetriever(
        faiss_index_path=indice_path,
        groq_api_key=os.getenv("GROQ_API_KEY"),
        groq_model=os.getenv("GROQ_MODEL_CUPS", "openai/gpt-oss-120b"),
    )

    # Procedimientos de prueba
    procedimientos = [
        "PUNCIÓN CISTERNAL, VIA LATERAL",
        "Reducción cerrada de fractura de falange proximal del 1er dedo en pie izquierdo",
        "Reducción cerrada falanges pie (una a dos)",
        "Lavado mas desbridamiento de lesión de tejido profundo de herida en pie izquierdo",
        # "Portátiles con fluoroscopia y/o intensificador de imagen (practicado en quirófanos)"
        # "TOMOGRAFÍA AXIAL COMPUTADA DE MIEMBROS INFERIORES Y ARTICULACIONES"
    ]

    df = retriever.asignar_codigos(procedimientos)
    print("Resultados CUPS:")
    print(df.to_string(index=False))


if __name__ == "__main__":
    main()
