"""Pipeline orchestrating STT -> LLM -> TTS."""

from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncIterator
from typing import Any

from roohai.audio import audio_to_wav_bytes, read_audio
from roohai.base import STTModel, TTSModel, LLMModel, VADModel
from roohai.config import load_config
from roohai.registry import registry

logger = logging.getLogger(__name__)

# Class-to-registry-name mapping for config-based instantiation.
# Maps the "class" field in agent config to the base class key used in _register_defaults().
# Variant names (e.g. "whisper-tiny") also map back to the base class ("whisper").
_CLASS_MAP: dict[str, dict[str, str]] = {
    "stt": {
        "whisper": "whisper", "wav2vec2": "wav2vec2", "deepgram": "deepgram",
        "nvidia-canary": "nvidia-canary", "nvidia-parakeet": "nvidia-parakeet",
    },
    "tts": {
        "speecht5": "speecht5", "bark": "bark", "piper": "piper",
        "cartesia": "cartesia", "deepgram-tts": "deepgram-tts",
    },
    "llm": {
        "bedrock": "bedrock", "local": "local", "openai": "openai",
        "anthropic": "anthropic", "gemini": "gemini", "ollama": "ollama",
    },
    "vad": {
        "silero": "silero",
    },
}


def _register_defaults() -> None:
    """Register all built-in model classes.

    Gracefully skips models whose optional dependencies are not installed.
    """
    from roohai.stt import _lazy_load as _stt_lazy
    from roohai.tts import _lazy_load as _tts_lazy
    from roohai.llm import _lazy_load as _llm_lazy
    from roohai.vad import _lazy_load as _vad_lazy

    stt_available = _stt_lazy()
    tts_available = _tts_lazy()
    llm_available = _llm_lazy()
    vad_available = _vad_lazy()

    _stt_map = {
        "whisper": "WhisperSTT", "wav2vec2": "Wav2Vec2STT", "deepgram": "DeepgramSTT",
        "nvidia-canary": "NvidiaCanarySTT", "nvidia-parakeet": "NvidiaParakeetSTT",
    }
    _tts_map = {
        "speecht5": "SpeechT5TTS", "bark": "BarkTTS", "piper": "PiperTTS",
        "cartesia": "CartesiaTTS", "deepgram-tts": "DeepgramTTS",
    }
    _llm_map = {
        "bedrock": "BedrockLLM", "local": "LocalLLM",
        "openai": "OpenAILLM", "anthropic": "AnthropicLLM",
        "gemini": "GeminiLLM", "ollama": "OllamaLLM",
    }
    _vad_map = {"silero": "SileroVAD"}

    _stt_classes = {k: stt_available[v] for k, v in _stt_map.items() if v in stt_available}
    _tts_classes = {k: tts_available[v] for k, v in _tts_map.items() if v in tts_available}
    _llm_classes = {k: llm_available[v] for k, v in _llm_map.items() if v in llm_available}
    _vad_classes = {k: vad_available[v] for k, v in _vad_map.items() if v in vad_available}

    return _stt_classes, _tts_classes, _llm_classes, _vad_classes


