import os
import asyncio
import google.generativeai as genai
from openai import OpenAI, APIError as OpenAIError
import gradio as gr

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

# Instrucciones del sistema por defecto (pueden ser modificadas por el usuario)
default_system_instruction = """
Actúa como un Ministro de la Corte Suprema de Chile, específicamente de la Primera Sala especializada en Derecho de Familia, para revisar y perfeccionar un borrador de resolución judicial. El objetivo es mejorar la gramática, redacción y estilo jurídico, manteniendo la estructura y contenido original.
Como Ministro de la Corte Suprema de Chile, especializado en Derecho de Familia, y con una vasta y reconocida trayectoria en la judicatura, tu misión es revisar y perfeccionar un borrador de resolución judicial, aplicando un tono formal, autoritario e impersonal, empleando el modo imperativo en expresiones de orden y mandato judicial. Proyectarás neutralidad y profesionalismo en todo el texto, evitando cualquier tono subjetivo o interpretativo. Usarás terminología jurídica específica del Derecho de Familia, asegurando que cada enunciado sea claro y preciso. Emplearás una redacción directa, evitando redundancias y expresiones superfluas. Utilizarás expresiones tradicionales y formales del ámbito judicial chileno.
El objetivo es elevar el texto a un estándar de excelencia en redacción jurídica, asegurando la máxima claridad, precisión, concisión y formalidad. **No debes modificar la estructura del borrador, tampoco agregar fundamentación o hechos. La mejora solo es gramatical, redaccional y estitica lenguistica juridica.**
"""

# Configuración del modelo de Google
google_model = genai.GenerativeModel(
    "gemini-exp-1114",
    generation_config={
        "temperature": 0.5,
        "top_p": 0.9,
        "top_k": 40,
        "max_output_tokens": 4000,
        "response_mime_type": "text/plain",
    },
)

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

# Función genérica para generar contenido
async def generate_content(client, model_name, system_instruction, borrador, is_nvidia=False):
    try:
        if is_nvidia:
            messages = [
                {"role": "system", "content": system_instruction},
                {"role": "user", "content": borrador}
            ]
            completion = await asyncio.to_thread(
                client.chat.completions.create,
                 model="meta/llama-3.1-405b-instruct",
                messages=messages,
                temperature=0.5,
                top_p=0.7,
                max_tokens=4000
            )
            return completion.choices[0].message.content
        else:  # Suponemos que es Google Gemini
            response = await asyncio.to_thread(client.generate_content, [system_instruction, borrador])
            return response.text
    except genai.types.GenerativeContentError as e:
        return f"Error en Google Gemini (Contenido): {str(e)}"
    except OpenAIError as e:
        return f"Error en NVIDIA API: {str(e.message)}"
    except Exception as e:
        return f"Error en {model_name}: {str(e)}"

# Función para combinar las respuestas de ambos modelos
async def combine_responses(borrador):
    system_instruction = default_system_instruction

    # Generar contenido con Google Gemini
    google_result = await generate_content(google_model, "gemini-exp-1114", system_instruction, borrador, is_nvidia=False)

    # Generar contenido con NVIDIA
    nvidia_result = await generate_content(nvidia_client, "gpt-4", system_instruction, borrador, is_nvidia=True)

    # Combinar las respuestas
    combined_result = f"**Google Gemini:**\n{google_result}\n\n**NVIDIA:**\n{nvidia_result}"
    return combined_result

# Función predict para procesar la entrada del usuario
async def predict(borrador):
    """
    borrador: str - El texto del borrador judicial que se debe procesar.
    """
    # Llama a la función para combinar las respuestas
    result = await combine_responses(borrador)
    return result

# Interfaz de usuario con Gradio
with gr.Blocks() as demo:
    gr.Markdown("### Mejorador de Resoluciones Judiciales para Derecho de Familia - Chile")

    # Campo de entrada para el borrador judicial
    borrador = gr.Textbox(
        label="Ingresa el borrador de la resolución judicial",
        placeholder="Escribe o pega el texto aquí...",
        lines=10
    )

    # Botón de procesamiento
    submit_btn = gr.Button("Procesar texto")

    # Campo de salida para el resultado procesado
    output = gr.Textbox(
        label="Resultado mejorado",
        placeholder="El resultado aparecerá aquí...",
        lines=10
    )

    # Conexión del botón con la función predict
    submit_btn.click(
        fn=lambda borrador: asyncio.run(predict(borrador)),
        inputs=borrador,
        outputs=output
    )

# Iniciar el servidor Gradio
if __name__ == "__main__":
    demo.launch()