import os
from openai import OpenAI
import gradio as gr
from gradio_agentchatbot import AgentChatbot, ChatMessage, ThoughtMetadata

# Configuración de claves de API
NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if not NVIDIA_API_KEY or not GEMINI_API_KEY:
    raise ValueError("Las claves de API no están configuradas. Asegúrate de establecer 'NVIDIA_API_KEY' y 'GEMINI_API_KEY'.")

# Configuración de clientes
nvidia_client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=NVIDIA_API_KEY,
)

gemini_client = OpenAI(
    api_key=GEMINI_API_KEY,
    base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
)

# Mensaje de sistema
system_message = {
    "role": "system",
    "content": """
Asume el rol de Ministro de la Corte Suprema de Chile especializado en Derecho de Familia. Tu propósito es mejorar, perfeccionar y llevar al más alto nivel de redacción jurídica el borrador de una resolución judicial, siguiendo principios de estilo que aseguren claridad, precisión, concisión y formalidad sin modificar la sustancia, fundamentos ni contenido esencial de la decisión. Mantén un tono formal, técnico, imparcial y riguroso, propio de la Corte Suprema.
[... más detalles como en el prompt proporcionado ...]
""".strip(),
}

# Función para interactuar con NVIDIA API
def nvidia_response(prompt: str, messages: list[ChatMessage]):
    if not prompt.strip():
        messages.append(ChatMessage(role="assistant", content="Por favor, ingresa un texto válido."))
        return messages

    # Agregar mensaje del usuario al historial
    metadata_user = ThoughtMetadata(tool_name="NVIDIA")
    messages.append(ChatMessage(role="user", content=prompt, thought_metadata=metadata_user))

    # Generar respuesta con NVIDIA API
    try:
        completion = nvidia_client.chat.completions.create(
            model="meta/llama-3.1-405b-instruct",
            messages=[system_message, {"role": "user", "content": prompt}],
            temperature=0.2,
            top_p=0.7,
            max_tokens=1024,
            stream=True,
        )
        response_content = ""
        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                response_content += chunk.choices[0].delta.content

        metadata_assistant = ThoughtMetadata(tool_name="NVIDIA")
        messages.append(ChatMessage(role="assistant", content=response_content, thought_metadata=metadata_assistant))
    except Exception as e:
        error_message = f"Error al procesar la solicitud con NVIDIA API: {e}"
        messages.append(ChatMessage(role="assistant", content=error_message))
    
    return messages

# Función para interactuar con Gemini API
def gemini_response(prompt: str, messages: list[ChatMessage]):
    if not prompt.strip():
        messages.append(ChatMessage(role="assistant", content="Por favor, ingresa un texto válido."))
        return messages

    # Agregar mensaje del usuario al historial
    metadata_user = ThoughtMetadata(tool_name="Gemini")
    messages.append(ChatMessage(role="user", content=prompt, thought_metadata=metadata_user))

    # Generar respuesta con Gemini API
    try:
        response = gemini_client.chat.completions.create(
            model="gemini-1.5-flash",
            messages=[system_message, {"role": "user", "content": prompt}],
            stream=True,
        )
        response_content = ""
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                response_content += chunk.choices[0].delta.content

        metadata_assistant = ThoughtMetadata(tool_name="Gemini")
        messages.append(ChatMessage(role="assistant", content=response_content, thought_metadata=metadata_assistant))
    except Exception as e:
        error_message = f"Error al procesar la solicitud con Gemini API: {e}"
        messages.append(ChatMessage(role="assistant", content=error_message))
    
    return messages

# Descripción de la aplicación
description = """
# Ministro IA de la Corte Suprema - Derecho de Familia
Este espacio utiliza NVIDIA y Gemini APIs para mejorar borradores de resoluciones judiciales en Derecho de Familia chileno.
"""

# Ejemplos
examples = [
    ["El juez decide que el papá puede ver a su hijo los fines de semana."],
    ["Se ordena pagarle una pensión a la mamá para los gastos del niño."],
]

# Interfaz de usuario con Gradio
with gr.Blocks() as demo:
    gr.Markdown(description)
    chatbot = AgentChatbot(label="Resolución Mejorada", avatar_images=("user.png", "bot.png"))
    text_input = gr.Textbox(
        lines=10,
        label="Borrador de la Resolución",
        placeholder="Ingrese el borrador aquí...",
    )

    # Botones y selector
    modelo_selector = gr.Radio(
        choices=["NVIDIA", "Gemini"],
        label="Selecciona el modelo",
        value="NVIDIA"
    )
    submit_button = gr.Button("Submit")
    clear_button = gr.Button("Limpiar")

    # Enlazar acciones a botones
    submit_button.click(
        lambda prompt, modelo, messages: nvidia_response(prompt, messages) if modelo == "NVIDIA" else gemini_response(prompt, messages),
        inputs=[text_input, modelo_selector, chatbot],
        outputs=[chatbot],
    )
    clear_button.click(lambda: ("", []), inputs=None, outputs=[text_input, chatbot])

    # Ejemplos
    gr.Examples(
        examples=examples,
        inputs=[text_input],
    )

# Ejecutar la aplicación
if __name__ == "__main__":
    demo.launch()