"""
@RAG_CONTEXT
{
  "archivo": "src/ocr_indexer/ai_refiner.py",
  "lineas": 364,
  "tokens_estimados": 3487,
  "hash_contenido": "f0846ae3ef3aeb4d932b7424a85a8c9c93155c70d9d5b8260f50b4cac005c595",
  "creacion": "2026-03-17",
  "modificacion": "2026-03-17",
  "tags": [
    "Image",
    "requests",
    "fitz",
    "BytesIO",
    "logging",
    "io",
    "PIL",
    "base64",
    "json",
    "os"
  ],
  "descripcion_corta": "Refinador OCR con IA para corrección de errores Tesseract.",
  "descripcion_larga": "Módulo Python que corrige errores de OCR de Tesseract mediante LLMs (OpenAI, DeepSeek, Anthropic, Gemini). Implementa dos flujos: corrección de texto existente usando prompts especializados y extracción directa de imágenes mediante modelos de visión. Gestiona configuración jerárquica (variables de entorno, archivo .env, config.json), codificación de imágenes en base64, y selección automática de modelos según proveedor. Incluye soporte para procesamiento de PDFs con PyMuPDF y redimensionamiento de imágenes para APIs."
}
"""
#!/usr/bin/env python3
"""
Módulo de refinado OCR con IA.
Corrige los errores típicos de Tesseract enviando el texto a un LLM.
Soporta: OpenAI, DeepSeek, Anthropic, Gemini.
Config vía .env (RAG_AI_PROVIDER + RAG_API_KEY) o config.json["ia"].
"""
import os
import json
import logging

log = logging.getLogger("ai_refiner")

_MODELOS = {
    "openai":    "gpt-4o-mini",
    "deepseek":  "deepseek-chat",
    "anthropic": "claude-haiku-4-5-20251001",
    "gemini":    "gemini-2.5-flash",
}

# Modelos con capacidad de visión (más potentes que los de texto para OCR directo)
_MODELOS_VISION = {
    "openai":    "gpt-4o",
    "anthropic": "claude-3-5-sonnet-20241022",
    "gemini":    "gemini-2.5-flash",
}
_PROVEEDORES_VISION = {"openai", "anthropic", "gemini"}  # deepseek no soporta visión

_PROMPT_VISION = (
    "Eres un sistema OCR de alta precisión. Extrae TODO el texto visible en esta imagen "
    "exactamente como aparece. Preserva la estructura: saltos de línea, párrafos, columnas, "
    "tablas. No interpretes, resumas ni corrijas el contenido. "
    "Devuelve ÚNICAMENTE el texto extraído, sin comentarios adicionales, sin markdown."
)

_PROMPT = (
    "Eres un asistente especializado en corrección de texto OCR.\n"
    "El siguiente texto fue extraído mecánicamente con Tesseract de un documento "
    "escaneado y puede contener errores típicos: caracteres confundidos (O↔0, I↔1↔l, "
    "rn↔m), palabras partidas por guiones, espacios faltantes o incorrectos, "
    "puntuación corrupta.\n\n"
    "Corrige todos los errores de OCR preservando fielmente el contenido original. "
    "No agregues ni elimines información. Responde ÚNICAMENTE con el texto corregido, "
    "sin explicaciones, sin comillas adicionales, sin comentarios.\n\n"
    "ARCHIVO: {nombre}\n"
    "TEXTO OCR:\n{texto}"
)


def _leer_env(engine_path: str) -> dict:
    """Lee KEY=VALUE desde el .env en el directorio padre del motor."""
    env_path = os.path.join(os.path.dirname(engine_path), '.env')
    vals = {}
    if not os.path.exists(env_path):
        return vals
    try:
        with open(env_path, encoding='utf-8', errors='replace') as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#') or '=' not in line:
                    continue
                k, _, v = line.partition('=')
                vals[k.strip()] = v.strip().strip('"').strip("'")
    except Exception:
        pass
    return vals


def get_ai_config(engine_path: str, config: dict):
    """
    Retorna (proveedor, api_key, modelo).
    Precedencia: os.environ → .env → config.json["ia"]
    """
    env = _leer_env(engine_path)
    ia_cfg = config.get("ia", {}) if isinstance(config.get("ia"), dict) else {}

    proveedor = (
        os.environ.get("RAG_AI_PROVIDER", "").strip().lower()
        or env.get("RAG_AI_PROVIDER", "").strip().lower()
        or ia_cfg.get("proveedor", "").strip().lower()
    )
    api_key = (
        os.environ.get("RAG_API_KEY", "").strip()
        or env.get("RAG_API_KEY", "").strip()
        or ia_cfg.get("api_key", "").strip()
    )
    modelo = ia_cfg.get("modelo", "").strip() or _MODELOS.get(proveedor, "")
    return proveedor, api_key, modelo


def ai_disponible(engine_path: str, config: dict) -> bool:
    """Retorna True si hay proveedor y API key configurados."""
    proveedor, api_key, _ = get_ai_config(engine_path, config)
    return bool(proveedor and api_key)


