import os
import asyncio
import google.generativeai as genai
import gradio as gr

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

# Instrucciones del sistema por defecto
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 estética lingüística jurídica.**
"""

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

# Configuración del modelo de Google LearnLM
google_learnlm_model = genai.GenerativeModel(
    "learnlm-1.5-pro-experimental",
    generation_config={
        "temperature": 0.5,
        "top_p": 0.9,
        "top_k": 40,
        "max_output_tokens": 5000,
        "response_mime_type": "text/plain",
    },
)

# Función genérica para generar contenido
async def generate_content(client, model_name, system_instruction, borrador):
    try:
        response = await asyncio.to_thread(client.generate_content, [system_instruction, borrador])
        return response.text
    except Exception as e:
        return f"Error en {model_name}: {str(e)}"

# Combina las respuestas de ambos modelos de Google
async def combine_responses(borrador):
    system_instruction = default_system_instruction

    # Generar contenido con Google Gemini
    google_gemini_result = await generate_content(google_gemini_model, "Google Gemini", system_instruction, borrador)

    # Generar contenido con Google LearnLM
    google_learnlm_result = await generate_content(google_learnlm_model, "Google LearnLM", system_instruction, borrador)

    # Combinar resultados
    combined_result = f"**Google Gemini:**\n{google_gemini_result}\n\n**Google LearnLM:**\n{google_learnlm_result}"
    return combined_result

# Función de predicción
async def predict(borrador):
    result = await combine_responses(borrador)
    return result

# Interfaz Gradio con botón funcional
with gr.Blocks() as demo:
    gr.Markdown("### Mejorador de resoluciones judiciales - Derecho de Familia en Chile")

    borrador = gr.Textbox(
        label="Borrador judicial",
        placeholder="Escribe o pega el texto aquí...",
        lines=10
    )
    output = gr.Textbox(
        label="Resultado mejorado",
        placeholder="El resultado aparecerá aquí...",
        lines=10
    )
    submit_btn = gr.Button("Enviar")  # Botón funcional para enviar el texto

    submit_btn.click(
        fn=lambda texto: asyncio.run(predict(texto)),
        inputs=borrador,
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()