core/engine.py

"""EuropaLex Inference Engine — Local model backends via llama-cli, llama-cpp-python, and Python packages."""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass
from pathlib import Path
from typing import ClassVar

import numpy as np
import torch

from core.types import CEFRLevel, EngineConfig, TextResult, ValidationError

logger = logging.getLogger(__name__)


@dataclass
class _EngineState:
    """Tracks which GPU engine is currently loaded."""

    translation_engine: LlamaCppTextEngine | None = None
    tts_engine: TTSEngine | None = None
    image_engine: ImageGenEngine | None = None


class MiniCPMTextEngine:
    """Generates English text using MiniCPM5-1B Q8_0 via llama-cpp-python.

    Lazy-loads the model on first call, unloads after completion to free memory.
    Uses create_chat_completion to format prompts with MiniCPM's required message tokens.
    Only one instance can be active at a time (enforced by EnginePool).

    Best for Phase 1 English text generation — ~1.1 GB RAM, no subprocess overhead.
    """

    def __init__(self, model_path: str, device: str = "cuda"):
        """Initialize the text engine.

        Args:
            model_path: Absolute path to the MiniCPM5-1B Q8_0 GGUF file.
            device: Device hint ('cuda', 'mps', or 'cpu').
        """
        self.model_path = Path(model_path)
        if not self.model_path.exists():
            raise FileNotFoundError(
                f"MiniCPM5-1B model not found at: {self.model_path}\n"
                f"Run: python models/download_models.py minicpm"
            )
        self.device = device
        self._llm = None
        self._loaded = False

    def _load_model(self) -> None:
        """Lazy-load the GGUF model via llama-cpp-python."""
        if self._loaded:
            return

        try:
            from llama_cpp import Llama
        except ImportError:
            raise ImportError(
                "llama-cpp-python package not installed. "
                "Run: pip install llama-cpp-python"
            )

        n_gpu = 99 if self.device == "cuda" else 0
        self._llm = Llama(
            model_path=str(self.model_path),
            n_gpu_layers=n_gpu,
            n_ctx=4096,
        )
        self._loaded = True
        logger.info("MiniCPMTextEngine loaded %s on %s", self.model_path.name, self.device)

    def generate(
        self,
        texts: list[str],
        scenario: str,
        cefr_level: CEFRLevel,
        batch_size: int | None = None,
        topic_description: str = "",
    ) -> TextResult:
        """Generate English sentences using the loaded GGUF model.

        Delegates to :func:`core.text_gen.generate_sentences` for LLM calling,
        retry loop, and extraction. Wraps result in ``TextResult``.

        Args:
            texts: Empty list (generation mode). Non-empty would be translation mode.
            scenario: Scenario/topic description for text generation.
            cefr_level: CEFR proficiency level (linguistic guidance only).
            batch_size: Number of sentences to generate.
            topic_description: Free-form description of topics/themes.

        Returns:
            TextResult with exactly one sentence per requested batch size.

        Raises:
            ValidationError: If generation fails after max attempts.
        """
        self._load_model()
        if batch_size is None:
            raise ValueError("batch_size is required for text generation")

        from core.text_gen import generate_sentences

        sentences = generate_sentences(scenario, cefr_level, batch_size, self._llm, topic_description)
        return TextResult(generated_texts=sentences)

    def unload(self) -> None:
        """Unload the model and free memory."""
        if self._llm is not None:
            del self._llm
            self._llm = None
            self._loaded = False
            try:
                torch.cuda.empty_cache()
            except Exception:
                pass
            logger.info("MiniCPMTextEngine unloaded")


