import pandas as pd
import json
import re
import logging
from unidecode import unidecode
from langchain_community.vectorstores import FAISS
from groq import Groq
import os
from dotenv import load_dotenv
from app.config import config, require_env
from modules.processing.embeddings import get_hf_embeddings

load_dotenv()
EMBEDDINGS_MODEL = os.getenv('EMBEDDINGS_MODEL', 'intfloat/e5-small-v2')

class CIE10Retriever:
    def __init__(self, faiss_index_path, estructura_json_path):
        self.embeddings = get_hf_embeddings(EMBEDDINGS_MODEL)
                                                #"intfloat/multilingual-e5-large-instruct")#"sentence-transformers/all-MiniLM-L6-v2")
        self.db = self._cargar_indice(faiss_index_path)
        self.estructura = self._cargar_estructura(estructura_json_path)
        self.normalizador = self._crear_normalizador()
    
    def _cargar_indice(self, path):
        try:
            return FAISS.load_local(path, self.embeddings, allow_dangerous_deserialization=True)
        except Exception as e:
            raise RuntimeError(f"Error cargando índice FAISS: {str(e)}")
    
    def _cargar_estructura(self, path):
        with open(path, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    def _crear_normalizador(self):
        return {
            'sinonimos': {
                r"\b(trazo\s+sugestivo\b|fx|fract)": "fractura",
                r"\b(dcha|der)\b": "derecho",
                r"\b(izq|izqda)\b": "izquierdo",
                r"\b(tac)\b": "tomografía axial computarizada",
                r"\b(5to|quinto)\b": "5"
            },
            'reemplazos': {
                "QUEMADURA POR FRICCION": "QUEMADURA POR FRICCIÓN",
                "TRAZO SUGESTIVO DE FRACTURA": "FRACTURA"
            }
        }
    
    def normalizar_consulta(self, texto):
        texto = texto.upper()
        for patron, reemplazo in self.normalizador['reemplazos'].items():
            texto = texto.replace(patron, reemplazo)
        for patron, reemplazo in self.normalizador['sinonimos'].items():
            texto = re.sub(patron, reemplazo, texto, flags=re.IGNORECASE)
        return unidecode(texto).strip()
    
    def buscar(self, consulta, k=50, lambda_param=0.7, nivel_minimo=3):
        consulta_norm = self.normalizar_consulta(consulta)
        
        resultados = self.db.max_marginal_relevance_search(
            consulta,
            k=k,
            lambda_param=lambda_param,
            filter={"nivel": {"$gte": nivel_minimo}}
        )
        
        return self._procesar_resultados(resultados)
    
    def _procesar_resultados(self, resultados):
        jerarquia = {
            'bloques': [],
            'categorias': [],
            'subcategorias': [],
            'detalles': []
        }
        
        for doc in resultados:
            metadata = doc.metadata
            nivel = metadata['nivel']
            entrada = {
                'codigo': metadata['codigo'],
                'descripcion': doc.page_content.split(":")[-1].strip(),
                'nivel': nivel,
                'ruta': metadata.get('ruta', ''),
                'score': self._calcular_score(nivel)
            }
            
            if nivel == 1:
                jerarquia['bloques'].append(entrada)
            elif nivel == 2:
                jerarquia['categorias'].append(entrada)
            else:
                jerarquia['subcategorias'].append(entrada)
            
            jerarquia['detalles'].append(entrada)
        
        return jerarquia
    
    def _calcular_score(self, nivel):
        # Prioriza subcategorías > categorías > bloques
        return {1: 0.3, 2: 0.6, 3: 1.0}.get(nivel, 0)
    
    def obtener_ruta_completa(self, codigo):
        for bloque in self.estructura['bloques']:
            if codigo == bloque['codigo']:
                return [bloque]
            for categoria in bloque['categorias']:
                if codigo == categoria['codigo']:
                    return [bloque, categoria]
                for subcat in categoria['subcategorias']:
                    if codigo == subcat['codigo']:
                        return [bloque, categoria, subcat]
        return []
    
    def validar_codigo(self, codigo):
        return any(self.obtener_ruta_completa(codigo))
    
    def mejores_resultados(self, resultados, top_n=3):
        return sorted(
            resultados['detalles'],
            key=lambda x: x['score'],
            reverse=True
        )[:top_n]


# Configurar cliente Groq
client = Groq(api_key=require_env(os.getenv("GROQ_API_KEY") or config.GROQ_API_KEY, "GROQ_API_KEY"))

# Función para generar respuesta utilizando el LLM de Groq
def generar_respuesta_groq(consulta, contexto, temperature=0.0):
    """
    Genera una respuesta utilizando el modelo LLM de Groq, basándose en la consulta y el contexto recuperado.
    """
    try:
        # Prompt refinado para asignar códigos CIE-10
        prompt = (
            "Eres un asistente médico experto en codificación CIE-10. "
            "Tu tarea es asignar el código CIE-10 correcto basándote en la consulta del usuario y en el contexto proporcionado.\n\n"
            "El contexto consiste en fragmentos de un corpus que relaciona códigos y descripciones médicas.\n"
            "Responde únicamente en el siguiente formato:\n\n"
            "Código: <código> - Descripción: <descripción>\n\n"
            "Si no puedes determinar un código con certeza, responde exactamente: 'No se pudo determinar el código CIE-10.'\n\n"
            f"Consulta: {consulta}\n\n"
            f"Contexto:\n{contexto}"
        )
        
        # Generar respuesta con Groq
        response = client.chat.completions.create(
            model=os.getenv("GROQ_MODEL_CIE10", "openai/gpt-oss-120b"),  # Modelo adecuado para consultas médicas
            temperature=temperature,
            top_p=0.1,
            messages=[
                {"role": "system", "content": "Eres un asistente médico experto."},
                {"role": "user", "content": prompt}
            ]
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"Error al generar la respuesta: {str(e)}"

def asignar_codigo_cie10(consulta, retriever, temperature=0.0, k=50, lambda_param=0.8):
    """
    Función que asigna el código CIE-10 basándose en una consulta.
    
    Integra la búsqueda de contexto utilizando la clase CIE10Retriever y la generación de respuesta
    con el LLM de Groq.
    
    Parámetros:
      - consulta: Consulta en texto sobre el diagnóstico o trauma.
      - retriever: Instancia de CIE10Retriever para recuperar el contexto.
      - temperature: Parámetro de temperatura para la generación de respuesta.
      - k: Número de documentos a recuperar.
      - lambda_param: Parámetro lambda para la búsqueda con MMR.
      
    Retorna:
      - La respuesta generada que indica el código CIE-10 asignado.
    """
    try:
        # Recuperar contexto relevante usando el método buscar del retriever
        resultados = retriever.buscar(consulta, k=k, lambda_param=lambda_param)
        
        # Construir un string de contexto a partir de los detalles obtenidos
        contexto = ""
        for detalle in resultados['detalles']:
            contexto += f"Código: {detalle['codigo']} - Descripción: {detalle['descripcion']}\n"
        
        print("Contexto recuperado:\n", contexto)
        
        # Generar la respuesta utilizando Groq
        respuesta = generar_respuesta_groq(consulta, contexto, temperature=temperature)
        return respuesta
    except Exception as e:
        return f"Error en la asignación del código CIE-10: {str(e)}"

# Uso del sistema
if __name__ == "__main__":
    # Instanciar el retriever con los índices y la estructura adecuados
    retriever = CIE10Retriever(
        faiss_index_path="faiss_principal_jerarquico",
        estructura_json_path="cie10_estructura_completa.json"
    )
    
    # Ejemplo de consulta
    consulta ="Tendinitis del tendón del músculo supraespinoso"
        #"TRAZO SUGESTIVO DE FRACTURA DE TUBEROSIDAD MAYOR DE HUMERO DERECHO"
        #"TRAZO SUGESTIVO DE FRACTURA DE CUPILA RADIAL DE CODO DERECHO"
        #"TRAZO SUGESTIVO DE FRACTURA DE RADIO DISTAL DE MUÑECA DERECHO"
        #"FRACTURA DE 5TO METATARSIANO DE PIE IZQUIERDO"
    #"DESBRIDAMIENTO DE LESION DE TEJIDO PROFUNDO DE HERIDA EN PIE IZQUIERDO UÑA "
    #"Trazo sugestivo de fractura de radio distal de muñeca derecha" S525
    #"quemadura  por friccion de tercer grado del HOMBRO" T223
    #"Trazo sugestivo de fractura de tuberosidad mayor de húmero derecho" #S422 mejor precision con llama-3.3-70b-versatile
    #"TRAUMA PIE IZQUIERDO" S99 - Descripción: OTROS TRAUMATISMO Y LOS NO ESPECIFICADOS DEL TOBILLO Y DEL PIE
    #"Trazo sugestivo de fractura de cupula radial de codo derecho" S521 con k= 50  
    #"Fractura de 5to metatarsiano de pie izquierdo" S923 
    #"TRAUMA DEL CODO" S500
    
    # Asignar código CIE-10 basado en la consulta
    codigo_asignado = asignar_codigo_cie10(consulta, retriever, temperature=0.0)
    
    print("\nRespuesta generada:\n", codigo_asignado)