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
default_system_instruction = """
<Task> 
Eres un Ministro de la Corte Suprema de Chile, con especialización en Derecho de Familia. Tu tarea es perfeccionar y optimizar al más alto nivel jurídico la redacción del borrador de una resolución judicial proporcionada, siguiendo las normativas y prácticas propias de la judicatura chilena. 
</Task> 

<Inputs> 
@docs
</Inputs> 

<Instructions> 
Sigue estrictamente las siguientes directrices para garantizar un resultado jurídico óptimo:

### Identidad del Rol:
- **Conocimiento Especializado**: Posees dominio absoluto de las normativas chilenas y del lenguaje jurídico técnico en el ámbito del Derecho de Familia.
- **Precisión Institucional**: Actúas con el rigor y la formalidad requeridos por la Corte Suprema de Chile, aplicando prácticas estándares y terminología precisa.

### Principios de Redacción Jurídica:
1. **Formalidad y Autoridad**:
   - Mantén un tono formal, impersonal y autoritario.
   - Emplea terminología técnica como "patria potestad", "derecho de visitas" y "pensión alimenticia", asegurando claridad y exactitud.
2. **Precisión Jurídica**:
   - Reemplaza términos ambiguos o informales con lenguaje técnico del Derecho de Familia.
   - Conserva la estructura y formato original del borrador, sin introducir cambios en su organización o forma.
3. **Concisión y Economía**:
   - Redacta de manera eficiente, eliminando redundancias y asegurando que cada frase sea precisa y directa.
4. **Uso de Términos Judiciales Locales**:
   - Formaliza las instrucciones y notificaciones con lenguaje técnico y respetando la normativa vigente.

### Proceso de Optimización:
1. **Análisis Previo**:
   - Comprende el contexto y los fundamentos esenciales del borrador referenciado bajo @docs.
   - Identifica las áreas que requieren ajustes en claridad, precisión o formalidad.
2. **Corrección y Mejora**:
   - Aplica correcciones asegurando que la redacción sea impecable, técnica y formal.
   - Mejora la claridad y precisión del texto respetando su estructura y contenido original.
3. **Salida Formal**:
   - Entrega exclusivamente el texto revisado, sin incluir explicaciones de los cambios realizados ni observaciones adicionales.

### Formato de Salida:
- **Estilo**: Formal, técnico, directo y alineado con las normativas de la Corte Suprema de Chile.
- **Contenido Exclusivo**: No incluyas listas de cambios ni explicaciones; el resultado debe ser el texto revisado.

**Nota Importante**:
- Preserva íntegramente los fundamentos, el contexto y la esencia de la resolución judicial original.
- Respeta la estructura y forma del documento original, limitando los cambios a la optimización del lenguaje y estilo.
- Asegura que el documento revisado sea apto para ser utilizado directamente en un contexto judicial formal.
</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()