import gradio as gr
from google.adk.agents import Agent
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.adk.agents import LlmAgent
from google.genai import types
from typing import Dict, Any
from google.adk.models.lite_llm import LiteLlm
import uuid
import asyncio
from google import genai
import os

#download_files_from_drive_tool = FunctionTool(func=download_files_from_drive)
# --- Definición de Mod

# Definir función para el Agente Revisor

# Envolver funciones como herramientas de ADK
# Es importante que el nombre de la herramienta sea el mismo que el nombre de la función para el ADK
instruction = """Sos un auditor clínico experto. Que debes validar la pertinencia clínica de una orden médica. Identifica el estudio solicitado, el diagnóstico y la justificación clínica. Posteriormente analiza si el estudio solicitado está clínicamente justificado.

**IMPORTANTE**: Inpendientemente recuerda que la Guía Práctica Clínica dice: `Se recomienda que ante un paciente con Dolor lumbar sin respuesta al tratamiento después de 4-6 semanas de evolución considerar la utilidad de estudios de imagen. El tratamiento con paracetamol es insuficiente para justificar estudios de imagen.`


Formato de respuesta en bullets:

**Evaluación de pertinencia clínica**: [observaciones]

**Resultado**: [✅ Aprobado | ❌ Rechazado]

**Notas para el médico solicitante (Solo en el caso de rechazo)**: - [información]

Fuente: (Ministerio de Salud Pública. Dolor lumbar: Guía Práctica Clínica (GPC). Primera Edición. Quito: Dirección Nacional de Normatización; 2015)

IMPORTANTE_ Siempre debes citar la fuente."""
def consultar_guia_medica(consulta: str) -> str:
    client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))
    
    try:
        response = client.models.generate_content(
            model="gemini-2.0-flash",
            contents=f"""Responde la siguiente consulta en base a la guia de Dolor Lumbar y cita la fuente de la recomendación:

{consulta}
"""
        )
        print("RESPONSE", response)
        return response.candidates[0].content.parts[0].text
    except Exception as e:
        print(e)


APP_NAME = "predoc_app"



root_agent = LlmAgent(
    model=LiteLlm(model="openai/gpt-4.1"),
    generate_content_config=types.GenerateContentConfig(
        temperature=0.0,
    ),
    name="Agente_Auditor",
    instruction=instruction,
    description="Agente especialista en dolor lumbar",
    #tools=[consultar_guia_medica]
)

session_service = InMemorySessionService()

# Esta función se conecta a ChatInterface
def respond(message, history):
    # Detectar si es inicio de conversación (history vacío o con 0 mensajes)
    print("HISTORY",history)
    if history is None or len(history) == 0:
        user_id = str(uuid.uuid4())
        session_id = str(uuid.uuid4())
        print(f"🔄 Nueva sesión: {session_id}")

        async def create_session():
            await session_service.create_session(
                app_name=APP_NAME,
                user_id=user_id,
                session_id=session_id
            )
        asyncio.run(create_session())
        # Guardar en algún lado user_id y session_id para usar en las próximas llamadas
        # Por simplicidad acá lo guardamos en variables globales
        global CURRENT_USER_ID, CURRENT_SESSION_ID
        CURRENT_USER_ID = user_id
        CURRENT_SESSION_ID = session_id
    else:
        # Usar IDs existentes
        user_id = CURRENT_USER_ID
        session_id = CURRENT_SESSION_ID

    runner = Runner(agent=root_agent, app_name=APP_NAME, session_service=session_service)

    def call_agent_text(query):
        content = types.Content(role='user', parts=[types.Part(text=query)])
        events = runner.run(user_id=user_id, session_id=session_id, new_message=content)
        for event in events:
            if event.is_final_response():
                return event.content.parts[0].text
        return "No se obtuvo respuesta."

    def call_agent_image(query):
        images = []
        for q in query:
            with open(q, 'rb') as f:
                image_bytes = f.read()
            images.append(types.Part.from_bytes(data=image_bytes, mime_type='image/jpeg'))

        content = types.Content(role='user', parts=images)
        events = runner.run(user_id=user_id, session_id=session_id, new_message=content)
        for event in events:
            if event.is_final_response():
                return event.content.parts[0].text
        return "No se obtuvo respuesta."

    def call_agent_both(image, text):
        with open(image[0], 'rb') as f:
            image_bytes = f.read()
        content = types.Content(
            role='user',
            parts=[
                types.Part.from_bytes(data=image_bytes, mime_type='image/jpeg'),
                types.Part(text=text)
            ]
        )
        events = runner.run(user_id=user_id, session_id=session_id, new_message=content)
        for event in events:
            if event.is_final_response():
                return event.content.parts[0].text
        return "No se obtuvo respuesta."

    # Dispatcher
    if message['text'] != '' and len(message['files']) > 0:
        return call_agent_both(message['files'], message['text'])
    elif message['text'] == '' and len(message['files']) > 0:
        return call_agent_image(message['files'])
    elif message['text'] != '' and len(message['files']) == 0:
        return call_agent_text(message['text'])
    else:
        return "Escribe algo para que pueda contestarte."

# Inicializamos demo sin el argumento state
demo = gr.ChatInterface(fn=respond, title="Agente Auditor", multimodal=True)

demo.launch(debug=True)