core/pipeline.py

"""
Pipeline completo: Imagen → OCR → Parsing → Traducción → Descripciones → Markdown.
"""

import json
import time
from typing import Tuple

from PIL import Image

from config import MAX_DISHES
from core.models import ModelManager
from core.ocr import extract_text
from core.translation import translate_text
from core.description import generate_description
from core.parsing import parse_menu_structure
from utils.markdown_output import generate_markdown


class MenuTranslatorPipeline:
    """Orquesta los 3 modelos para traducir menús completos."""

    def __init__(self):
        self.manager = ModelManager()
        self.manager.load_all()

    def process_menu(
        self,
        image: Image.Image,
        target_lang: str = "English",
        detail_level: int = 2,
        include_cultural_notes: bool = True,
    ) -> Tuple[str, str, str]:
        """
        Pipeline completo.

        Returns:
            (markdown_output, stats_json, raw_spanish_text)
        """
        total_start = time.time()
        stats = {}

        print("\n" + "=" * 70)
        print("🚀 INICIANDO PROCESAMIENTO DE MENÚ")
        print("=" * 70)

        # --- PASO 1: OCR ---
        print("\n📷 PASO 1/4: Extracción de texto (OCR)...")
        ocr_result = extract_text(image, self.manager)

        if not ocr_result["success"]:
            return (
                f"❌ Error en OCR: {ocr_result.get('error', 'Desconocido')}",
                json.dumps({"error": "OCR failed"}, indent=2),
                "",
            )

        spanish_text = ocr_result["text"]
        stats["ocr"] = {
            "time": f"{ocr_result['processing_time']:.2f}s",
            "text_length": len(spanish_text),
            "model": ocr_result["model"],
        }
        print(f"✅ Texto extraído: {len(spanish_text)} caracteres")

        # --- PASO 2: PARSING ---
        print("\n🔍 PASO 2/4: Analizando estructura del menú...")
        dishes = parse_menu_structure(spanish_text)

        stats["parsing"] = {
            "dishes_found": len(dishes),
            "sections": len(set(d["section"] for d in dishes)),
        }
        print(
            f"✅ Encontrados {len(dishes)} platos en "
            f"{stats['parsing']['sections']} secciones"
        )

        # --- PASO 3: TRADUCCIÓN + DESCRIPCIONES ---
        print("\n🌐 PASO 3/4: Traduciendo y generando descripciones...")
        enriched_dishes, translation_time, description_time = (
            self._enrich_dishes(dishes, target_lang, detail_level)
        )

        stats["translation"] = {
            "time": f"{translation_time:.2f}s",
            "model": self.manager.translation_model.__class__.__name__,
        }
        stats["description"] = {
            "time": f"{description_time:.2f}s",
            "model": self.manager.description_model.__class__.__name__,
        }

        # --- PASO 4: MARKDOWN ---
        print("\n📝 PASO 4/4: Generando output en Markdown...")
        markdown = generate_markdown(
            enriched_dishes, target_lang, include_cultural_notes
        )

        total_time = time.time() - total_start
        stats["total"] = {
            "time": f"{total_time:.2f}s",
            "device": self.manager.device,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        }

        print(f"\n✅ PROCESAMIENTO COMPLETADO en {total_time:.2f}s")
        print("=" * 70 + "\n")

        return markdown, json.dumps(stats, indent=2), spanish_text

    def _enrich_dishes(self, dishes, target_lang, detail_level):
        """Genera descripciones traducidas para cada plato (nombres sin traducir)."""
        translation_time = 0.0
        description_time = 0.0
        enriched = []

        for i, dish in enumerate(dishes[:MAX_DISHES], 1):
            print(f"  [{i}/{min(len(dishes), MAX_DISHES)}] {dish['name'][:40]}...")

            # Generar descripción en español
            t0 = time.time()
            desc_es = generate_description(
                dish["name"], self.manager,
                detail_level=detail_level,
            )
            description_time += time.time() - t0

            # Traducir solo la descripción al idioma destino
            t0 = time.time()
            desc_result = translate_text(desc_es, self.manager, target_lang)
            translation_time += time.time() - t0

            enriched.append({
                **dish,
                "description": desc_result["translated_text"],
            })

        return enriched, translation_time, description_time