# ========================================================================================================================================
#                                               PROYECTO DE ATENCION MÉDICA PERSONALIZADA 
#                                                         Bogotá Colombia 2024
#
#                                          --- ADAPTADO PARA USAR GOOGLE GEMINI API ---
# ========================================================================================================================================
import streamlit as st
import google.generativeai as genai
import os
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
st.set_page_config(
    page_title="Consulta Médica Colombiana IA",
    page_icon="🩺",
    layout="wide"
)
st.title("🩺 Consulta Médica Colombiana IA")
st.caption("Desarrollado con Google Gemini")


# --- CONFIGURACIÓN Gemini ---
# Carga la API key desde los secretos de Streamlit o variables de entorno
GEMINI_API_KEY = st.secrets.get("GEMINI_API_KEY", os.getenv("GEMINI_API_KEY"))

if not GEMINI_API_KEY:
    st.error("GEMINI_API_KEY no encontrada. Por favor configúrala en tus secretos de Streamlit o variables de entorno.")
    logger.error("GEMINI_API_KEY no encontrada.")
    st.stop()

try:
    genai.configure(api_key=GEMINI_API_KEY)
    logger.info("✅ Configuración de Gemini API realizada.")
except Exception as e:
    error_msg = f"❌ Error al configurar Gemini API: {str(e)}"
    logger.error(error_msg)
    st.error(error_msg)
    st.stop()

# --- CARGA DEL MODELO ---
@st.cache_resource
def get_gemini_model():
    """Carga y cachea el modelo generativo de Gemini."""
    logger.info("🔄 Cargando modelo Gemini...")
    # NOTA: "gemini-2.5-flash-lite" puede ser un nombre de modelo no disponible públicamente aún.
    # Un modelo rápido y común es "gemini-1.5-flash-latest". 
    # Si el nombre que pusiste da error, prueba a cambiarlo por "gemini-1.5-flash-latest".
    try:
        model = genai.GenerativeModel("gemini-2.5-flash-lite") # Cambiado a un modelo conocido. Usa "gemini-2.5-flash-lite" si estás seguro de su disponibilidad.
        logger.info("✅ Modelo Gemini cargado correctamente.")
        return model
    except Exception as e:
        error_msg = f"❌ Error al inicializar modelo Gemini: {str(e)}"
        logger.error(error_msg)
        st.error(error_msg)
        return None

model = get_gemini_model()
if model is None:
    st.stop()


# --- INICIALIZACIÓN DEL CHAT EN LA SESIÓN ---
# Se inicializa el historial para mostrar en la UI y el objeto de chat de Gemini
if "history" not in st.session_state:
    st.session_state.history = []
    
if "chat" not in st.session_state:
    # El historial del chat de Gemini se inicializa con un prompt de sistema
    # para darle a la IA su rol y sus instrucciones.
    system_prompt = (
        "Eres un asistente médico virtual diseñado para Colombia. Tu principal objetivo es "
        "proporcionar orientación general y educativa sobre temas de salud, siempre en español. "
        "puedes proporcionar diagnósticos médicos, pero solo recomedar examenes o tratamientos "
        "ejemplo tus sintomas pueden corresponde a un dengue clasico "
        "Tu función es informativa. Debes adaptar la complejidad de tu lenguaje al nivel de estudio "
        "del usuario y considerar su ubicación (departamento) si es relevante. Siempre, al final de "
        "tu respuesta, incluye una advertencia clara: 'Recuerda que esta es una IA y no reemplaza "
        "la consulta con un médico real. En caso de emergencia, contacta a un profesional de la salud.'"
    )
    st.session_state.chat = model.start_chat(history=[
        {'role': 'user', 'parts': [system_prompt]},
        {'role': 'model', 'parts': ["Entendido. Estoy listo para ayudar como asistente médico virtual informativo para Colombia."]}
    ])

# --- INTERFAZ DE USUARIO ---
with st.sidebar:
    st.header("Información del Paciente")
    nivel_estudio = st.selectbox("Nivel de estudio:", ["Primaria", "Secundaria", "Técnico", "Pregrado", "Posgrado"])
    departamento = st.selectbox("Departamento de Colombia:", ["Bogotá DC","Amazonas", "Antioquia", "Arauca", "Atlántico", "Bolívar", "Boyacá", "Caldas", "Caquetá", "Casanare", "Cauca", "Cesar", "Chocó", "Córdoba", "Cundinamarca", "Guainía", "Guaviare", "Huila", "La Guajira", "Magdalena", "Meta", "Nariño", "Norte de Santander", "Putumayo", "Quindío", "Risaralda", "San Andrés y Providencia", "Santander", "Sucre", "Tolima", "Valle del Cauca", "Vaupés", "Vichada"])
    
    st.warning("Esta IA no reemplaza a un profesional médico.")
    
    if st.button("Limpiar Conversación"):
        st.session_state.history = []
        st.session_state.chat = model.start_chat(history=[]) # Reinicia el chat de Gemini
        st.rerun()

for message in st.session_state.history:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Capturar la entrada del usuario
if prompt := st.chat_input("Describe tus síntomas o consulta médica... (edad, sexo, duración, antecedentes, etc.)"):
    # Construir el prompt completo para enviar al modelo
    full_prompt = (
        f"Contexto del usuario: Nivel de estudio '{nivel_estudio}', vive en el departamento de '{departamento}'.\n\n"
        f"Consulta del usuario: {prompt}"
    )

    # Añadir y mostrar el mensaje del usuario en la UI
    st.session_state.history.append({"role": "user", "content": full_prompt})
    with st.chat_message("user"):
        st.markdown(prompt) # Mostramos solo la consulta, no el contexto interno

    # Generar y mostrar la respuesta de la IA
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        try:
            # Enviar el prompt al chat de Gemini y obtener la respuesta en streaming
            responses = st.session_state.chat.send_message(
                full_prompt,
                stream=True,
                generation_config=genai.types.GenerationConfig(
                    # Puedes ajustar estos parámetros
                    temperature=0.7, 
                    max_output_tokens=1024
                )
            )
            
            for response in responses:
                # Quita los mensajes de seguridad del stream si los hubiera
                if response.parts:
                    full_response += response.text
                    message_placeholder.markdown(full_response + "▌") # Muestra un cursor parpadeante
            
            message_placeholder.markdown(full_response)
            logger.info(f"Respuesta generada para: {prompt[:50]}...")
            
        except Exception as e:
            full_response = f"Lo siento, ocurrió un error al procesar tu solicitud: {str(e)}"
            message_placeholder.error(full_response)
            logger.error(f"Error en la llamada a Gemini API: {str(e)}")

    # Añadir la respuesta completa de la IA al historial de la UI
    st.session_state.history.append({"role": "assistant", "content": full_response})