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")

# 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=model_name,
                messages=messages,
                temperature=0.5,
                top_p=0.7,
                max_tokens=4000
            )
            return completion.choices[0].message.content
        else:  # Suponemos que es Google Gemini
            client.system_instruction = system_instruction
            response = await asyncio.to_thread(client.generate_content, 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 predict para procesar la entrada del usuario
async def predict(api_choice, borrador):
    """
    api_choice: str - Define si se usará 'NVIDIA' o 'Google Gemini'.
    borrador: str - El texto del borrador judicial que se debe procesar.
    """
    system_instruction = default_system_instruction
    model_name = "gemini-exp-1114" if api_choice == "Google Gemini" else "gpt-4"

    # Verifica si se selecciona NVIDIA o Google Gemini
    if api_choice == "NVIDIA":
        client = nvidia_client
        is_nvidia = True
    elif api_choice == "Google Gemini":
        client = google_model
        is_nvidia = False
    else:
        return "Error: Debe seleccionar una API válida ('NVIDIA' o 'Google Gemini')."

    # Llama a la función genérica de generación
    result = await generate_content(client, model_name, system_instruction, borrador, is_nvidia)
    return result

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

    # Selección de API (NVIDIA o Google Gemini)
    api_choice = gr.Radio(
        choices=["NVIDIA", "Google Gemini"],
        label="Selecciona el modelo API",
        value="Google Gemini"
    )

    # 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 api_choice, borrador: asyncio.run(predict(api_choice, borrador)),
        inputs=[api_choice, borrador],
        outputs=output
    )

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