import pandas as pd
import json
import re
import logging
from unidecode import unidecode
from langchain_community.vectorstores import FAISS
from modules.processing.embeddings import get_hf_embeddings


from groq import Groq
import os
from groq import Groq
import pandas as pd
from dotenv import load_dotenv
from pathlib import Path

load_dotenv()
EMBEDDINGS_MODEL = os.getenv('EMBEDDINGS_MODEL', 'intfloat/e5-small-v2')
GROQ_MODEL_CUPS = os.getenv("GROQ_MODEL_CUPS", "openai/gpt-oss-120b")

class CUPSRetriever:
    def __init__(self, faiss_index_path, groq_api_key=None, groq_model=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)
        api_key = groq_api_key or os.getenv("GROQ_API_KEY")
        if not api_key:
            raise RuntimeError("Se requiere GROQ_API_KEY")
        self.groq = Groq(api_key=api_key)
        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
            logging.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):
        self._require_db()
        return self.db.max_marginal_relevance_search(consulta, k=k, lambda_param=0.7)

    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.'"
        )
        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="cups_faiss",
        groq_api_key=os.getenv("GROQ_API_KEY"),
        groq_model=os.getenv("GROQ_MODEL_CUPS", "openai/gpt-oss-120b"),
    )

    # Procedimientos de prueba
    procedimientos = [
        "PUNCION 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)"
       # "TOMOGRAFIA 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()