def register_all_models() -> None:
    """Register every built-in model and its variants."""
    stt_classes, tts_classes, llm_classes, vad_classes = _register_defaults()

    # --- STT variants ---
    WhisperSTT = stt_classes["whisper"]
    Wav2Vec2STT = stt_classes["wav2vec2"]
    DeepgramSTT = stt_classes["deepgram"]

    registry.register_stt("whisper-tiny", WhisperSTT, meta={
        "display_name": "Whisper Tiny",
        "description": "Fastest Whisper model. Good for low-latency, English-focused tasks.",
        "estimated_download_mb": 150,
        "hf_models": ["openai/whisper-tiny"],
        "provider": "whisper",
    }, model_id="openai/whisper-tiny")

    registry.register_stt("whisper-base", WhisperSTT, meta={
        "display_name": "Whisper Base",
        "description": "Small Whisper model. Better accuracy than Tiny with modest resource use.",
        "estimated_download_mb": 290,
        "hf_models": ["openai/whisper-base"],
        "provider": "whisper",
    }, model_id="openai/whisper-base")

    registry.register_stt("whisper-small", WhisperSTT, meta={
        "display_name": "Whisper Small",
        "description": "Mid-size Whisper. Strong multilingual accuracy, reasonable speed.",
        "estimated_download_mb": 950,
        "hf_models": ["openai/whisper-small"],
        "provider": "whisper",
    }, model_id="openai/whisper-small")

    registry.register_stt("whisper-medium", WhisperSTT, meta={
        "display_name": "Whisper Medium",
        "description": "Large Whisper model. High accuracy, requires more memory.",
        "estimated_download_mb": 3000,
        "hf_models": ["openai/whisper-medium"],
        "provider": "whisper",
    }, model_id="openai/whisper-medium")

    registry.register_stt("whisper-large", WhisperSTT, meta={
        "display_name": "Whisper Large V3",
        "description": "Largest Whisper model. Best accuracy, needs GPU with 10GB+ VRAM.",
        "estimated_download_mb": 6000,
        "hf_models": ["openai/whisper-large-v3"],
        "provider": "whisper",
    }, model_id="openai/whisper-large-v3")

    registry.register_stt("wav2vec2", Wav2Vec2STT, meta={
        "display_name": "Wav2Vec2 Base",
        "description": "Facebook's speech recognition model. English-only, runs locally.",
        "provider": "wav2vec2",
    })

    registry.register_stt("wav2vec2-large", Wav2Vec2STT, meta={
        "display_name": "Wav2Vec2 Large",
        "description": "Larger Wav2Vec2 model. Better accuracy, more resource-intensive.",
        "estimated_download_mb": 1200,
        "hf_models": ["facebook/wav2vec2-large-960h"],
        "provider": "wav2vec2",
    }, model_id="facebook/wav2vec2-large-960h")

    registry.register_stt("deepgram-nova-3", DeepgramSTT, meta={
        "display_name": "Deepgram Nova-3",
        "description": "Latest Deepgram model. Best accuracy and speed. Cloud, requires API key.",
        "provider": "deepgram",
    }, model="nova-3")

    registry.register_stt("deepgram-nova-2", DeepgramSTT, meta={
        "display_name": "Deepgram Nova-2",
        "description": "Previous-gen Deepgram model. Stable and reliable. Cloud, requires API key.",
        "provider": "deepgram",
    }, model="nova-2")

    NvidiaCanarySTT = stt_classes.get("nvidia-canary")
    NvidiaParakeetSTT = stt_classes.get("nvidia-parakeet")

    if NvidiaCanarySTT:
        registry.register_stt("nvidia-canary", NvidiaCanarySTT, meta={
            "display_name": "NVIDIA Canary-Qwen 2.5B",
            "description": "NVIDIA Canary-Qwen 2.5B speech model. Multilingual, requires CUDA GPU.",
            "estimated_download_mb": 5000,
            "provider": "nvidia-canary",
        }, model_id="nvidia/canary-qwen-2.5b")

    if NvidiaParakeetSTT:
        registry.register_stt("nvidia-parakeet", NvidiaParakeetSTT, meta={
            "display_name": "NVIDIA Parakeet TDT 0.6B V2",
            "description": "NVIDIA Parakeet TDT 0.6B V2. Fast, accurate, supports timestamps. Requires CUDA GPU.",
            "estimated_download_mb": 2400,
            "provider": "nvidia-parakeet",
        }, model_id="nvidia/parakeet-tdt-0.6b-v2")

    # --- TTS variants ---
    SpeechT5TTS = tts_classes.get("speecht5")
    BarkTTS = tts_classes.get("bark")
    PiperTTS = tts_classes.get("piper")
    CartesiaTTS = tts_classes.get("cartesia")

    if SpeechT5TTS:
        registry.register_tts("speecht5", SpeechT5TTS, meta={
            "display_name": "SpeechT5",
            "description": "Microsoft's text-to-speech model. Fast, lightweight, runs locally.",
            "provider": "speecht5",
        })

    if BarkTTS:
        registry.register_tts("bark-small", BarkTTS, meta={
            "display_name": "Bark Small",
            "description": "Suno's generative TTS (small). Expressive speech, moderate resource use.",
            "estimated_download_mb": 1800,
            "hf_models": ["suno/bark-small"],
            "provider": "bark",
        }, model_id="suno/bark-small")

        registry.register_tts("bark-large", BarkTTS, meta={
            "display_name": "Bark Large",
            "description": "Suno's generative TTS (full). Most expressive, needs GPU with 8GB+ VRAM.",
            "estimated_download_mb": 5000,
            "hf_models": ["suno/bark"],
            "provider": "bark",
        }, model_id="suno/bark")

    if PiperTTS:
        registry.register_tts("piper", PiperTTS, meta={
            "display_name": "Piper",
            "description": "Fast, lightweight local TTS. Multiple voices available.",
            "provider": "piper",
        })

    if CartesiaTTS:
        registry.register_tts("cartesia", CartesiaTTS, meta={
            "display_name": "Cartesia Sonic",
            "description": "Cloud TTS with fast, natural voices. Requires API key.",
            "provider": "cartesia",
        })

    DeepgramTTS = tts_classes.get("deepgram-tts")

    if DeepgramTTS:
        registry.register_tts("deepgram-aura-2", DeepgramTTS, meta={
            "display_name": "Deepgram Aura-2",
            "description": "Cloud TTS with 90+ natural voices across 7 languages. Requires API key.",
            "provider": "deepgram-tts",
        })

    # --- LLM variants ---
    for name, cls in llm_classes.items():
        registry.register_llm(name, cls)

    # --- VAD ---
    for name, cls in vad_classes.items():
        registry.register_vad(name, cls)


