"""Drop-in replacement for NAMHost using the C++ nam_engine subprocess.

Same interface as NAMHost (load_model, process, is_loaded, etc.) but
spawns the C++ NeuralAudio engine for ~34x faster inference.
"""

from __future__ import annotations

import logging
import os
import time
from pathlib import Path
from typing import Optional
from dataclasses import dataclass, field

import numpy as np

from .nam_engine import NAMEngineProcess

logger = logging.getLogger(__name__)

MODELS_DIR = Path(__file__).parent.parent / "models" / "nam"


@dataclass
class NAMFastModel:
    """Metadata matching the NAMModel dataclass from nam_host.py."""
    name: str
    path: str
    size_mb: float
    architecture: str = "LSTM"
    params_k: float = 0.0
    receptive_field: int = 0
    sample_rate: int = 48000
    compatible: bool = True

    @property
    def family(self) -> str:
        if self.size_mb < 0.1:
            return "nano"
        elif self.size_mb < 1.0:
            return "feather"
        elif self.size_mb < 4.0:
            return "lite"
        else:
            return "standard"

    @property
    def estimated_latency_ms(self) -> str:
        return "0.05-0.2 ms (C++ NeuralAudio engine)"


class FastNAMHost:
    """NAM model host using the C++ nam_engine subprocess.

    Same API surface as NAMHost (from nam_host.py), but uses the
    NeuralAudio C++ engine for much faster inference.

    Parameters
    ----------
    models_dir : str | Path
        Directory scanned for available .nam models.
    block_size : int
        Audio block size (must match the pipeline's JACK buffer).
    """

    def __init__(
        self,
        models_dir: str | Path = MODELS_DIR,
        block_size: int = 256,
    ):
        self._models_dir = Path(models_dir)
        self._block_size = block_size
        self._engine: Optional[NAMModel] = None  # Using current naming matching nam_host
        self._loaded_path: Optional[str] = None
        self._loaded_model: Optional[NAMFastModel] = None

        self._models_dir.mkdir(parents=True, exist_ok=True)

    # ── Properties ─────────────────────────────────────────────────

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

    @property
    def current_model(self) -> Optional[NAMFastModel]:
        return self._loaded_model

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

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

    def set_block_size(self, block_size: int) -> None:
        """Update block size. Reloads current model if loaded."""
        if block_size == self._block_size:
            return
        self._block_size = block_size
        if self._loaded_path:
            logger.info("Block size changed to %d — reloading model %s", block_size, self._loaded_path)
            self.load_model(self._loaded_path)

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

    def load_model(self, model_path: str) -> bool:
        """Load a .nam model into the C++ engine.

        Returns True on success, False on error.
        """
        path = Path(model_path)
        if not path.exists() or path.suffix.lower() not in (".nam",):
            logger.error("Model not found or invalid: %s", model_path)
            return False

        # Stop any existing engine
        self.unload()

        size_mb = path.stat().st_size / (1024 * 1024)

        # Create and start the engine
        engine = NAMEngineProcess(str(path), self._block_size)
        if not engine.start():
            logger.error("Failed to start NAM engine for: %s", model_path)
            return False

        self._engine = engine
        self._loaded_path = str(path)
        self._loaded_model = NAMFastModel(
            name=path.stem,
            path=str(path),
            size_mb=size_mb,
            architecture="LSTM",
        )

        logger.info(
            "Loaded NAM model via C++ engine: %s (%.1f KB, static=%s, engine=NeuralAudio)",
            path.stem,
            size_mb * 1024,
            engine.is_static,
        )
        return True

    def unload(self) -> None:
        """Unload the current model and stop the engine."""
        if self._engine is not None:
            self._engine.stop()
            self._engine = None
        self._loaded_path = None
        self._loaded_model = None
        logger.info("NAM model unloaded")

    # ── Warm-up ────────────────────────────────────────────────────

    def warm_up(self, block_size: int = 256) -> None:
        """Run a dry inference to warm caches."""
        if self._engine is None or not self._engine.is_loaded:
            return
        dummy = np.zeros(block_size, dtype=np.float32)
        for _ in range(5):
            self._engine.process(dummy)

    # ── Inference ──────────────────────────────────────────────────

    def process(self, audio_block: np.ndarray) -> np.ndarray:
        """Run a block of audio through the NAM model.

        Args:
            audio_block: float32 numpy array, shape (N,) or (1, N).

        Returns:
            Processed audio, same shape, float32.
        """
        if self._engine is None or not self._engine.is_loaded:
            return audio_block  # passthrough

        return self._engine.process(audio_block)

    # ── Model switching (crossfade compatible) ─────────────────────

    _crossfade_buf = None  # For pipeline crossfade compatibility

    def apply_crossfade(self, buf: np.ndarray) -> np.ndarray:
        """Passthrough — crossfade not needed with fast C++ switching."""
        return buf

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

    def list_available_models(self) -> list[NAMFastModel]:
        """Scan models_dir for .nam files and return metadata."""
        models: list[NAMFastModel] = []
        for f in sorted(self._models_dir.glob("*.nam")):
            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
