src/logic.py

"""
Este módulo contiene toda la lógica de negocio para la aplicación de contestación de tutelas.
Se encarga de interactuar con la API de Gemini, procesar los archivos PDF y DOCX,
y orquestar el flujo de trabajo completo.
"""

import os
import json
import re
import tempfile
import pathlib
from docx import Document
from docx.shared import Inches
import google.generativeai as genai
import gradio as gr

# --- CONFIGURACIÓN GLOBAL ---

API_KEY = os.environ.get("GEMINI_API_KEY")
if API_KEY:
    genai.configure(api_key=API_KEY)

EXAMPLE_PDF_PATH = "ejemplo_tutela.pdf"

# --- PLANTILLAS DE PROMPTS ---

EXTRACT_PROMPT = '''
Del siguiente documento de acción de tutela, extrae de forma precisa y concisa la siguiente información:
1.  **juzgado**: Nombre completo del juzgado al que se dirige la tutela.
2.  **ciudad**: La ciudad donde se presenta la tutela.
3.  **radicado**: El número de radicado o referencia del caso.
4.  **accionante**: El nombre completo del demandante o accionante.
5.  **direccion_accionante**: La dirección física del accionante.
6.  **email**: La lista de correos electrónicos de notificación.
7.  **datos**: Una lista de las partes involucradas (demandante, demandado).
8.  **hechos**: Un resumen claro y numerado de los hechos que motivan la tutela.
9.  **argumentos**: Un resumen de los argumentos legales del accionante.
10. **peticiones**: Un resumen de lo que el accionante solicita al juez.

Devuelve la respuesta como un objeto JSON válido con las claves especificadas.
'''

ENCABEZADO_PROMPT = '''
Con base en el documento de ejemplo, redacta el encabezado para una contestación de tutela.
Usa el mismo estilo, tono y formato del ejemplo.
Información de la tutela a contestar:
- Juzgado: {juzgado}
- Ciudad: {ciudad}
- Radicado: {radicado}
- Accionante: {accionante}
- Accionado: (El accionado es la entidad que representas, no lo inventes, usa un placeholder como [Nombre de la Entidad Accionada])

Devuelve un JSON con la clave "encabezado" y el texto correspondiente.
'''

HECHOS_PROMPT = '''
Actuando como un abogado experto, redacta la sección de \'HECHOS\' para una contestación de tutela, basándote estrictamente en el estilo y formato del documento de ejemplo.
**Instrucción clave: Debes responder a CADA UNO de los hechos numerados presentados por el accionante.** La respuesta debe seguir el formato "Al hecho PRIMERO, es cierto.", "Al hecho SEGUNDO, no es cierto por...", o una variación similar que se alinee con el ejemplo.

**Hechos del Accionante a Responder:**
{hechos}

**Contexto General de las Partes:**
{datos}

Devuelve un JSON válido con una única clave "hechos" que contenga el texto completo de la sección.
'''

FUNDAMENTOS_PROMPT = '''
Actuando como un abogado experto, redacta la sección de \'FUNDAMENTOS DE DERECHO\' para una contestación de tutela, basándote estrictamente en el estilo y formato del documento de ejemplo.
**Instrucción clave: Debes rebatir, uno por uno, los argumentos legales del accionante.** La estructura de tu respuesta debe reflejar una contra-argumentación directa a cada punto presentado.

**Argumentos del Accionante a Rebatir:**
{argumentos}

**Contexto General de las Partes:**
{datos}

Devuelve un JSON válido con una única clave "fundamentos" que contenga el texto completo de la sección.
'''

PETICIONES_PROMPT = '''
Actuando como un abogado experto, redacta la sección de \'PETICIONES\' para una contestación de tutela, basándote estrictamente en el estilo y formato del documento de ejemplo.
**Instrucción clave: Debes pronunciarte sobre CADA UNA de las peticiones del accionante.** La respuesta debe ser una oposición directa y formal a sus pretensiones, siguiendo el tono del ejemplo.

**Peticiones del Accionante a Responder:**
{peticiones}

**Contexto General de las Partes:**
{datos}

Devuelve un JSON válido con una única clave "peticiones" que contenga el texto completo de la sección.
'''


# --- FUNCIONES AUXILIARES ---

def _safe_json_loads(text: str) -> dict:
    """
    Analiza de forma robusta un string que debería ser JSON.
    Intenta limpiar el texto y encontrar un objeto JSON válido.
    """
    if not isinstance(text, str):
        text = str(text)

    # Elimina las vallas de formato de código (```json ... ```) si existen.
    text = re.sub(r"^```(?:json)?\s*|\s*```$", "", text, flags=re.S).strip()

    try:
        return json.loads(text)
    except json.JSONDecodeError as e:
        # Si falla, intenta encontrar el primer y más grande bloque JSON (entre { y }).
        match = re.search(r"\{.*\}", text, re.DOTALL)
        if match:
            json_like_string = match.group(0)
            try:
                return json.loads(json_like_string)
            except json.JSONDecodeError:
                raise ValueError(f"Error al decodificar el JSON extraído: {json_like_string}") from e
        raise ValueError(f"No se pudo encontrar un objeto JSON válido en la respuesta: {text}") from e