class LlamaCppTextEngine:
    """Generates text using llama-cpp-python (GGUF models, lazy-load + unload).

    Uses the llama-cpp-python Python bindings instead of spawning subprocesses.
    Lazy-loads the model on first call, unloads after completion to free VRAM.
    Only one instance can be active at a time (enforced by EnginePool).

    Best for smaller GGUF models (e.g. tiny-aya-water Q4_K_M ~2 GB) where
    keeping the model in Python memory is efficient and avoids subprocess overhead.
    """

    def __init__(
        self,
        model_path: str,
        device: str = "cuda",
        target_language: str = "Latvian",
    ):
        """Initialize the translation engine.

        Args:
            model_path: Absolute path to the GGUF model file.
            device: Device hint (informational; n_gpu_layers=99 used for CUDA).
            target_language: Target language for translations (e.g. "Latvian").
        """
        self.model_path = Path(model_path)
        if not self.model_path.exists():
            raise FileNotFoundError(f"Model not found: {self.model_path}")
        self.device = device
        self.target_language = target_language
        self._llm = None
        self._loaded = False

    def _load_model(self) -> None:
        """Lazy-load the GGUF model via llama-cpp-python."""
        if self._loaded:
            return

        try:
            from llama_cpp import Llama
        except ImportError:
            raise ImportError(
                "llama-cpp-python package not installed. "
                "Run: pip install llama-cpp-python"
            )

        n_gpu = 99 if self.device == "cuda" else 0
        self._llm = Llama(
            model_path=str(self.model_path),
            n_gpu_layers=n_gpu,
            n_ctx=4096,
        )
        self._loaded = True
        logger.info("LlamaCppTextEngine loaded %s on %s", self.model_path.name, self.device)

    def _translate_single(
        self,
        text: str,
        cefr_level: CEFRLevel,
        topic_description: str = "",
        target_language: str | None = None,
    ) -> str:
        """Translate a single English sentence with retry loop.

        Uses ``create_chat_completion`` so the model's chat template (with
        special tokens like ``<|USER_TOKEN|>``) is applied correctly. Raw
        prompt strings bypass the chat template and produce poor output.

        Wraps the LLM call in a retry loop (max 3 attempts). Returns the
        translated string or falls back to the original English text on failure.

        Args:
            text: Single English sentence to translate.
            cefr_level: CEFR proficiency level (linguistic guidance only).
            topic_description: Free-form description of topics/themes (for context).
            target_language: Override target language for this call only.
                Defaults to ``self.target_language`` from config.

        Returns:
            Translated string, or the original English text as fallback.
        """
        self._load_model()
        effective_lang = target_language or self.target_language
        base_messages = [
            {
                "role": "system",
                "content": (
                    f"You are a professional translator. Translate English sentences into "
                    f"{effective_lang} at the specified CEFR level. Output ONLY the translation, "
                    f"one line. No explanations, no notes, no source text repetition."
                ),
            }
        ]
        last_messages: list = []

        for attempt in range(1, 4):
            messages = list(base_messages)
            if attempt > 1 and last_messages:
                # Append failed output + retry instruction in conversation context
                messages.extend(last_messages)

            messages.append({
                "role": "user",
                "content": self._build_single_translation_prompt(
                    text, cefr_level, topic_description, effective_lang,
                ),
            })

            output = self._llm.create_chat_completion(
                messages=messages,
                max_tokens=128,
                temperature=0.3,
            )

            raw_text = output.get("choices", [{}])[0].get("message", {}).get("content", "")
            last_messages = [
                {"role": "assistant", "content": raw_text},
            ]
            line = raw_text.strip()

            # Validate: must be a single non-empty line, not repetitive garbage
            if line and self._is_valid_translation(line):
                logger.info(
                    "LlamaCppTextEngine: translated '%s' on attempt %d -> '%s'",
                    text[:30], attempt, line[:40],
                )
                return line

            # Invalid output — retry with stricter prompt in context
            if attempt < 3:
                last_messages.append({
                    "role": "user",
                    "content": (
                        f"That output was invalid. Translate ONLY this sentence into {effective_lang}:\n"
                        f"{text}\n\n"
                        f"Output ONE line only — the translation. Nothing else."
                    ),
                })
                logger.warning(
                    "LlamaCppTextEngine attempt %d: invalid output for '%s' — retrying",
                    attempt, text[:30],
                )
            else:
                logger.warning(
                    "LlamaCppTextEngine: exhausted retries for '%s'. Falling back to English.",
                    text[:30],
                )

        # Exhausted retries — fall back to original English text
        logger.info("LlamaCppTextEngine: fallback to English for '%s'", text[:30])
        return text

    def _is_valid_translation(self, line: str) -> bool:
        """Check if a translation output is valid (single line, not repetitive garbage).

        Args:
            line: The raw model output to validate.

        Returns:
            True if the output looks like a reasonable translation.
        """
        if not line or len(line) < 2:
            return False

        # Reject if it contains multiple lines (model generated too much)
        if "\n" in line:
            return False

        # Reject if it's just the English text back (no translation happened)
        lower = line.lower()
        if any(word in lower for word in ["translate", "translation", "english"]):
            return False

        # Reject very short outputs that are likely noise
        words = line.split()
        if len(words) < 1:
            return False

        return True

    def generate(
        self,
        texts: list[str],
        scenario: str,
        cefr_level: CEFRLevel,
        batch_size: int | None = None,
        topic_description: str = "",
    ) -> TextResult:
        """Generate translations using the loaded GGUF model.

        Translates each sentence individually for better quality with small models.
        Falls back to English text on failure (after max retries).

        Args:
            texts: English sentences to translate.
            scenario: Scenario/topic description (contextual).
            cefr_level: CEFR proficiency level (linguistic guidance only).
            batch_size: Number of translations expected.
            topic_description: Free-form description of topics/themes (contextual).

        Returns:
            TextResult with one translation per input text.

        Raises:
            ValidationError: If generation fails after max attempts and no lines produced.
        """
        self._load_model()
        if batch_size is None:
            raise ValueError("batch_size is required for translation")

        translations = []
        for text in texts:
            translated = self._translate_single(text, cefr_level, topic_description)
            translations.append(translated)

        return TextResult(generated_texts=translations)

    def _build_single_translation_prompt(
            self,
            text: str,
            cefr_level: CEFRLevel,
            topic_description: str = "",
            target_language: str | None = None,
    ) -> str:
        """Build prompt for translating a single sentence.

        Optimized for small models (tiny-aya-water ~3.3B params). Produces
        natural, idiomatic output by emphasizing how native speakers actually
        phrase things — not literal word-for-word translation.

        Uses CEFR linguistic guidance only — no hardcoded topics.

        Args:
            text: English sentence to translate.
            cefr_level: CEFR proficiency level.
            topic_description: Free-form context for the translation.
            target_language: Language to translate into. Defaults to ``self.target_language``.
        """
        target_lang = target_language or self.target_language
        cefr_desc = cefr_level.description()
        topic_hint = f" Context: {topic_description}." if topic_description else ""
        return (
            f"Translate the following English sentence into {target_lang}.{topic_hint}\n"
            f"CEFR linguistic guidance: {cefr_desc}.\n\n"
            f"CRITICAL — NATURAL LANGUAGE RULES:\n"
            f"1. Produce how a native speaker at this CEFR level would naturally express this idea in {target_lang}.\n"
            f"2. Do NOT translate word-for-word. Capture the meaning and rephrase it naturally in {target_lang}.\n"
            f"3. Use common idiomatic expressions, colloquial phrasing, and everyday vocabulary appropriate for the level.\n"
            f"4. Follow the grammar patterns typical of {target_lang} — not English sentence structure.\n"
            f"5. If the English uses an awkward or literal construction, render it as a native speaker would say it.\n\n"
            f"Rules:\n"
            f"1. Output ONLY the translated sentence — one line, nothing else.\n"
            f"2. Do NOT include explanations, notes, labels, or quotation marks.\n"
            f"3. Do NOT repeat the English text in your output.\n"
            f"4. Match the CEFR linguistic complexity for the target level.\n\n"
            f"English: {text}\n"
            f"{target_lang}:"
        )

    def unload(self) -> None:
        """Unload the model and free GPU memory."""
        if self._llm is not None:
            del self._llm
            self._llm = None
            self._loaded = False
            try:
                torch.cuda.empty_cache()
            except Exception:
                pass
            logger.info("LlamaCppTextEngine unloaded")