def _register_from_config(config: dict[str, Any]) -> None:
    """Register models defined in config.

    The "class" field can be either a base class name (e.g. "whisper") or a
    variant name (e.g. "whisper-tiny").  We resolve variants through _CLASS_MAP.
    """
    stt_classes, tts_classes, llm_classes, vad_classes = _register_defaults()
    models_cfg = config.get("models", {})

    def _resolve(class_name: str, model_type: str, class_dict: dict) -> type | None:
        """Resolve a class name (possibly a variant) to the actual model class."""
        if class_name in class_dict:
            return class_dict[class_name]
        # Try resolving through _CLASS_MAP (e.g. "whisper-tiny" -> "whisper")
        base = _CLASS_MAP.get(model_type, {}).get(class_name)
        if base and base in class_dict:
            return class_dict[base]
        return None

    for name, model_cfg in models_cfg.items():
        model_type = model_cfg.get("type")
        class_name = model_cfg.get("class")
        kwargs = {k: v for k, v in model_cfg.items() if k not in ("type", "class")}

        # If the name already exists in the registry (from register_all_models)
        # and no extra kwargs are provided, skip — the existing registration
        # already has the correct model_id and other variant-specific params.
        # If extra kwargs are present (e.g., api_key from secrets), merge them
        # on top of the existing registration's kwargs.
        store_map = {"stt": registry._stt, "tts": registry._tts, "llm": registry._llm, "vad": registry._vad}
        store = store_map.get(model_type)
        if store is not None and name in store and not kwargs:
            continue  # already registered with correct variant kwargs

        if store is not None and name in store and kwargs:
            # Merge agent config kwargs on top of existing registration
            existing_cls, existing_kwargs, existing_meta = store[name]
            merged = {**existing_kwargs, **kwargs}
            register_fn = {
                "stt": registry.register_stt,
                "tts": registry.register_tts,
                "llm": registry.register_llm,
                "vad": registry.register_vad,
            }[model_type]
            register_fn(name, existing_cls, meta=existing_meta, **merged)
            continue

        # New name not in registry — resolve class and register.
        # If class_name is itself a registered variant (e.g. "whisper-large"),
        # inherit its kwargs (model_id, etc.) as the base.
        if store is not None and class_name in store:
            base_cls, base_kwargs, _base_meta = store[class_name]
            merged = {**base_kwargs, **kwargs}
            register_fn = {
                "stt": registry.register_stt,
                "tts": registry.register_tts,
                "llm": registry.register_llm,
                "vad": registry.register_vad,
            }[model_type]
            register_fn(name, base_cls, **merged)
        elif model_type == "stt":
            cls = _resolve(class_name, "stt", stt_classes)
            if cls:
                registry.register_stt(name, cls, **kwargs)
        elif model_type == "tts":
            cls = _resolve(class_name, "tts", tts_classes)
            if cls:
                registry.register_tts(name, cls, **kwargs)
        elif model_type == "llm":
            cls = _resolve(class_name, "llm", llm_classes)
            if cls:
                registry.register_llm(name, cls, **kwargs)
        elif model_type == "vad":
            cls = _resolve(class_name, "vad", vad_classes)
            if cls:
                registry.register_vad(name, cls, **kwargs)


