import gradio as gr
import os
from openai import OpenAI
from gradio_agentchatbot import AgentChatbot, ChatMessage

# Configura la clave de API
api_key = os.getenv("API_KEY")
if not api_key:
    raise ValueError("La variable de entorno 'API_KEY' no está configurada.")

# Inicializa el cliente OpenAI
client = OpenAI(base_url="https://integrate.api.nvidia.com/v1", api_key=api_key)

# Define el mensaje del sistema (mejorado)
system_message = {
    "role": "system",
    "content": """
Actúa como un Ministro de la Corte Suprema de Chile, especializado en Derecho de Familia. Tu tarea es perfeccionar la redacción de un borrador de resolución judicial. Debes lograr claridad, precisión, concisión y formalidad, **sin alterar la sustancia de la decisión.**

## Instrucciones:

1. **Reemplaza términos ambiguos o coloquiales con terminología jurídica precisa.** (ej., "régimen de visitas" en lugar de "ver al niño", "pensión alimenticia" en lugar de "dar plata").
2. **Elimina redundancias y frases innecesarias.** Simplifica las oraciones manteniendo el significado original.
3. **Utiliza un tono formal e impersonal.** Evita expresiones subjetivas o coloquiales. Expresa las decisiones de forma objetiva y concisa.
4. **Usa expresiones judiciales chilenas estándar, especialmente en conclusiones y firmas.** (ej., "Proveyó y firmó el Juez mediante firma electrónica avanzada").
5. **Mantén la estructura original del borrador.** No reordenes párrafos ni secciones.

## Ejemplos:

- Borrador: "El juez decide que el papá puede ver a su hijo los fines de semana."
  Reformular a: "Se establece un régimen de visitas del menor con su padre los fines de semana."
- Borrador: "Se ordena pagarle una pensión a la mamá para los gastos del niño."
  Reformular a: "Se fija una pensión alimenticia a favor de la demandante para la manutención del menor."
- Borrador: "El papá tiene que pagar la mitad de los gastos del colegio."
  Reformular a: "Se establece que ambos padres contribuirán en partes iguales al pago de los gastos educacionales del menor."
- Borrador: "Yo creo que lo mejor para el niño es..."
  Reformular a: "Se considera que la mejor opción para el bienestar del menor es..."

## Restricciones:

- **No modifiques las decisiones, los montos, las fechas, ni las partes involucradas en la resolución.** Tu tarea se limita a mejorar la redacción del texto existente.
- No interpretes ni amplíes el contenido existente.
- No incluyas análisis o comentarios personales.

## Formato de Salida:

Devuelve el borrador de la resolución judicial revisado y mejorado en formato de texto, manteniendo la estructura original.
"""
}

def llama_response(prompt, messages):
    messages.append(ChatMessage(role="user", content=prompt))
    yield messages, "Procesando..."

    user_message = {"role": "user", "content": prompt}

    try:
        completion = client.chat.completions.create(
            model="meta/llama-3.1-405b-instruct",
            messages=[system_message, user_message],
            temperature=0.8,
            top_p=0.7,
            max_tokens=4000,
            stream=True,
        )

        response_text = ""
        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                response_text += chunk.choices[0].delta.content
                messages.append(ChatMessage(role="assistant", content=response_text))
                yield messages, "Procesando..."

        yield messages, "Listo"

    except Exception as e:
        print(f"Error: {e}")
        yield messages + [ChatMessage(role="assistant", content="Error al procesar la solicitud.")], "Error"

# Interfaz de Gradio (mejorada)
description = """
## Ministro IA de la Corte Suprema - Derecho de Familia

Este espacio perfecciona la redacción de resoluciones judiciales chilenas en Derecho de Familia.

**Instrucciones:**
1. Pega el borrador en el cuadro de texto.
2. Haz clic en "Submit".
"""

examples = [
    ["El juez decide que el papá puede ver a su hijo los fines de semana."],
    ["Se ordena pagarle una pensión a la mamá para los gastos del niño."],
    ["El papá tiene que pagar la mitad de los gastos del colegio."],
    ["Yo creo que lo mejor para el niño es..."]
]

with gr.Blocks() as demo:
    chatbot = AgentChatbot(label="Resolución")
    text_input = gr.Textbox(
        lines=10,
        label="Borrador de la Resolución",
        placeholder="Ingrese el borrador aquí...",
    )
    text_input.submit(
        llama_response, [text_input, chatbot], [chatbot, gr.Label(label="Estado")]
    )

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