import gradio as gr
import os
import json
import logging
from groq import Groq
from langchain_groq import ChatGroq

# Cargar variables de entorno desde archivo .env


# Configuración de logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Pesos para el cálculo del RAID
RAID_WEIGHTS = {
    "Dolor": 0.21,
    "Discapacidad Funcional": 0.16,
    "Fatiga": 0.15,
    "Sueño": 0.12,
    "Bienestar Físico": 0.12,
    "Bienestar Emocional": 0.12,
    "Afronte": 0.12
}

# Plantilla de prompt para evaluación RAID
RAID_PROMPT_TEMPLATE = """
Analiza la siguiente conversación o texto del paciente con artritis reumatoide.

Si el texto está vacío o no contiene información relevante, responde únicamente con el siguiente mensaje:
    "No hay información suficiente para calcular el RAID."

En caso de haber información, extrae y asigna una puntuación del 0 al 10 para cada una de las siguientes categorías, basándote en los detalles proporcionados:

1. Dolor:
   - Asigne el número que mejor describa el dolor que sintió debido a su Artritis durante la última semana.
   - Escala: 0 = Ninguno, 10 = Extremo.
2. Discapacidad Funcional:
   - Asigne el número que mejor describa la dificultad que tuvo para realizar actividades diarias debido a su Artritis Reumatoidea durante la última semana.
   - Escala: 0 = Sin dificultad, 10 = Dificultad extrema.
3. Fatiga:
   - Asigne el número que mejor describa la fatiga que experimentó debido a su Artritis Reumatoidea durante la última semana.
   - Escala: 0 = Sin fatiga, 10 = Totalmente exhausto.
4. Sueño:
   - Asigne el número que mejor describa las dificultades para dormir que experimentó durante la última semana.
   - Escala: 0 = Sin dificultad, 10 = Dificultad extrema.
5. Bienestar Físico:
   - Asigne el número que mejor describa su nivel de bienestar físico en relación a su Artritis Reumatoidea durante la última semana.
   - Escala: 0 = Muy bueno, 10 = Muy malo.
6. Bienestar Emocional:
   - Asigne el número que mejor describa su nivel de bienestar emocional en relación a su Artritis Reumatoidea durante la última semana.
   - Escala: 0 = Muy bueno, 10 = Muy malo.
7. Afronte:
   - Asigne el número que mejor describa cómo enfrentó o se adaptó a su enfermedad durante la última semana.
   - Escala: 0 = Muy bien, 10 = Muy mal.

Utiliza los siguientes pesos para calcular el valor final del RAID:
{weights}

La fórmula de cálculo es:
    RAID final = (Dolor × 0.21) + (Discapacidad Funcional × 0.16) + (Fatiga × 0.15) + (Bienestar Físico × 0.12) + (Sueño × 0.12) + (Bienestar Emocional × 0.12) + (Afronte × 0.12)

El valor final estará en el rango de 0 a 10, donde valores más altos indican un peor estado.

Por favor, responde únicamente con un JSON válido que incluya:
    - "raid_scores": un objeto que contenga las puntuaciones asignadas para cada categoría.
    - "raid_total": el valor final calculado.
Si falta información para alguna categoría, asigna 0 a esa categoría y añade una clave "nota" con un mensaje indicando que la información es parcial.

Ejemplo de respuesta JSON:
{{
    "raid_scores": {{
        "Dolor": 8,
        "Discapacidad Funcional": 6,
        "Fatiga": 7,
        "Sueño": 5,
        "Bienestar Físico": 6,
        "Bienestar Emocional": 7,
        "Afronte": 4
    }},
    "raid_total": 6.45
}}

Texto del paciente:
{patient_text}
"""

# Inicialización del cliente Groq
# Inicialización del cliente
api_key = os.environ.get("GROQ_API_KEY")
if not api_key:
    raise ValueError("GROQ_API_KEY no está configurada en las variables de entorno.")
client = Groq(api_key=api_key)
model_name = "llama-3.3-70b-versatile"
chat_groq = ChatGroq(model=model_name)


def transcribe_audio(audio_filepath):
    """Transcribe el audio usando Whisper de Groq"""
    if not audio_filepath:
        return ""
    try:
        with open(audio_filepath, "rb") as file:
            transcription = client.audio.transcriptions.create(
                file=(audio_filepath, file.read()),
                model="whisper-large-v3",
                response_format="json",
                temperature=0.0
            )
        logger.info("Transcripción de audio exitosa.")
        return transcription.text
    except Exception as e:
        logger.error("Error en transcripción de audio: %s", e, exc_info=True)
        return ""