def _safe_json_field(raw, field):
    """Extrae de forma segura un campo de una respuesta JSON."""
    if raw is None:
        return ""
    if isinstance(raw, list):
        return str(raw)
    if isinstance(raw, dict):
        return str(raw.get(field, ""))

    try:
        obj = _safe_json_loads(str(raw))
        if isinstance(obj, dict):
            return str(obj.get(field, ""))
        return str(obj) # Devuelve el objeto parseado si no es un diccionario
    except ValueError:
        return str(raw) # Si todo falla, devuelve el texto original

def _add_section(doc: Document, title: str, body: str):
    """Añade una sección con título y cuerpo a un documento de Word."""
    doc.add_heading(title, level=1)
    body = str(body) if body is not None else ""
    blocks = [b.strip() for b in body.strip().split("\n\n") if b.strip()]
    if not blocks:
        doc.add_paragraph("(sin contenido)")
        return
    for b in blocks:
        doc.add_paragraph(b)

# --- LÓGICA PRINCIPAL CON LA API DE GEMINI ---

def extract_tutela_parts(uploaded_file_path: str) -> dict:
    """Usa Gemini para extraer las partes clave de una tutela."""
    model = genai.GenerativeModel("gemini-1.5-flash")
    
    uploaded_file = genai.upload_file(path=uploaded_file_path, display_name="Tutela PDF")
    
    try:
        response = model.generate_content(
            [EXTRACT_PROMPT, uploaded_file],
            generation_config=genai.types.GenerationConfig(
                response_mime_type="application/json",
                response_schema={
                    "type": "object",
                    "properties": {
                        "datos": {"type": "array", "items": {"type": "string"}},
                        "hechos": {"type": "array", "items": {"type": "string"}},
                        "argumentos": {"type": "array", "items": {"type": "string"}},
                        "peticiones": {"type": "array", "items": {"type": "string"}},
                        "juzgado": {"type": "string"},
                        "direccion_juzgado": {"type": "string"},
                        "email": {"type": "array", "items": {"type": "string"}},
                        "radicado": {"type": "string"},
                        "accionante": {"type": "string"},
                        "direccion_accionante": {"type": "string"},
                        "ciudad": {"type": "string"}
                    },
                    "required": ["juzgado", "ciudad", "radicado", "accionante", "hechos", "argumentos", "peticiones"]
                }
            )
        )
        return _safe_json_loads(response.text)
    finally:
        genai.delete_file(uploaded_file.name)

def generate_response_section(section_name: str, prompt_template: str, context_data: dict) -> str:
    """Función genérica para generar una sección de la contestación."""
    model = genai.GenerativeModel("gemini-1.5-flash")
    
    example_filepath = pathlib.Path(EXAMPLE_PDF_PATH)
    if not example_filepath.is_file():
        raise FileNotFoundError(f"El archivo de ejemplo '{EXAMPLE_PDF_PATH}' no se encontró.")

    system_instruction = prompt_template.format(**context_data)
    
    example_file = genai.upload_file(path=str(example_filepath), display_name="Ejemplo Tutela")
    
    try:
        response = model.generate_content(
            [system_instruction, example_file],
            generation_config=genai.types.GenerationConfig(
                response_mime_type="application/json",
                response_schema={"type": "object", "properties": {section_name: {"type": "string"}}}
            )
        )
        return _safe_json_field(response.text, section_name)
    finally:
        try:
            genai.delete_file(example_file.name)
        except Exception:
            pass

def write_tutela_docx(encabezado: str, hechos: str, peticiones: str, fundamentos: str) -> str:
    """Crea un archivo .docx con el encabezado y las secciones generadas."""
    doc = Document()
    for s in doc.sections:
        s.top_margin, s.bottom_margin, s.left_margin, s.right_margin = [Inches(1)] * 4

    # Añadir encabezado
    if encabezado:
        doc.add_paragraph(encabezado)
        doc.add_paragraph("\n") # Espacio

    _add_section(doc, "HECHOS", hechos)
    _add_section(doc, "FUNDAMENTOS DE DERECHO", fundamentos)
    _add_section(doc, "PETICIONES", peticiones)
    
    with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tmp:
        doc.save(tmp.name)
        return tmp.name

# --- FUNCIÓN PRINCIPAL DEL FLUJO DE TRABAJO ---

def full_workflow(uploaded_file_path: str):
    """Orquesta todo el proceso: desde la subida del PDF hasta la generación del DOCX."""
    if not uploaded_file_path:
        raise gr.Error("Por favor, sube un archivo de tutela en formato PDF.")

    try:
        # 1. Extraer las partes de la tutela del usuario.
        extracted_data = extract_tutela_parts(uploaded_file_path)
        
        # 2. Generar cada sección de la contestación.
        encabezado_generado = generate_response_section("encabezado", ENCABEZADO_PROMPT, extracted_data)
        hechos_generados = generate_response_section("hechos", HECHOS_PROMPT, extracted_data)
        fundamentos_generados = generate_response_section("fundamentos", FUNDAMENTOS_PROMPT, extracted_data)
        peticiones_generadas = generate_response_section("peticiones", PETICIONES_PROMPT, extracted_data)

        # 3. Crear el archivo .docx con los textos generados.
        docx_path = write_tutela_docx(encabezado_generado, hechos_generados, peticiones_generadas, fundamentos_generados)

        # 4. Devolver todos los resultados a la interfaz de Gradio.
        return hechos_generados, fundamentos_generados, peticiones_generadas, docx_path

    except FileNotFoundError as e:
        print(f"Error: {e}")
        raise gr.Error(f"Error de configuración del servidor: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        raise gr.Error(f"Ha ocurrido un error inesperado: {e}")