"""NAM A2 model host — load, configure, and run inference on RPi 4B.

Leverages the `neural-amp-modeler` (nam) Python package for model loading
and inference. Supports Linear, WaveNet, and LSTM architectures.

On RPi 4B:
- Python + PyTorch: good for Feather/Nano models only (~0.5-3ms at 256-block)
- For Standard/Lite models, use `neural-amp-modeler-lv2` compiled natively
  (NeuralAudio engine, LV2 plugin, ~1-2ms at 256-block)
- For A2 Slimmable runtime quality dialing, port to OpenSauce/nam-rs
"""

from __future__ import annotations

import json
import logging
import time
import warnings
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Optional

import numpy as np

logger = logging.getLogger(__name__)

DEFAULT_NAM_DIR = Path.home() / ".pedal" / "nam"

# ── Enum ────────────────────────────────────────────────────────────────


class ModelSwitchMode(Enum):
    """How to handle audio when switching between NAM models.

    INSTANT:  Immediate switch — may produce a click/pop.
    CROSSFADE: Smooth 256-sample fade between old and new output.
    PAUSE:    Brief silence (~1 block) during model load.
    """
    INSTANT = "instant"
    CROSSFADE = "crossfade"
    PAUSE = "pause"

    def __str__(self) -> str:
        return self.value


# ── Metadata ─────────────────────────────────────────────────────────────


@dataclass
class NAMModel:
    """Metadata for a loaded NAM model."""
    name: str
    path: str
    size_mb: float
    architecture: str
    params_k: float = 0.0
    receptive_field: int = 0
    sample_rate: int = 48000
    compatible: bool = True

    @property
    def family(self) -> str:
        """Categorise the model by file size."""
        if self.size_mb < 0.1:
            return "nano"
        elif self.size_mb < 1.0:
            return "feather"
        elif self.size_mb < 4.0:
            return "lite"
        elif self.size_mb < 10.0:
            return "standard"
        else:
            return "heavy"

    @property
    def estimated_latency_ms(self) -> str:
        """Return estimated per-block latency on RPi 4B at 256-block / 48 kHz."""
        estimates = {
            "nano": "0.1-0.2 ms (always safe)",
            "feather": "0.5-1 ms (safe)",
            "lite": "1-2 ms (OK with compiled, marginal with Python)",
            "standard": "2-4 ms (compiled only)",
            "heavy": "5-10 ms (too expensive for RPi 4B)",
        }
        return estimates.get(self.family, "unknown")


# ── Host ─────────────────────────────────────────────────────────────────