def _encode_image(ruta_o_pil, max_dim: int = 2048):
    """
    Codifica una imagen en base64 JPEG para enviar a APIs de visión.
    Acepta ruta de archivo (str) o PIL Image. Para PDFs renderiza la primera página.
    Redimensiona si max(w, h) > max_dim. Retorna (base64_str, "image/jpeg").
    """
    import base64
    from io import BytesIO
    from PIL import Image

    if isinstance(ruta_o_pil, str):
        ext = os.path.splitext(ruta_o_pil)[1].lower()
        if ext == '.pdf':
            try:
                import fitz
                doc = fitz.open(ruta_o_pil)
                pix = doc[0].get_pixmap(dpi=150)
                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            except Exception as e:
                log.error(f"Error renderizando PDF para visión: {e}")
                raise
        else:
            img = Image.open(ruta_o_pil).convert("RGB")
    else:
        img = ruta_o_pil.convert("RGB")

    if max(img.size) > max_dim:
        img.thumbnail((max_dim, max_dim))

    buf = BytesIO()
    img.save(buf, format='JPEG', quality=85)
    b64 = base64.b64encode(buf.getvalue()).decode()
    return b64, "image/jpeg"


def vision_disponible(engine_path: str, config: dict) -> bool:
    """
    Retorna True si el modo visión está activado y el proveedor lo soporta.
    DeepSeek no soporta visión → siempre False aunque use_vision=true.
    """
    ia_cfg = config.get("ia", {}) if isinstance(config.get("ia"), dict) else {}
    if not ia_cfg.get("use_vision"):
        return False
    proveedor, api_key, _ = get_ai_config(engine_path, config)
    if proveedor == "deepseek":
        log.warning("DeepSeek no soporta visión. Usando Tesseract como fallback.")
        return False
    return bool(proveedor in _PROVEEDORES_VISION and api_key)


def call_vision(ruta: str, nombre: str, engine_path: str, config: dict) -> str:
    """
    Envía la imagen al LLM de visión para extraer el texto directamente.
    Retorna el texto extraído, o "" en caso de error (activa fallback a Tesseract).
    Timeout: 60 s.
    """
    try:
        import requests
    except ImportError:
        log.error("'requests' no instalado.")
        return ""

    proveedor, api_key, _ = get_ai_config(engine_path, config)
    modelo = _MODELOS_VISION.get(proveedor, "")
    if not modelo:
        log.warning(f"Sin modelo de visión para proveedor '{proveedor}'")
        return ""

    try:
        b64, media_type = _encode_image(ruta)
    except Exception as e:
        log.error(f"Error codificando imagen '{nombre}': {e}")
        return ""

    try:
        if proveedor == "openai":
            headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
            payload = {
                "model": modelo,
                "messages": [{"role": "user", "content": [
                    {"type": "text", "text": _PROMPT_VISION},
                    {"type": "image_url", "image_url": {"url": f"data:{media_type};base64,{b64}"}},
                ]}],
            }
            res = requests.post("https://api.openai.com/v1/chat/completions",
                                json=payload, headers=headers, timeout=60)
            res.raise_for_status()
            return res.json()["choices"][0]["message"]["content"].strip()

        elif proveedor == "anthropic":
            headers = {"x-api-key": api_key, "anthropic-version": "2023-06-01",
                       "content-type": "application/json"}
            payload = {
                "model": modelo,
                "max_tokens": 4096,
                "messages": [{"role": "user", "content": [
                    {"type": "image", "source": {"type": "base64",
                                                  "media_type": media_type, "data": b64}},
                    {"type": "text", "text": _PROMPT_VISION},
                ]}],
            }
            res = requests.post("https://api.anthropic.com/v1/messages",
                                json=payload, headers=headers, timeout=60)
            res.raise_for_status()
            return res.json()["content"][0]["text"].strip()

        elif proveedor == "gemini":
            url = (f"https://generativelanguage.googleapis.com/v1beta/models"
                   f"/{modelo}:generateContent?key={api_key}")
            payload = {"contents": [{"parts": [
                {"text": _PROMPT_VISION},
                {"inlineData": {"mimeType": media_type, "data": b64}},
            ]}]}
            res = requests.post(url, json=payload,
                                headers={"Content-Type": "application/json"}, timeout=60)
            res.raise_for_status()
            return res.json()["candidates"][0]["content"]["parts"][0]["text"].strip()

        else:
            log.warning(f"Proveedor sin soporte de visión: '{proveedor}'")
            return ""

    except Exception as e:
        log.error(f"Error visión IA para '{nombre}': {e}")
        return ""