class Rooh:
    def __init__(self, config_path: str = "config.yaml", config_dict: dict | None = None) -> None:
        if config_dict is not None:
            self.config = config_dict
        else:
            self.config = load_config(config_path)
        register_all_models()
        _register_from_config(self.config)

        pipeline_cfg = self.config.get("pipeline", {})
        self._stt_name = pipeline_cfg.get("stt")
        self._tts_name = pipeline_cfg.get("tts")
        self._llm_name = pipeline_cfg.get("llm")
        self._vad_name = pipeline_cfg.get("vad")

        self.stt: STTModel | None = None
        self.tts: TTSModel | None = None
        self.llm: LLMModel | None = None
        self.vad: VADModel | None = None
        self.system_prompt: str | None = self.config.get("system_prompt")
        self.llm_streaming: bool = self.config.get("llm_streaming", True)
        self.vad_threshold: float = float(pipeline_cfg.get("vad_threshold", 0.7))
        self.barge_in_enabled: bool = pipeline_cfg.get("barge_in_enabled", True)
        self.silence_duration_ms: int = int(pipeline_cfg.get("silence_duration_ms", 1000))
        self._llm_hook = None
        self._llm_stream_hook = None
        self._barge_in_hook = None
        self._session_disconnect_hook = None

    @staticmethod
    def _check_model_dependencies(model_type: str, model_name: str) -> None:
        """Ensure pip dependencies for a model are installed before loading.

        Runs ``pip install`` with the full version spec for each declared
        dependency.  Pip skips already-satisfied packages instantly, and
        upgrades outdated ones automatically.
        """
        import subprocess
        import sys

        meta = registry.get_meta(model_type, model_name)
        pip_deps = meta.get("pip_dependencies", [])
        if not pip_deps:
            return

        logger.info(
            "Ensuring dependencies for %s model '%s': %s",
            model_type, model_name, pip_deps,
        )
        try:
            subprocess.check_call(
                [sys.executable, "-m", "pip", "install", "--quiet", *pip_deps],
                stdout=subprocess.DEVNULL,
            )
        except subprocess.CalledProcessError as exc:
            raise RuntimeError(
                f"Failed to install dependencies for {model_type} model '{model_name}': {pip_deps}. "
                f"Install manually with: pip install {' '.join(pip_deps)}"
            ) from exc

    def load(self) -> None:
        """Load all configured models.

        Before loading each model, checks that its declared pip dependencies
        are installed.  Raises RuntimeError with install instructions if any
        are missing.
        """
        models_to_load = [
            ("stt", self._stt_name, registry.create_stt, "stt"),
            ("tts", self._tts_name, registry.create_tts, "tts"),
            ("llm", self._llm_name, registry.create_llm, "llm"),
            ("vad", self._vad_name, registry.create_vad, "vad"),
        ]

        for model_type, model_name, create_fn, attr in models_to_load:
            if not model_name:
                continue
            self._check_model_dependencies(model_type, model_name)
            logger.info("Loading %s model: %s", model_type.upper(), model_name)
            try:
                model = create_fn(model_name)
                model.load()
            except RuntimeError:
                raise
            except Exception as exc:
                logger.error(
                    "Failed to load %s model '%s': %s", model_type, model_name, exc,
                )
                raise RuntimeError(
                    f"Failed to load {model_type} model '{model_name}': {exc}"
                ) from exc
            setattr(self, attr, model)
            logger.info("Loaded %s model: %s", model_type.upper(), model_name)

    def set_system_prompt(self, prompt: str | None) -> None:
        """Set or clear the system prompt for LLM calls."""
        self.system_prompt = prompt

    def set_llm_hook(self, hook, stream_hook=None) -> None:
        """Set or clear the LLM hook.

        When set, the hook replaces the built-in LLM call in :meth:`chat`.
        Pass ``None`` to restore default behaviour.

        Args:
            hook: Async callable ``(message, *, history, system_prompt) -> str``.
            stream_hook: Optional async generator ``(message, *, history, system_prompt)``
                that yields ``str`` chunks.  When provided and ``llm_streaming``
                is ``True``, :meth:`chat_stream` will use it instead of falling
                back to the batch hook.
        """
        self._llm_hook = hook
        self._llm_stream_hook = stream_hook

    def set_barge_in_hook(self, hook) -> None:
        """Set or clear the barge-in hook.

        When set, the hook is called whenever a barge-in event occurs.
        Pass ``None`` to remove the hook.

        Args:
            hook: Async callable ``(session_id, generation_id) -> None``.
        """
        self._barge_in_hook = hook

    def set_session_disconnect_hook(self, hook) -> None:
        """Set or clear the session disconnect hook.

        Called when a WebSocket/WebRTC session ends. Used by the Strands
        hook to clean up per-session agent memory.

        Args:
            hook: Callable ``(session_id: str) -> None``.
        """
        self._session_disconnect_hook = hook

    def on_session_disconnect(self, session_id: str) -> None:
        """Notify the pipeline that a session has disconnected.

        Called by the streaming/webrtc handlers on connection close.
        """
        if self._session_disconnect_hook is not None:
            try:
                self._session_disconnect_hook(session_id)
            except Exception:
                logger.debug("Session disconnect hook error for %s", session_id, exc_info=True)

    async def process_audio(self, audio_bytes: bytes, system_prompt: str | None = None) -> tuple[str, str, bytes]:
        """Full pipeline: audio -> transcription -> LLM response -> speech audio bytes."""
        transcription = await self.transcribe(audio_bytes)
        response_text = await self.chat(transcription, system_prompt=system_prompt)
        response_audio = await self.synthesize(response_text)
        return transcription, response_text, response_audio

    async def transcribe(self, audio_bytes: bytes) -> str:
        """Transcribe audio bytes to text."""
        if self.stt is None:
            raise RuntimeError("No STT model loaded")
        sample_rate = self.config.get("audio", {}).get("sample_rate", 16000)
        audio = await asyncio.to_thread(read_audio, audio_bytes, sample_rate)
        return await asyncio.to_thread(self.stt.transcribe, audio, sample_rate)

    async def chat(
        self, text: str, history: list[dict] | None = None,
        system_prompt: str | None = None, session_id: str | None = None,
    ) -> str:
        """Send text to LLM and get response.

        By default routes through the Strands Agent hook (set during agent
        activation).  Falls back to the raw LLM model if no hook is set.
        """
        prompt = system_prompt or self.system_prompt
        if self._llm_hook is not None:
            return await self._llm_hook(text, history=history, system_prompt=prompt, session_id=session_id)
        # ── Legacy direct LLM path (kept for reference / fallback) ───────
        # If no Strands hook is configured, fall back to raw LLM model.
        # This path is used when Strands is explicitly disabled or for
        # local-only models that don't support tool use.
        if self.llm is None:
            raise RuntimeError("No LLM model loaded and no Strands hook configured")
        return await self.llm.chat(text, history=history, system_prompt=prompt, session_id=session_id)

    async def chat_stream(
        self, text: str, history: list[dict] | None = None,
        system_prompt: str | None = None, session_id: str | None = None,
    ) -> AsyncIterator[str]:
        """Yield LLM response chunks.

        By default routes through the Strands Agent hook.  Falls back to
        stream hook -> batch hook -> raw LLM model.
        """
        prompt = system_prompt or self.system_prompt
        if self._llm_stream_hook is not None:
            async for chunk in self._llm_stream_hook(
                text, history=history, system_prompt=prompt, session_id=session_id,
            ):
                yield chunk
            return
        if self._llm_hook is not None:
            result = await self._llm_hook(text, history=history, system_prompt=prompt, session_id=session_id)
            yield result
            return
        # ── Legacy direct LLM streaming path (kept for reference / fallback) ─
        # if self.llm is None:
        #     raise RuntimeError("No LLM model loaded")
        # async for chunk in self.llm.chat_stream(text, history=history, system_prompt=prompt, session_id=session_id):
        #     yield chunk
        if self.llm is None:
            raise RuntimeError("No LLM model loaded and no Strands hook configured")
        async for chunk in self.llm.chat_stream(text, history=history, system_prompt=prompt, session_id=session_id):
            yield chunk

    async def synthesize(self, text: str) -> bytes:
        """Synthesize text to WAV audio bytes."""
        if self.tts is None:
            raise RuntimeError("No TTS model loaded")
        audio_array, sample_rate = await asyncio.to_thread(self.tts.synthesize, text)
        return audio_to_wav_bytes(audio_array, sample_rate)

    def swap_model(self, model_type: str, model_name: str) -> None:
        """Hot-swap a model at runtime."""
        self._check_model_dependencies(model_type, model_name)

        create_fns = {
            "stt": registry.create_stt,
            "tts": registry.create_tts,
            "llm": registry.create_llm,
            "vad": registry.create_vad,
        }
        if model_type not in create_fns:
            raise ValueError(f"Unknown model type: {model_type}")

        current = getattr(self, model_type)
        if current is not None:
            current.unload()

        logger.info("Loading %s model: %s", model_type.upper(), model_name)
        try:
            model = create_fns[model_type](model_name)
            model.load()
        except RuntimeError:
            raise
        except Exception as exc:
            logger.error(
                "Failed to swap %s model '%s': %s", model_type, model_name, exc,
            )
            raise RuntimeError(
                f"Failed to swap {model_type} model '{model_name}': {exc}"
            ) from exc

        setattr(self, model_type, model)
        setattr(self, f"_{model_type}_name", model_name)
        logger.info("Swapped %s model to: %s", model_type.upper(), model_name)

    @property
    def stt_supports_streaming(self) -> bool:
        """Check if the current STT model supports streaming."""
        return self.stt is not None and getattr(self.stt, "META", {}).get("streaming", False)

    @classmethod
    def from_config_dict(cls, config: dict[str, Any]) -> "Rooh":
        """Create a Rooh instance from an in-memory config dict (no file needed)."""
        instance = object.__new__(cls)
        instance.config = config
        _register_from_config(config)

        pipeline_cfg = config.get("pipeline", {})
        instance._stt_name = pipeline_cfg.get("stt")
        instance._tts_name = pipeline_cfg.get("tts")
        instance._llm_name = pipeline_cfg.get("llm")
        instance._vad_name = pipeline_cfg.get("vad")

        instance.stt = None
        instance.tts = None
        instance.llm = None
        instance.vad = None
        instance.system_prompt = config.get("system_prompt")
        instance.llm_streaming = config.get("llm_streaming", True)
        instance.vad_threshold = float(pipeline_cfg.get("vad_threshold", 0.7))
        instance.silence_duration_ms = int(pipeline_cfg.get("silence_duration_ms", 1000))
        instance.barge_in_enabled = pipeline_cfg.get("barge_in_enabled", True)
        instance._llm_hook = None
        instance._llm_stream_hook = None
        instance._barge_in_hook = None
        instance._session_disconnect_hook = None
        return instance

    @classmethod
    def builder(cls) -> "RoohBuilder":
        """Return a fluent builder for constructing a Rooh pipeline in code.

        Example::

            pipeline = (
                Rooh.builder()
                .stt("deepgram", model_id="nova-3", language="en", api_key="dg-xxx")
                .llm("bedrock", model_id="anthropic.claude-3-5-sonnet-20241022-v2:0", region="us-west-2")
                .tts("cartesia", voice_id="79a125e8-...", api_key="sk-xxx")
                .vad("silero")
                .vad_threshold(0.7)
                .barge_in(True)
                .silence_duration(1000)
                .system_prompt("You are a helpful voice assistant.")
                .build()
            )
            pipeline.load()
        """
        return RoohBuilder()