def evaluate_raid(patient_text):
    """Evalúa el RAID score basado en el texto del paciente"""
    try:
        # Limitar texto para contexto
        if len(patient_text) > 4000:
            logger.warning("Texto demasiado largo, se truncará a 4000 caracteres")
            patient_text = patient_text[:4000]
            
        prompt = RAID_PROMPT_TEMPLATE.format(
            patient_text=patient_text,
            weights=json.dumps(RAID_WEIGHTS, ensure_ascii=False)
        )
        logger.info("Evaluando RAID con prompt")
        
        response = chat_groq.invoke(prompt)
        content = response.content
        logger.info("Respuesta de chat_groq recibida")
        
        # Extraer JSON de la respuesta
        start = content.find('{')
        end = content.rfind('}') + 1
        json_str_candidate = content[start:end]
        try:
            json_response = json.loads(json_str_candidate)
        except json.JSONDecodeError:
            # Intentar extraer contenido entre marcadores ```json
            try:
                json_str_candidate = content.split('```json')[1].split('```')[0].strip()
                json_response = json.loads(json_str_candidate)
            except Exception as e:
                logger.error("Error al extraer JSON usando marcadores: %s", e)
                json_response = {}
        
        if not isinstance(json_response, dict):
            logger.warning("Respuesta JSON no es un diccionario. Se usará un diccionario vacío.")
            json_response = {}
        
        # Obtener o inicializar raid_scores
        raid_scores = json_response.get('raid_scores', {})
        if not isinstance(raid_scores, dict):
            logger.warning("El campo 'raid_scores' no es un diccionario, se establecerá como vacío.")
            raid_scores = {}
        
        # Completar las categorías faltantes con 0
        for category in RAID_WEIGHTS:
            if category not in raid_scores:
                logger.warning("Categoría '%s' no encontrada en la respuesta. Se asignará 0.", category)
                raid_scores[category] = 0
        
        # Calcular el puntaje total usando las categorías disponibles
        calculated_total = sum(raid_scores[k] * RAID_WEIGHTS[k] for k in RAID_WEIGHTS)
        raid_total = round(calculated_total, 2)
        
        # Actualizar json_response
        json_response['raid_scores'] = raid_scores
        json_response['raid_total'] = raid_total
            
        return json_response
    
    except Exception as e:
        logger.error("Error en evaluación RAID: %s", e, exc_info=True)
        # Retorna resultados parciales con mensaje de error
        partial_result = {
            "raid_scores": {k: 0 for k in RAID_WEIGHTS},
            "raid_total": 0,
            "error": str(e)
        }
        return partial_result

def format_raid_results(raid_data):
    """Formatea los resultados del RAID para visualización"""
    if not raid_data:
        return "No se pudo calcular el RAID score"
    
    scores = raid_data.get('raid_scores', {})
    total = raid_data.get('raid_total', 0)
    
    breakdown_lines = []
    for category, weight in RAID_WEIGHTS.items():
        score = scores.get(category, 0)
        breakdown_lines.append(f"- {category}: {score} (Peso: {weight*100}%)")
    breakdown = "\n".join(breakdown_lines)
    
    interpretation = "Interpretación del puntaje:\n"
    if total >= 7:
        interpretation += "RAID alto: Impacto significativo en la calidad de vida. Considerar ajuste terapéutico."
    elif total >= 4:
        interpretation += "RAID moderado: Impacto notable. Monitoreo cercano recomendado."
    else:
        interpretation += "RAID bajo: Buen control sintomático. Mantener seguimiento."
    
    result = f"""
**Puntaje Total RAID: {total:.2f}/10**

**Desglose:**
{breakdown}

{interpretation}
"""
    if "error" in raid_data:
        result += f"\n\n**Nota:** Se produjo un error durante el cálculo del RAID: {raid_data['error']}"
    return result

def process_input(audio, text_input):
    """Procesa la entrada de audio y texto para calcular el RAID"""
    patient_text = ""
    
    # Procesar audio si está disponible
    if audio:
        transcription = transcribe_audio(audio)
        if transcription:
            patient_text += transcription + "\n\n"
    
    # Añadir texto manual si está disponible
    if text_input:
        patient_text += text_input
    
    # Si no hay texto para analizar, retornar mensaje
    if not patient_text.strip():
        return "No se proporcionó texto ni audio para analizar. Por favor, añada información del paciente."
    
    # Evaluar RAID con el texto disponible
    raid_results = evaluate_raid(patient_text)
    formatted_results = format_raid_results(raid_results)
    
    return formatted_results

# Configuración del tema
theme = gr.themes.Base(
    primary_hue=gr.themes.Color(
        c100="#fce7f3", c200="#fbcfe8", c300="#f9a8d4",
        c400="#f472b6", c50="#fdf2f8", c500="#ff4da6",
        c600="#db2777", c700="#ff299b", c800="#f745b6",
        c900="#ff38ac", c950="#e1377e"
    ),
    secondary_hue="pink",
    neutral_hue="neutral",
)

# Interfaz de Gradio simple para el cálculo de RAID
with gr.Blocks(theme=theme) as iface:
    gr.Markdown("# Calculadora RAID para Artritis Reumatoide")
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(
                sources=["microphone", "upload"], 
                type="filepath", 
                label="Audio del paciente"
            )
            text_input = gr.Textbox(
                label="Texto del paciente (opcional)",
                placeholder="Ingrese aquí cualquier información adicional del paciente...",
                lines=5
            )
            submit_btn = gr.Button("Calcular RAID", variant="primary")
        
        with gr.Column():
            raid_output = gr.Markdown(
                "### Resultados RAID\n\nEspere mientras se procesa la información...",
            )
    
    submit_btn.click(
        fn=process_input,
        inputs=[audio_input, text_input],
        outputs=raid_output
    )
    
    # También calcular automáticamente cuando se cambia el audio
    audio_input.change(
        fn=process_input,
        inputs=[audio_input, text_input],
        outputs=raid_output
    )

if __name__ == "__main__":
    iface.launch()