def call_ai(texto: str, nombre: str, proveedor: str, api_key: str, modelo: str) -> str:
    """
    Envía el texto OCR al LLM para corrección.
    Retorna el texto corregido, o el texto original en caso de error.
    Timeout: 45 s. Sin dependencias externas excepto requests.
    """
    if not texto.strip():
        return texto

    try:
        import requests
    except ImportError:
        log.error("'requests' no instalado. Instala con: pip install requests")
        return texto

    prompt = _PROMPT.format(nombre=nombre, texto=texto[:6000])

    try:
        if proveedor in ("openai", "deepseek"):
            url = (
                "https://api.openai.com/v1/chat/completions"
                if proveedor == "openai"
                else "https://api.deepseek.com/chat/completions"
            )
            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            }
            payload = {
                "model": modelo,
                "messages": [{"role": "user", "content": prompt}],
            }
            res = requests.post(url, json=payload, headers=headers, timeout=45)
            res.raise_for_status()
            return res.json()["choices"][0]["message"]["content"].strip()

        elif proveedor == "anthropic":
            headers = {
                "x-api-key": api_key,
                "anthropic-version": "2023-06-01",
                "content-type": "application/json",
            }
            payload = {
                "model": modelo,
                "max_tokens": 4096,
                "messages": [{"role": "user", "content": prompt}],
            }
            res = requests.post(
                "https://api.anthropic.com/v1/messages",
                json=payload, headers=headers, timeout=45,
            )
            res.raise_for_status()
            return res.json()["content"][0]["text"].strip()

        elif proveedor == "gemini":
            url = (
                f"https://generativelanguage.googleapis.com/v1beta/models"
                f"/{modelo}:generateContent?key={api_key}"
            )
            payload = {"contents": [{"parts": [{"text": prompt}]}]}
            res = requests.post(
                url, json=payload,
                headers={"Content-Type": "application/json"}, timeout=45,
            )
            res.raise_for_status()
            return res.json()["candidates"][0]["content"]["parts"][0]["text"].strip()

        else:
            log.warning(f"Proveedor IA no soportado: '{proveedor}'")
            return texto

    except Exception as e:
        log.error(f"Error llamando IA para '{nombre}': {e}")
        return texto


def refinar_fase3(pendientes, db, index_path, engine_path, config, log_ref, actualizar_status):
    """
    Refina el texto OCR de `pendientes` usando IA.

    Estrategia (en orden de prioridad):
    1. Visión disponible → call_vision() re-extrae la imagen directamente (más preciso).
       Si visión falla → fallback a corrección de texto con call_ai().
    2. Sin visión pero IA configurada → call_ai() corrige el texto Tesseract.
    En ambos casos reemplaza texto_ocr y marca refinado_ia=True.
    """
    proveedor, api_key, modelo = get_ai_config(engine_path, config)
    use_vision = vision_disponible(engine_path, config)
    total = len(pendientes)

    modo = "visión" if use_vision else "texto"
    modelo_display = _MODELOS_VISION.get(proveedor, modelo) if use_vision else modelo

    actualizar_status({
        "fase": "Fase 3: Refinado IA",
        "total_ia": total,
        "procesados_ia": 0,
        "proveedor": proveedor,
        "modelo": modelo_display,
        "modo": modo,
    })
    print(f"  Modo: {modo} | Proveedor: {proveedor} | Modelo: {modelo_display}")

    for i, filepath in enumerate(pendientes, 1):
        entry = db.get(filepath, {})
        nombre = entry.get("archivo", {}).get("nombre", os.path.basename(filepath))
        ext = entry.get("archivo", {}).get("extension", "")

        log_ref.info(f"IA [{modo}] [{i}/{total}]: {nombre}")
        print(f"  IA [{modo}] [{i}/{total}]: {nombre}.{ext}")

        if use_vision:
            texto_refinado = call_vision(filepath, nombre, engine_path, config)
            if not texto_refinado:
                log_ref.warning(f"Visión falló para '{nombre}', fallback a corrección de texto")
                texto_ocr = entry.get("contenido", {}).get("texto_ocr", "")
                texto_refinado = call_ai(texto_ocr, nombre, proveedor, api_key, modelo)
        else:
            texto_ocr = entry.get("contenido", {}).get("texto_ocr", "")
            texto_refinado = call_ai(texto_ocr, nombre, proveedor, api_key, modelo)

        if filepath in db:
            if isinstance(db[filepath].get("contenido"), dict):
                db[filepath]["contenido"]["texto_ocr"] = texto_refinado
            db[filepath]["refinado_ia"] = True

        actualizar_status({
            "fase": "Fase 3: Refinado IA",
            "total_ia": total,
            "procesados_ia": i,
        })

        try:
            with open(index_path, 'w') as f:
                json.dump(db, f, indent=4)
        except Exception as e:
            log_ref.error(f"Error guardando tras refinar {filepath}: {e}")

    actualizar_status({"fase": "Completado", "total": len(db)})
    log_ref.info(f"Fase 3 completa: {total} archivos refinados ({modo}).")
    print(f"✅ Fase 3 completa: {total} archivos refinados con {proveedor}/{modelo_display} [{modo}].")
