"""
@RAG_CONTEXT
{
  "archivo": "src/ocr_indexer/ai_refiner.py",
  "lineas": 204,
  "tokens_estimados": 1833,
  "hash_contenido": "1b1685f3976f6b9cb0f80710f0163e07076e1d1fe8208fd6c841845d2b9f3bea",
  "creacion": "2026-03-17",
  "modificacion": "2026-03-17",
  "tags": [
    "logging",
    "json",
    "requests",
    "os"
  ],
  "descripcion_corta": "Refinador OCR con LLM para corrección automática de texto.",
  "descripcion_larga": "Módulo Python que corrige errores de OCR de Tesseract enviando texto a LLMs (OpenAI, DeepSeek, Anthropic, Gemini). Lee configuración desde variables de entorno, archivo .env o config.json, formatea prompts específicos para corrección OCR, y maneja llamadas API con timeout de 45 segundos. Incluye detección de disponibilidad de IA y fallback al texto original en errores."
}
"""
#!/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",
}

_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 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):
    """
    Itera sobre los filepaths de `pendientes`, llama call_ai() para cada uno,
    actualiza db[filepath]["contenido"]["texto_ocr"] con el resultado,
    marca db[filepath]["refinado_ia"] = True, y guarda index.json incrementalmente.
    """
    proveedor, api_key, modelo = get_ai_config(engine_path, config)
    total = len(pendientes)

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

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

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

        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.")
    print(f"✅ Fase 3 completa: {total} archivos refinados con {proveedor}/{modelo}.")
