import requests
import json
import gradio as gr
import spaces

# Configuración de la API de NVIDIA
base_url = "https://integrate.api.nvidia.com/v1"
api_key = "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"
headers = {
    "Authorization": f"Bearer {api_key}",
    "Content-Type": "application/json",
    "Accept": "application/json"
}

# Instrucciones para el modelo
instrucciones = """
Eres un Ministro de la Corte Suprema de Chile, especializado en derecho civil. Tu tarea principal es mejorar borradores de resoluciones judiciales y responder preguntas generales, asegurando un tono formal y técnico cuando sea necesario, y proporcionando respuestas claras y precisas. Tus habilidades incluyen:
1. Utilizar un tono formal y técnico para garantizar claridad y coherencia.
2. Usar terminología legal específica, como 'folio', 'autos','resuelve', 'artículo', 'Código de Procedimiento Civil', 'audiencia', 'notificación', 'prueba', 'nulidad', 'oficiar' y 'notificar'.
3. Crear oraciones complejas y subordinadas, manteniendo un tono formal e imperativo.
4. Proporcionar referencias detalladas a regulaciones y precedentes judiciales, citando consistentemente artículos de leyes y sentencias anteriores.
5. Mantener un tono uniforme y detallado en todo el documento, centrado en la precisión terminológica.
"""

instrucciones_legales = """
Revisas el documento para asegurar:
- Comprensión del contexto y propósito.
- Identificación del tipo de documento judicial (sentencia, apelación).
- Progresión lógica de ideas y coherencia entre secciones.
- Claridad y precisión en cada párrafo y oración.
- Uso consistente de terminología legal.
- Corrección gramatical y claridad en la estructura de las oraciones.
- Uso adecuado de conectores para transiciones suaves.
- Precisión en términos legales y sugerencias de expresiones más efectivas.
- Consideración de sugerencias y validaciones, clasificándolas según su relevancia.
- Revisión integral y presentación concisa y unificada de la información.
"""

def mejorar_resolucion(input_text):
    # Construcción del prompt con instrucciones unificadas
    prompt = f"Adopta el rol de Ministro de la Corte Suprema de Chile y responde a la siguiente pregunta/pregunta con base en las instrucciones e instrucciones legales proporcionadas:\n{instrucciones}\n{instrucciones_legales}\nTexto o pregunta:\n{input_text}\nRespuesta:"
    
    payload = {
        "model": "nvidia/nemotron-4-340b-instruct",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.5,
        "top_p": 0.7,
        "max_tokens": 1024,
        "stream": False
    }

    try:
        response = requests.post(f"{base_url}/chat/completions", headers=headers, data=json.dumps(payload))
        response.raise_for_status()
        response_body = response.json()
        return response_body['choices'][0]['message']['content']
    except requests.exceptions.HTTPError as errh:
        return f"Error HTTP: {errh}"
    except requests.exceptions.ConnectionError as errc:
        return f"Error de conexión: {errc}"
    except requests.exceptions.Timeout as errt:
        return f"Error de tiempo de espera: {errt}"
    except requests.exceptions.RequestException as err:
        return f"Error de solicitud: {err}"

# Definición de la interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Mejora de Resoluciones Judiciales y Respuestas con Nemotron")
    input_text = gr.Textbox(label="Introduce tu resolución judicial o pregunta")
    output_text = gr.Textbox(label="Respuesta mejorada o respuesta corta")
    submit_button = gr.Button("Enviar")
    submit_button.click(fn=mejorar_resolucion, inputs=input_text, outputs=output_text)

# Lanzamiento de la aplicación
demo.launch()