class RoohBuilder:
    """Fluent builder for constructing a :class:`Rooh` pipeline entirely in code.

    Each method returns ``self`` so calls can be chained.  Call :meth:`build`
    at the end to get a configured (but not yet loaded) ``Rooh`` instance.
    """

    def __init__(self) -> None:
        self._pipeline_cfg: dict[str, Any] = {}
        self._models_cfg: dict[str, dict[str, Any]] = {}
        self._system_prompt: str | None = None
        self._llm_hook = None
        self._llm_stream_hook = None
        self._barge_in_hook = None
        self._session_disconnect_hook = None
        # Custom model classes registered via custom_stt/tts/llm/vad.
        # List of (name, cls, kwargs) tuples to register before build.
        self._custom_models: list[tuple[str, str, type, dict[str, Any]]] = []

    # ── Model stages ──────────────────────────────────────────────────────

    def stt(self, provider: str, *, name: str | None = None, **kwargs) -> "RoohBuilder":
        """Add a speech-to-text model.

        Args:
            provider: Model class — one of:

                ``"whisper"``
                    model_id (str): HuggingFace model ID. Default ``"openai/whisper-tiny"``.
                    language (str): Language code. Default ``"en"``.

                ``"deepgram"``
                    model_id (str): Deepgram model. Default ``"nova-3"``.
                    language (str): Language code. Default ``"en"``.
                    api_key (str): Deepgram API key. Falls back to ``DEEPGRAM_API_KEY`` env var.

            name: Optional registry name. Defaults to *provider*.
            **kwargs: Passed directly to the model constructor.
        """
        model_name = name or provider
        self._models_cfg[model_name] = {"type": "stt", "class": provider, **kwargs}
        self._pipeline_cfg["stt"] = model_name
        return self

    def llm(self, provider: str, *, name: str | None = None, **kwargs) -> "RoohBuilder":
        """Add a large language model.

        Args:
            provider: Model class — one of:

                ``"bedrock"``
                    model_id (str): Bedrock model ID. Default ``"anthropic.claude-3-haiku-20240307-v1:0"``.
                    region (str): AWS region. Default ``"us-east-1"``.
                    auth_mode (str): ``"environment"`` (default) | ``"access_keys"`` | ``"api_key"``.
                    aws_access_key_id (str): Required when auth_mode is ``"access_keys"``.
                    aws_secret_access_key (str): Required when auth_mode is ``"access_keys"``.
                    bedrock_api_key (str): Required when auth_mode is ``"api_key"``.

                ``"openai"``
                    model_id (str): OpenAI model. Default ``"gpt-4o"``.
                    api_key (str): Falls back to ``OPENAI_API_KEY`` env var.

                ``"anthropic"``
                    model_id (str): Anthropic model. Default ``"claude-sonnet-4-20250514"``.
                    api_key (str): Falls back to ``ANTHROPIC_API_KEY`` env var.

                ``"gemini"``
                    model_id (str): Gemini model. Default ``"gemini-2.5-flash"``.
                    api_key (str): Falls back to ``GOOGLE_API_KEY`` env var.

                ``"ollama"``
                    model_id (str): Ollama model name. Default ``"llama3"``.
                    host (str): Ollama server URL. Default ``"http://localhost:11434"``.

            name: Optional registry name. Defaults to *provider*.
            **kwargs: Passed directly to the model constructor.
        """
        model_name = name or provider
        self._models_cfg[model_name] = {"type": "llm", "class": provider, **kwargs}
        self._pipeline_cfg["llm"] = model_name
        return self

    def tts(self, provider: str, *, name: str | None = None, **kwargs) -> "RoohBuilder":
        """Add a text-to-speech model.

        Args:
            provider: Model class — one of:

                ``"cartesia"``
                    voice_id (str): Cartesia voice ID. Default is a built-in English voice.
                    model_id (str): Cartesia model. Default ``"sonic-2"``.
                    api_key (str): Falls back to ``CARTESIA_API_KEY`` env var.

                ``"deepgram-tts"``
                    model (str): Deepgram TTS model. Default ``"aura-2-thalia-en"``.
                    api_key (str): Falls back to ``DEEPGRAM_API_KEY`` env var.

                ``"piper"``
                    voice (str): Piper voice name. Default ``"en_US-lessac-medium"``.

                ``"bark"``
                    voice (str): Bark voice preset. Default ``"v2/en_speaker_6"``.

                ``"speecht5"``
                    No additional kwargs required.

            name: Optional registry name. Defaults to *provider*.
            **kwargs: Passed directly to the model constructor.
        """
        model_name = name or provider
        self._models_cfg[model_name] = {"type": "tts", "class": provider, **kwargs}
        self._pipeline_cfg["tts"] = model_name
        return self

    def vad(self, provider: str = "silero", *, name: str | None = None, **kwargs) -> "RoohBuilder":
        """Add a voice activity detection model.

        Args:
            provider: Model class (default ``"silero"``).
            name: Optional registry name. Defaults to *provider*.
            **kwargs: Passed to the model constructor.
        """
        model_name = name or provider
        self._models_cfg[model_name] = {"type": "vad", "class": provider, **kwargs}
        self._pipeline_cfg["vad"] = model_name
        return self

    # ── Custom model classes ──────────────────────────────────────────────

    def custom_stt(self, name: str, cls: type, **kwargs) -> "RoohBuilder":
        """Register a custom STT model class and set it as the active STT.

        The class must subclass :class:`roohai.base.STTModel` and implement
        ``load()``, ``transcribe()``, ``unload()``, and ``is_loaded``.

        Args:
            name: Registry name for this model (e.g. ``"my-stt"``).
            cls: The model class (not an instance).
            **kwargs: Default constructor arguments (e.g. ``api_key``, ``model``).

        Example::

            from my_module import ElevenLabsSTT

            pipeline = (
                Rooh.builder()
                .custom_stt("elevenlabs-stt", ElevenLabsSTT, api_key="sk-xxx")
                .llm("bedrock")
                .tts("piper")
                .build()
            )
        """
        self._custom_models.append((name, "stt", cls, kwargs))
        self._pipeline_cfg["stt"] = name
        return self

    def custom_tts(self, name: str, cls: type, **kwargs) -> "RoohBuilder":
        """Register a custom TTS model class and set it as the active TTS.

        The class must subclass :class:`roohai.base.TTSModel` and implement
        ``load()``, ``synthesize()``, ``unload()``, and ``is_loaded``.

        Args:
            name: Registry name for this model (e.g. ``"my-tts"``).
            cls: The model class (not an instance).
            **kwargs: Default constructor arguments (e.g. ``api_key``, ``voice``).

        Example::

            from my_module import ElevenLabsTTS

            pipeline = (
                Rooh.builder()
                .stt("deepgram")
                .llm("bedrock")
                .custom_tts("elevenlabs", ElevenLabsTTS, api_key="sk-xxx", voice_id="abc")
                .build()
            )
        """
        self._custom_models.append((name, "tts", cls, kwargs))
        self._pipeline_cfg["tts"] = name
        return self

    def custom_llm(self, name: str, cls: type, **kwargs) -> "RoohBuilder":
        """Register a custom LLM model class and set it as the active LLM.

        The class must subclass :class:`roohai.base.LLMModel` and implement
        ``load()``, ``chat()``, ``unload()``, and ``is_loaded``.
        Optionally override ``chat_stream()`` for token-level streaming.

        Args:
            name: Registry name for this model (e.g. ``"my-llm"``).
            cls: The model class (not an instance).
            **kwargs: Default constructor arguments (e.g. ``api_key``, ``model``).

        Example::

            from my_module import MyLLM

            pipeline = (
                Rooh.builder()
                .stt("deepgram")
                .custom_llm("my-llm", MyLLM, api_key="sk-xxx", model="gpt-5")
                .tts("cartesia")
                .build()
            )
        """
        self._custom_models.append((name, "llm", cls, kwargs))
        self._pipeline_cfg["llm"] = name
        return self

    def custom_vad(self, name: str, cls: type, **kwargs) -> "RoohBuilder":
        """Register a custom VAD model class and set it as the active VAD.

        The class must subclass :class:`roohai.base.VADModel` and implement
        ``load()``, ``is_speech()``, ``unload()``, and ``is_loaded``.

        Args:
            name: Registry name for this model (e.g. ``"my-vad"``).
            cls: The model class (not an instance).
            **kwargs: Default constructor arguments.

        Example::

            from my_module import WebRTCVAD

            pipeline = (
                Rooh.builder()
                .stt("deepgram")
                .llm("bedrock")
                .tts("cartesia")
                .custom_vad("webrtc-vad", WebRTCVAD, aggressiveness=3)
                .build()
            )
        """
        self._custom_models.append((name, "vad", cls, kwargs))
        self._pipeline_cfg["vad"] = name
        return self

    # ── Pipeline settings ─────────────────────────────────────────────────

    def vad_threshold(self, value: float) -> "RoohBuilder":
        """Set VAD sensitivity (0.0–1.0). Default 0.7."""
        self._pipeline_cfg["vad_threshold"] = value
        return self

    def barge_in(self, enabled: bool = True) -> "RoohBuilder":
        """Enable or disable barge-in."""
        self._pipeline_cfg["barge_in_enabled"] = enabled
        return self

    def silence_duration(self, ms: int) -> "RoohBuilder":
        """Set silence duration in milliseconds before ending an utterance."""
        self._pipeline_cfg["silence_duration_ms"] = ms
        return self

    def system_prompt(self, prompt: str) -> "RoohBuilder":
        """Set the system prompt for the LLM."""
        self._system_prompt = prompt
        return self

    # ── Hooks ─────────────────────────────────────────────────────────────

    def llm_hook(self, hook, stream_hook=None) -> "RoohBuilder":
        """Set a custom LLM hook (replaces the LLM model call)."""
        self._llm_hook = hook
        self._llm_stream_hook = stream_hook
        return self

    def barge_in_hook(self, hook) -> "RoohBuilder":
        """Set a barge-in event hook."""
        self._barge_in_hook = hook
        return self

    def session_disconnect_hook(self, hook) -> "RoohBuilder":
        """Set a session disconnect hook."""
        self._session_disconnect_hook = hook
        return self

    # ── Build ─────────────────────────────────────────────────────────────

    def build(self) -> "Rooh":
        """Construct the :class:`Rooh` pipeline. Call ``.load()`` to load models."""
        # Register custom model classes before building so _register_from_config
        # can find them in the registry.
        if self._custom_models:
            register_all_models()  # ensure defaults are loaded first
            for name, model_type, cls, kwargs in self._custom_models:
                register_fn = {
                    "stt": registry.register_stt,
                    "tts": registry.register_tts,
                    "llm": registry.register_llm,
                    "vad": registry.register_vad,
                }[model_type]
                register_fn(name, cls, **kwargs)

        config: dict[str, Any] = {
            "pipeline": self._pipeline_cfg,
            "models": self._models_cfg,
        }
        if self._system_prompt is not None:
            config["system_prompt"] = self._system_prompt

        pipeline = Rooh(config_dict=config)

        if self._llm_hook is not None:
            pipeline.set_llm_hook(self._llm_hook, stream_hook=self._llm_stream_hook)
        if self._barge_in_hook is not None:
            pipeline.set_barge_in_hook(self._barge_in_hook)
        if self._session_disconnect_hook is not None:
            pipeline.set_session_disconnect_hook(self._session_disconnect_hook)

        return pipeline
