"""NAM engine router — switch between C++ subprocess and PyTorch in-process.

Provides a unified interface that wraps either FastNAMHost (C++ subprocess)
or NAMHost (PyTorch in-process) and allows runtime switching.

Usage:
    router = NAMEngineRouter(engine_mode='cpp')  # or 'pytorch'
    router.load_model('/path/to/model.nam')
    out = router.process(audio_block)
    router.set_engine('pytorch')  # switches at runtime, reloads current model
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Optional

import numpy as np

logger = logging.getLogger(__name__)


class NAMEngineRouter:
    """Wraps either FastNAMHost or NAMHost with a common interface.

    Parameters
    ----------
    engine_mode : str
        ``'cpp'`` for C++ subprocess (FastNAMHost), ``'pytorch'`` for
        in-process PyTorch (NAMHost).
    models_dir : str | Path
        Directory scanned for available .nam models.
    block_size : int
        Audio block size in samples.
    """

    ENGINE_MODES = ("cpp", "pytorch")

    def __init__(
        self,
        engine_mode: str = "cpp",
        models_dir: str | Path | None = None,
        block_size: int = 256,
    ):
        if engine_mode not in self.ENGINE_MODES:
            raise ValueError(f"engine_mode must be one of {self.ENGINE_MODES}, got {engine_mode!r}")

        self._models_dir = Path(models_dir) if models_dir else (
            Path(__file__).parent.parent / "models" / "nam"
        )
        self._block_size = block_size
        self._engine_mode = engine_mode
        self._engine: object = None  # FastNAMHost or NAMHost instance
        self._loaded_path: Optional[str] = None

        self._create_engine()

    # ── Engine lifecycle ────────────────────────────────────────────

    def _create_engine(self) -> None:
        """Create the engine instance for the current mode."""
        if self._engine_mode == "cpp":
            from .nam_engine_host import FastNAMHost
            self._engine = FastNAMHost(
                models_dir=str(self._models_dir),
                block_size=self._block_size,
            )
            logger.info("NAM engine: C++ subprocess (FastNAMHost)")
        else:
            from .nam_host import NAMHost, ModelSwitchMode
            self._engine = NAMHost(
                models_dir=str(self._models_dir),
                device="cpu",
                switch_mode=ModelSwitchMode.CROSSFADE,
            )
            logger.info("NAM engine: PyTorch in-process (NAMHost)")

    def set_engine(self, mode: str) -> bool:
        """Switch engine type at runtime.

        Unloads the current model and reloads it in the new engine.
        Returns True on success, False if the model couldn't be reloaded.
        """
        if mode == self._engine_mode:
            return True
        if mode not in self.ENGINE_MODES:
            raise ValueError(f"engine_mode must be one of {self.ENGINE_MODES}, got {mode!r}")

        # Save currently loaded model path
        old_path = self._loaded_path
        self.unload()

        self._engine_mode = mode
        self._create_engine()

        # Reload current model if one was loaded
        if old_path:
            logger.info("Reloading model %s in new engine %s", old_path, mode)
            return self.load_model(old_path)
        return True

    @property
    def engine_mode(self) -> str:
        """Current engine mode: 'cpp' or 'pytorch'."""
        return self._engine_mode

    # ── Properties delegated to engine ──────────────────────────────

    @property
    def is_loaded(self) -> bool:
        return self._engine is not None and self._engine.is_loaded

    @property
    def current_model(self):
        return getattr(self._engine, 'current_model', None) if self._engine else None

    @property
    def avg_inference_ms(self) -> float:
        if self._engine is None:
            return 0.0
        return getattr(self._engine, 'avg_inference_ms', 0.0)

    @property
    def block_size(self) -> int:
        return self._block_size

    # ── Crossfade (compatible with both engines) ────────────────────

    @property
    def _crossfade_buf(self):
        """For pipeline crossfade compatibility.
        PyTorch NAMHost has this natively; FastNAMHost has None."""
        if hasattr(self._engine, '_crossfade_buf'):
            return self._engine._crossfade_buf
        return None

    def apply_crossfade(self, buf: np.ndarray) -> np.ndarray:
        if hasattr(self._engine, 'apply_crossfade'):
            return self._engine.apply_crossfade(buf)
        return buf

    # ── Model loading ───────────────────────────────────────────────

    def load_model(self, model_path: str) -> bool:
        if self._engine is None:
            return False
        ok = self._engine.load_model(model_path)
        if ok:
            self._loaded_path = model_path
        return ok

    def unload(self) -> None:
        if self._engine is not None:
            self._engine.unload()
        self._loaded_path = None

    def set_block_size(self, block_size: int) -> None:
        self._block_size = block_size
        if self._engine is not None and hasattr(self._engine, 'set_block_size'):
            self._engine.set_block_size(block_size)

    def warm_up(self) -> None:
        if self._engine is not None and hasattr(self._engine, 'warm_up'):
            self._engine.warm_up(block_size=self._block_size)

    def process(self, audio_block: np.ndarray) -> np.ndarray:
        if self._engine is None:
            return audio_block
        return self._engine.process(audio_block)

    # ── Model discovery ─────────────────────────────────────────────

    def list_available_models(self):
        """List available .nam files from models_dir and ~/.pedal/models/."""
        models = []
        if self._engine is not None:
            models = self._engine.list_available_models()
        # Also scan the Tone3000 download directory for models not in
        # the primary models_dir
        from pathlib import Path
        extra = Path.home() / ".pedal" / "models"
        if extra != self._models_dir and extra.is_dir():
            from .nam_engine_host import NAMFastModel
            seen = set(m.path for m in models)
            for f in sorted(extra.glob("*.nam")):
                if str(f) not in seen:
                    size_mb = f.stat().st_size / (1024 * 1024)
                    models.append(NAMFastModel(
                        name=f.stem,
                        path=str(f),
                        size_mb=size_mb,
                        architecture="LSTM",
                    ))
        return models
