"""
@RAG_CONTEXT
{
  "archivo": "src/ocr_indexer/indexer.py",
  "lineas": 287,
  "tokens_estimados": 2766,
  "hash_contenido": "33591b78a9d4610e18485ce96b0f785dbaebd147a4ea5c300c264313c250f7ae",
  "creacion": "2026-03-13",
  "modificacion": "2026-03-17",
  "tags": [
    "sys",
    "time",
    "hashlib",
    "logging",
    "shutil",
    "fitz",
    "PIL",
    "json",
    "pypdf",
    "pytesseract",
    "tempfile",
    "DecompressionBombWarning",
    "PIL.Image",
    ".ai_refiner",
    "datetime",
    "warnings",
    "Image",
    "ai_disponible",
    "os"
  ],
  "descripcion_corta": "Indexador OCR con Proxy para imágenes grandes",
  "descripcion_larga": "Script Python que indexa documentos mediante OCR (Tesseract), procesando imágenes y PDFs. Implementa un patrón Proxy que detecta imágenes grandes (>50MP) mediante DecompressionBombWarning, reduciéndolas a 2500x2500 píxeles en memoria temporal antes del OCR. Incluye funciones para cálculo de hash MD5, filtrado blacklist, y manejo de configuración mediante archivos JSON. Depende de Pillow, pytesseract, pypdf/pymupdf y tesseract instalado en sistema."
}
"""
#!/usr/bin/env python3
import os
import sys
import json
import hashlib
import logging
import tempfile
import time
import warnings
from datetime import datetime

try:
    from PIL import Image
    from PIL.Image import DecompressionBombWarning
    import pytesseract
    import pypdf
    Image.MAX_IMAGE_PIXELS = 50_000_000  # 50 MP — límite seguro
    # Convertir Warning en excepción para que el Proxy se active a partir de 50 MP
    warnings.filterwarnings("error", category=DecompressionBombWarning)
except ImportError:
    print("Faltan dependencias. Instala con: pip install Pillow pytesseract pypdf")
    sys.exit(1)

import shutil as _shutil
if not _shutil.which("tesseract"):
    print("ERROR: Tesseract no está instalado en el sistema.")
    print("Instálalo con: ocr-indexer install-deps")
    print("  o manualmente: apt-get install -y tesseract-ocr tesseract-ocr-spa")
    sys.exit(1)

log = logging.getLogger("indexer")


def calcular_hash(ruta):
    hasher = hashlib.md5()
    try:
        with open(ruta, 'rb') as f:
            buf = f.read(65536)
            while len(buf) > 0:
                hasher.update(buf)
                buf = f.read(65536)
        return hasher.hexdigest()
    except Exception as e:
        log.error(f"Error calculando hash de {ruta}: {e}")
        return None


def es_permitido(ruta, blacklist):
    for bloqueado in blacklist:
        if f"/{bloqueado}/" in ruta or ruta.endswith(f"/{bloqueado}"):
            return False
    return True


def procesar_ocr(ruta):
    """OCR normal. Si la imagen supera el límite de píxeles (50MP), activa el Patrón Proxy."""
    try:
        img = Image.open(ruta)
        return pytesseract.image_to_string(img, lang='spa+eng').strip()
    except DecompressionBombWarning:
        log.warning(f"Imagen gigante detectada, activando Proxy: {ruta}")
        tmp_path = None
        try:
            Image.MAX_IMAGE_PIXELS = None          # desactivar límite solo para esta imagen
            img = Image.open(ruta)
            img.thumbnail((2500, 2500))            # reducir en memoria antes de guardar
            with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
                tmp_path = tmp.name
            img.save(tmp_path, 'JPEG')
            del img                                # liberar RAM de la imagen original
            return pytesseract.image_to_string(Image.open(tmp_path), lang='spa+eng').strip()
        except Exception as e:
            log.error(f"Error en Proxy para {ruta}: {e}")
            return ""
        finally:
            Image.MAX_IMAGE_PIXELS = 50_000_000   # restaurar límite global siempre
            warnings.filterwarnings("error", category=DecompressionBombWarning)
            if tmp_path and os.path.exists(tmp_path):
                try:
                    os.remove(tmp_path)
                except Exception as e:
                    log.warning(f"No se pudo eliminar temporal {tmp_path}: {e}")
    except Exception as e:
        log.error(f"Error OCR en {ruta}: {e}")
        return ""