class NAMHost:
    """Hosts NAM models for real-time amp simulation.

    Loads ``.nam`` files (JSON format with weights) and runs inference
    through PyTorch.  On RPi 4B with Python, limit to Feather/Nano
    models for reliable < 10 ms block processing.

    Parameters
    ----------
    models_dir : str | Path
        Directory scanned for available models.
    device : str
        Torch device string (``'cpu'``, ``'cuda'``, …).
    switch_mode : ModelSwitchMode
        Transition style when switching models.
    """

    def __init__(
        self,
        models_dir: str | Path = DEFAULT_NAM_DIR,
        device: str = "cpu",
        switch_mode: ModelSwitchMode = ModelSwitchMode.CROSSFADE,
    ) -> None:
        self._models_dir = Path(models_dir)
        self._device = device
        self._switch_mode = switch_mode

        self._loaded_model: Optional[NAMModel] = None
        self._inference_model = None       # nn.Module instance
        self._torch = None                 # lazy import
        self._torch_device = None          # resolved device object

        # Timing stats
        self._timing_samples: list[float] = []

        # Crossfade state
        self._crossfade_len: int = 256
        self._crossfade_buf: Optional[np.ndarray] = None
        self._crossfade_pos: int = 0

        # Simple model cache (path -> model instance)
        self._model_cache: dict[str, object] = {}

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

        logger.info(
            "NAMHost(models_dir=%s, device=%s, switch_mode=%s)",
            self._models_dir, device, switch_mode,
        )

    # ── Lazy torch import ──────────────────────────────────────────

    def _import_torch(self):
        if self._torch is None:
            import torch
            self._torch = torch
            self._torch_device = torch.device(self._device)

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

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

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

    @property
    def avg_inference_ms(self) -> float:
        """Rolling average inference time in ms (last 50 blocks)."""
        if not self._timing_samples:
            return 0.0
        return float(np.mean(self._timing_samples[-50:]))

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

    def load_model(self, model_path: str) -> bool:
        """Load a ``.nam`` model file and build its inference model.

        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

        # Read file
        try:
            with open(path, "r") as f:
                data = json.load(f)
        except (json.JSONDecodeError, OSError) as e:
            logger.error("Failed to parse .nam file: %s", e)
            return False

        architecture = data.get("architecture", "Linear")
        config = data.get("config", {})
        size_mb = path.stat().st_size / (1024 * 1024)

        # Build metadata
        self._loaded_model = NAMModel(
            name=path.stem,
            path=str(path),
            size_mb=size_mb,
            architecture=architecture,
            receptive_field=config.get("receptive_field", 0),
        )

        # Build PyTorch inference model
        model_ok = self._build_inference(data)

        if model_ok:
            # Store param count in metadata
            params = sum(p.numel() for p in self._inference_model.parameters())
            self._loaded_model.params_k = round(params / 1000, 1)

        logger.info(
            "Loaded NAM model: %s (%.1f MB, %s, %s family, rf=%d, params=%.1fK)",
            self._loaded_model.name,
            size_mb,
            architecture,
            self._loaded_model.family,
            self._loaded_model.receptive_field,
            self._loaded_model.params_k,
        )
        return True

    def _build_inference(self, data: dict) -> bool:
        """Instantiate a PyTorch model from a NAM config dict.

        Uses the ``nam`` package's ``init_from_nam()`` factory.
        Falls back gracefully if the package is unavailable or the
        architecture is unsupported.
        """
        self._import_torch()

        path = data.get("path", "")
        cache_key = str(path)

        # Check cache first
        if cache_key and cache_key in self._model_cache:
            self._inference_model = self._model_cache[cache_key]
            self._inference_model.eval()
            return True

        try:
            from nam.models import init_from_nam

            model = init_from_nam(data)
            model.eval()

            # Move to target device
            if str(self._torch_device) != "cpu":
                model = model.to(self._torch_device)

            # Cache it
            if cache_key:
                self._model_cache[cache_key] = model

            self._inference_model = model
            return True

        except ImportError:
            logger.warning("nam package not installed; inference unavailable")
            self._inference_model = None
            return False
        except Exception as exc:
            logger.warning("Failed to build inference model: %s", exc)
            self._inference_model = None
            return False

    def unload(self) -> None:
        """Unload the current model and free resources."""
        self._loaded_model = None
        if self._inference_model is not None:
            self._inference_model = self._inference_model.to("cpu")
            del self._inference_model
            self._inference_model = None
        self._torch = None
        self._torch_device = None
        self._timing_samples.clear()
        self._crossfade_buf = None
        self._crossfade_pos = 0
        logger.info("NAM model unloaded")

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

    def warm_up(self, block_size: int = 256) -> None:
        """Run a dry inference to warm caches / GPU.

        Safe to call even when no model is loaded.
        """
        if self._inference_model is None:
            return
        self._import_torch()
        try:
            dummy = self._torch.randn(1, block_size, device=self._torch_device)
            with self._torch.no_grad():
                self._inference_model(dummy)
            logger.debug("Model warm-up complete (block=%d)", block_size)
        except Exception as exc:
            logger.warning("Warm-up failed (non-fatal): %s", exc)

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

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

        Handles 1-D (``(N,)``) and 2-D (``(1, N)``) float32 input.
        If no model is loaded, passes audio through unchanged.

        Parameters
        ----------
        audio_block : np.ndarray
            PCM samples, float32 in [-1, 1].

        Returns
        -------
        np.ndarray
            Processed audio, same shape.
        """
        if self._inference_model is None:
            # Pass-through when no model is loaded
            return audio_block

        self._import_torch()
        start = time.perf_counter()

        with self._torch.no_grad():
            x = self._torch.from_numpy(audio_block.astype(np.float32))

            # Ensure shape (1, T) for ConvNet-like models
            was_1d = x.dim() == 1
            if was_1d:
                x = x.unsqueeze(0)

            # Move to device
            if str(self._torch_device) != "cpu":
                x = x.to(self._torch_device)

            y = self._inference_model(x)

            # Squeeze back
            if was_1d:
                y = y.squeeze(0)

            # Move back to CPU numpy
            if str(self._torch_device) != "cpu":
                y = y.cpu()
            y = y.numpy()

        # Trim / pad to match input length
        if y.shape[-1] > audio_block.shape[-1]:
            y = y[..., :audio_block.shape[-1]]
        elif y.shape[-1] < audio_block.shape[-1]:
            pad_len = audio_block.shape[-1] - y.shape[-1]
            y = np.pad(y, ((0, 0),) * (y.ndim - 1) + ((0, pad_len),), mode="constant")

        # Track timing
        elapsed_ms = (time.perf_counter() - start) * 1000
        self._timing_samples.append(elapsed_ms)
        if len(self._timing_samples) > 200:
            self._timing_samples = self._timing_samples[-100:]

        return y.astype(np.float32)

    # ── Model switching ────────────────────────────────────────────

    def switch_model(
        self,
        model_path: str,
        mode: Optional[ModelSwitchMode] = None,
    ) -> bool:
        """Load a new model with a configurable transition.

        Parameters
        ----------
        model_path : str
            Path to the new ``.nam`` file.
        mode : ModelSwitchMode or None
            Override the default switch mode for this call.

        Returns
        -------
        bool
            ``True`` on success.
        """
        effective_mode = mode if mode is not None else self._switch_mode

        if effective_mode == ModelSwitchMode.PAUSE:
            # Replace immediately (caller is expected to mute output briefly)
            return self.load_model(model_path)

        if effective_mode == ModelSwitchMode.CROSSFADE and self._inference_model is not None:
            # Snapshot the current output buffer for crossfade
            if self._crossfade_buf is None:
                self._crossfade_buf = np.zeros(self._crossfade_len, dtype=np.float32)
            self._crossfade_pos = 0

        return self.load_model(model_path)

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

    def list_available_models(self) -> list[NAMModel]:
        """Scan ``models_dir`` for ``.nam`` files and return metadata."""
        models: list[NAMModel] = []
        for f in sorted(self._models_dir.glob("*.nam")):
            try:
                with open(f) as fp:
                    data = json.load(fp)
                size_mb = f.stat().st_size / (1024 * 1024)
                models.append(
                    NAMModel(
                        name=f.stem,
                        path=str(f),
                        size_mb=size_mb,
                        architecture=data.get("architecture", "?"),
                        receptive_field=data.get("config", {}).get("receptive_field", 0),
                    )
                )
            except (json.JSONDecodeError, OSError) as exc:
                logger.debug("Skipping %s: %s", f.name, exc)
        return models

    # ── Crossfade helper (used by pipeline) ─────────────────────────

    def apply_crossfade(self, buf: np.ndarray) -> np.ndarray:
        """Apply a smooth crossfade from the model-switch boundary.

        Intended to be called by the pipeline after ``process()`` when
        a crossfade is active.
        """
        if self._crossfade_buf is None or self._crossfade_pos >= self._crossfade_len:
            return buf

        remain = min(self._crossfade_len - self._crossfade_pos, len(buf))
        fade_out = np.cos(np.linspace(0, np.pi / 2, remain)) ** 2
        fade_in = np.sin(np.linspace(0, np.pi / 2, remain)) ** 2

        buf[:remain] = (
            fade_out * self._crossfade_buf[:remain]
            + fade_in * buf[:remain]
        )

        self._crossfade_pos += remain
        return buf


# ── Standalone helpers ────────────────────────────────────────────────────


def available_models(models_dir: str | Path = DEFAULT_NAM_DIR) -> list[dict]:
    """Lightweight model listing — returns dicts (lighter than NAMModel).

    Returns
    -------
    list[dict]
        Each entry: ``{name, path, architecture, size_mb, feather, receptive_field}``.
    """
    host = NAMHost(models_dir=models_dir)
    results: list[dict] = []
    for m in host.list_available_models():
        results.append({
            "name": m.name,
            "path": m.path,
            "architecture": m.architecture,
            "size_mb": round(m.size_mb, 2),
            "feather": m.size_mb < 10.0,
            "receptive_field": m.receptive_field,
        })
    return results


def process_with_model(
    model_path: str,
    audio: np.ndarray,
    device: str = "cpu",
) -> np.ndarray:
    """Convenience: load a model, process one block, return audio."""
    host = NAMHost(device=device)
    host.load_model(model_path)
    return host.process(audio)