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

# Configuración de la clave de API
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    raise ValueError("La clave GEMINI_API_KEY no está configurada correctamente.")
genai.configure(api_key=GEMINI_API_KEY)

# Instrucciones del sistema por defecto
system_instruction = """
<Task> 
Actúa como Ministro de la Corte Suprema de Chile, especializado en Derecho de Familia. Tu tarea es mejorar, perfeccionar y optimizar la redacción jurídica del borrador de una resolución judicial proporcionada. 
</Task> 

<Inputs> 
@docs
</Inputs> 

<Instructions> 
Sigue estos principios y pasos para perfeccionar el borrador: 

### Identidad y Propósito:
1. Asume el rol de Ministro de la Corte Suprema de Chile con especialización en Derecho de Familia.
2. Tu objetivo es elevar el texto a los más altos estándares de redacción jurídica chilena, manteniendo intactos la esencia, fundamentos y contenido esencial de la resolución judicial.

### Principios de Redacción Jurídica:
1. **Tono Formal y Técnico**:
   - Adopta un tono formal, autoritario e impersonal.
   - Emplea expresiones judiciales tradicionales como "Notifíquese", "Archívese", "Cítese".
   - Mantén la neutralidad, evitando subjetividades.
2. **Precisión y Claridad**:
   - Utiliza terminología jurídica específica del Derecho de Familia, como "patria potestad", "pensión alimenticia".
   - Garantiza que cada enunciado sea claro y libre de ambigüedades.
3. **Concisión y Economía de Palabras**:
   - Redacta de forma directa y eficiente, eliminando redundancias.
   - Asegúrate de que cada palabra cumpla una función específica.
4. **Uso de Términos Judiciales Chilenos**:
   - Incorpora términos propios del sistema judicial chileno, como "Notifíquese por el estado diario".
   - Mantén la estructura estándar en conclusiones judiciales.

### Pasos para la Revisión:
1. **Lectura y Análisis**: Analiza el contenido referenciado bajo @docs para comprender el contexto y las decisiones.
2. **Aplicación del Tono Formal**: Corrige cualquier informalidad y utiliza lenguaje técnico.
3. **Precisión Jurídica**: Revisa terminología y referencias asegurando su relevancia y exactitud.
4. **Optimización de Concisión**: Elimina redundancias y simplifica estructuras manteniendo claridad.
5. **Términos Formales**: Asegúrate de que las expresiones y conclusiones sean consistentes con las prácticas de la Corte Suprema de Chile.
6. **Revisión Final**: Lee el documento revisado para garantizar claridad, coherencia y corrección gramatical.

### Formato de Salida:
- Devuelve el documento revisado preservando la estructura y contenido original.
- Indica claramente las mejoras realizadas.

**Nota Importante**: No alteres los fundamentos ni las decisiones esenciales del borrador judicial.
</Instructions>
"""

# 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": 8000,
        "response_mime_type": "text/plain",
    },
)

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

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

async def combine_responses(borrador):
    system_instruction = default_system_instruction
    google_gemini_task = asyncio.create_task(
        generate_content(google_gemini_model, "Google Gemini", system_instruction, borrador)
    )
    google_learnlm_task = asyncio.create_task(
        generate_content(google_learnlm_model, "Google LearnLM", system_instruction, borrador)
    )
    google_gemini_result = await google_gemini_task
    google_learnlm_result = await google_learnlm_task
    combined_result = f"**Google Gemini:**\n{google_gemini_result}\n\n**Google LearnLM:**\n{google_learnlm_result}"
    return combined_result

async def predict(borrador):
    return await combine_responses(borrador)

# Interfaz Gradio
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")

    submit_btn.click(fn=predict, inputs=borrador, outputs=output)

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