# TTSEngine has been extracted to core.audio_gen
from core.audio_gen import TTSEngine  # noqa: F401


# ImageGenEngine has been extracted to core.image_gen
from core.image_gen import ImageGenEngine  # noqa: F401


class EnginePool:
    """Singleton managing mutual exclusion between GPU inference engines.

    Ensures only one GPU model (LlamaCppTextEngine, TTSEngine, or ImageGenEngine)
    is loaded at a time. Text engines that use llama-cli subprocesses do not
    consume persistent VRAM.

    Usage:
        pool = EnginePool.get(config)
        text_result = pool.get_translation_engine().generate(texts, scenario, cefr_level)
        # ... later ...
        audio_result = pool.get_tts_engine().synthesize(translations, output_dir)
    """

    _instance: ClassVar[EnginePool | None] = None
    _config: EngineConfig
    _state: _EngineState

    def __new__(cls) -> EnginePool:
        if cls._instance is None:
            raise RuntimeError(
                "EnginePool must be created via EnginePool.get(config), not directly."
            )
        return cls._instance

    @classmethod
    def get(cls, config: EngineConfig) -> EnginePool:
        """Get or create the EnginePool singleton.

        Args:
            config: Validated engine configuration.

        Returns:
            The singleton EnginePool instance.
        """
        if cls._instance is None:
            instance = super().__new__(cls)
            instance._config = config
            instance._state = _EngineState()
            cls._instance = instance
            logger.info("EnginePool initialized (device=%s)", config.device)
        return cls._instance

    @classmethod
    def reset(cls) -> None:
        """Reset the singleton (useful for testing). Unloads all engines."""
        if cls._instance is not None:
            cls._instance._unload_translation()
            cls._instance._unload_tts()
            cls._instance._unload_image()
            cls._instance = None

    def get_english_engine(self) -> MiniCPMTextEngine:
        """Get a fresh English text generation engine (MiniCPM5-1B).

        Unloads any active GPU engines before returning.
        Returns a new MiniCPMTextEngine instance each call (stateless after unload).
        """
        self._ensure_exclusive("text")
        return MiniCPMTextEngine(
            model_path=self._config.minicpm_model_path,
            device=self._config.device,
        )

    def get_translation_engine(self) -> LlamaCppTextEngine:
        """Get or create the translation engine (tiny-aya-water via llama-cpp-python).

        Unloads any active GPU engines before loading. The same instance is returned
        on subsequent calls until explicitly unloaded.
        """
        self._ensure_exclusive("translation")
        if self._state.translation_engine is None:
            self._state.translation_engine = LlamaCppTextEngine(
                model_path=self._config.llm_model_path,
                device=self._config.device,
                target_language=self._config.target_language,
            )
        return self._state.translation_engine

    def get_tts_engine(self) -> TTSEngine:
        """Get or create the TTS engine.

        Unloads any active GPU engines before loading TTS.
        The same TTSEngine instance is returned on subsequent calls until unloaded.
        """
        self._ensure_exclusive("tts")
        if self._state.tts_engine is None:
            self._state.tts_engine = TTSEngine(device=self._config.device)
        return self._state.tts_engine

    def get_image_engine(self) -> ImageGenEngine:
        """Get or create the image generation engine.

        Unloads any active GPU engines before loading images.
        The same ImageGenEngine instance is returned on subsequent calls until unloaded.
        """
        self._ensure_exclusive("image")
        if self._state.image_engine is None:
            self._state.image_engine = ImageGenEngine(device=self._config.device)
        return self._state.image_engine

    def _ensure_exclusive(self, target: str) -> None:
        """Unload any active GPU engine that conflicts with the target."""
        if target == "text":
            self._unload_translation()
            self._unload_tts()
            self._unload_image()
            self._unload_english()
        elif target == "translation":
            self._unload_tts()
            self._unload_image()
        elif target == "tts":
            self._unload_translation()
            self._unload_image()
        elif target == "image":
            self._unload_translation()
            self._unload_tts()

    def _unload_translation(self) -> None:
        """Unload the translation engine if active."""
        if self._state.translation_engine is not None:
            self._state.translation_engine.unload()
            self._state.translation_engine = None

    def _unload_tts(self) -> None:
        """Unload the TTS engine if active."""
        if self._state.tts_engine is not None:
            self._state.tts_engine.unload()
            self._state.tts_engine = None

    def _unload_image(self) -> None:
        """Unload the image engine if active."""
        if self._state.image_engine is not None:
            self._state.image_engine.unload()
            self._state.image_engine = None

    def _unload_english(self) -> None:
        """Unload the English text engine if active."""
        # MiniCPMTextEngine instances are per-call (stateless), but we track
        # any loaded model state to ensure clean GPU memory.
        try:
            torch.cuda.empty_cache()
        except Exception:
            pass