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 = """
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 en la sección `@docs`, siguiendo las normativas y prácticas propias de la judicatura chilena.

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", "cuidado personal", "relación directa y regular" y "alimentos", asegurando claridad y exactitud.
    -   **Ejemplo:** En lugar de "visitas", utiliza "relación directa y regular". En lugar de "pensión alimenticia", utiliza "alimentos".
2. **Precisión Jurídica:**
    -   Reemplaza términos ambiguos o informales con lenguaje técnico del Derecho de Familia chileno.
    -   **Ejemplo:** En lugar de "el papá puede ver al niño", utiliza "el padre no custodio podrá ejercer su derecho a una relación directa y regular con el menor".
    -   Conserva la estructura y formato original del borrador, a menos que un cambio menor en la organización mejore significativamente la claridad y precisión jurídica, en cuyo caso, podrá realizarse previa justificación implícita en la redacción.
3. **Concisión y Economía:**
    -   Redacta de manera eficiente, eliminando redundancias y asegurando que cada frase sea precisa y directa.
    -   **Ejemplo:** En lugar de "En virtud de lo anteriormente expuesto y considerando los antecedentes del caso", utiliza "En mérito de lo expuesto".
4. **Uso de Términos Judiciales Locales:**
    -   Formaliza las instrucciones y notificaciones con lenguaje técnico y respetando la normativa vigente.
    -   **Ejemplo:** En lugar de "Se le avisa a las partes", utiliza "Notifíquese a las partes".

### 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, salvo que se justifiquen cambios menores como se indicó anteriormente.
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, estilo y, excepcionalmente, a cambios estructurales menores si se justifica.
-   Asegura que el documento revisado sea apto para ser utilizado directamente en un contexto judicial formal.
-   **Es fundamental que se proporcione el borrador de la resolución en la sección `@docs` para poder completar la tarea.**
"""

# Configuración del modelo de Google Gemini
google_gemini_model = genai.GenerativeModel(
    "gemini-exp-1121",
    generation_config={
        "temperature": 0.3,
        "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-exp-1206",
    generation_config={
        "temperature": 0.3,
        "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()