def procesar_pdf(ruta):
    try:
        import fitz  # pymupdf
        doc = fitz.open(ruta)
        pix = doc[0].get_pixmap(dpi=150)
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        return pytesseract.image_to_string(img, lang='spa+eng').strip()
    except Exception as e:
        log.error(f"Error procesando PDF {ruta}: {e}")
        return ""

def main():
    if len(sys.argv) < 2:
        sys.exit(1)

    engine_path = sys.argv[1]
    solo_fase2 = len(sys.argv) > 2 and sys.argv[2] == "--fase2"

    config_path = os.path.join(engine_path, 'config.json')
    index_path = os.path.join(engine_path, 'ocr-indexer.json')
    status_path = os.path.join(engine_path, 'status.json')
    log_path = os.path.join(engine_path, 'indexer.log')
    pid_path = os.path.join(engine_path, 'fase2.pid')

    logging.basicConfig(
        filename=log_path,
        level=logging.INFO,
        format='%(asctime)s [%(levelname)s] %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S',
        force=True  # forzar reconfiguración aunque pytesseract/PIL hayan pre-configurado el root logger
    )

    def actualizar_status(datos):
        try:
            with open(status_path, 'w') as f:
                json.dump(datos, f, indent=4)
        except Exception as e:
            log.error(f"Error escribiendo status.json: {e}")

    try:
        with open(config_path, 'r') as f:
            config = json.load(f)
        root_dir = config.get("root", os.path.dirname(engine_path))
        blacklist_raiz = config.get("blacklist", [])
    except Exception as e:
        log.critical(f"Error leyendo config.json: {e}")
        sys.exit(1)

    try:
        with open(index_path, 'r') as f:
            db = json.load(f)
    except Exception as e:
        log.warning(f"ocr-indexer.json no encontrado o corrupto, iniciando vacío: {e}")
        db = {}

    extensiones_validas = {'.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.pdf'}

    if solo_fase2:
        print("--- MODO VIGÍA: SALTANDO A FASE 2 ---")
    else:
        # ==========================================
        # FASE 1: ESCANEO RÁPIDO Y CREACIÓN DE JSON
        # ==========================================
        print("--- INICIANDO FASE 1: ESCANEO DE ESTRUCTURA ---")
        actualizar_status({"fase": "Escaneando directorios...", "inicio": time.time()})
    
        for actual_dir, dirs, files in os.walk(root_dir):
            dirs[:] = [d for d in dirs if d not in blacklist_raiz]

            # Absorción de sub-motores
            sub_engine_json = os.path.join(actual_dir, '.ocr-indexer', 'ocr-indexer.json')
            if actual_dir != root_dir and os.path.exists(sub_engine_json):
                try:
                    with open(sub_engine_json, 'r') as f:
                        sub_db = json.load(f)
                    for filepath, data in sub_db.items():
                        if es_permitido(filepath, blacklist_raiz) and "integridad" in data:
                            data["estado"] = "completado"
                            db[filepath] = data
                except Exception as e:
                    log.warning(f"Error absorbiendo sub-motor en {actual_dir}: {e}")

            for filename in files:
                ext = os.path.splitext(filename)[1].lower()
                if ext not in extensiones_validas:
                    continue

                filepath = os.path.join(actual_dir, filename)
                hash_actual = calcular_hash(filepath)
                if not hash_actual:
                    continue

                if filepath in db and db[filepath].get("integridad", {}).get("hash_md5") == hash_actual:
                    if "estado" not in db[filepath]:
                        db[filepath]["estado"] = "completado"
                    continue

                db[filepath] = {
                    "archivo": {"nombre": os.path.splitext(filename)[0], "extension": ext.replace(".", ""), "directorio": actual_dir},
                    "fechas": {"indexacion": datetime.now().isoformat()},
                    "integridad": {"hash_md5": hash_actual},
                    "contenido": {"texto_ocr": ""},
                    "estado": "pendiente"
                }

        try:
            with open(index_path, 'w') as f:
                json.dump(db, f, indent=4)
        except Exception as e:
            log.critical(f"Error guardando ocr-indexer.json tras Fase 1: {e}")
        print("FASE 1 COMPLETADA: Base de datos estructurada.")

    # ==========================================
    # FASE 2: PROCESAMIENTO OCR PESADO
    # ==========================================
    print("--- INICIANDO FASE 2: EXTRACCIÓN OCR ---")
    pendientes = [k for k, v in db.items() if v.get("estado") == "pendiente"]
    total_pendientes = len(pendientes)
    
    if total_pendientes == 0:
        actualizar_status({"fase": "Completado", "total": len(db)})
        print("No hay archivos nuevos para procesar.")
        sys.exit(0)

    # Escribir PID file para que el vigía pueda detectar si ya hay una Fase 2 activa
    try:
        with open(pid_path, 'w') as f:
            f.write(str(os.getpid()))
    except Exception as e:
        log.warning(f"No se pudo escribir fase2.pid: {e}")

    inicio_ocr = time.time()
    procesados = 0

    actualizar_status({
        "fase": "Procesando OCR",
        "inicio_ocr": inicio_ocr,
        "total_ocr": total_pendientes,
        "procesados": 0
    })

    try:
        for filepath in pendientes:
            log.info(f"OCR: {os.path.basename(filepath)}")
            ext = db[filepath]["archivo"]["extension"]
            nombre_f = db[filepath]["archivo"]["nombre"]

            print(f"OCR: {nombre_f}.{ext}")
            texto_extraido = procesar_pdf(filepath) if ext == "pdf" else procesar_ocr(filepath)

            db[filepath]["contenido"]["texto_ocr"] = texto_extraido
            db[filepath]["estado"] = "completado"
            procesados += 1

            actualizar_status({
                "fase": "Procesando OCR",
                "inicio_ocr": inicio_ocr,
                "total_ocr": total_pendientes,
                "procesados": procesados
            })

            # Guardado incremental tras cada archivo — ningún OOM kill destruye el progreso
            try:
                with open(index_path, 'w') as f:
                    json.dump(db, f, indent=4)
            except Exception as e:
                log.error(f"Error en guardado incremental tras {filepath}: {e}")

        actualizar_status({"fase": "Completado", "total": len(db)})
        log.info("Escaneo total finalizado.")
        print("🏁 Escaneo total finalizado.")

        # ==========================================
        # FASE 3: REFINADO IA (opcional)
        # Visión: re-extrae imagen con LLM de visión (más preciso que Tesseract)
        # Texto:  envía el OCR mecánico al LLM para corrección
        # ==========================================
        try:
            from .ai_refiner import ai_disponible, vision_disponible, refinar_fase3
            if ai_disponible(engine_path, config) or vision_disponible(engine_path, config):
                pendientes_sin_refinar = [p for p in pendientes if not db.get(p, {}).get("refinado_ia")]
                if pendientes_sin_refinar:
                    print("--- INICIANDO FASE 3: REFINADO IA ---")
                    refinar_fase3(pendientes_sin_refinar, db, index_path, engine_path, config, log, actualizar_status)
                else:
                    print("Fase 3 omitida: todos los archivos ya refinados.")
            else:
                print("Fase 3 omitida: configura RAG_AI_PROVIDER y RAG_API_KEY para activarla.")
        except Exception as e:
            log.error(f"Error en Fase 3 (Refinado IA): {e}")
            print(f"⚠️  Fase 3 omitida por error: {e}")
    finally:
        if os.path.exists(pid_path):
            try:
                os.remove(pid_path)
            except Exception:
                pass

if __name__ == "__main__":
    main()
