"""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."""
    from roohai.stt import WhisperSTT, Wav2Vec2STT, DeepgramSTT, NvidiaCanarySTT, NvidiaParakeetSTT
    from roohai.tts import SpeechT5TTS, BarkTTS, PiperTTS, CartesiaTTS, DeepgramTTS
    from roohai.llm import BedrockLLM, LocalLLM, OpenAILLM, AnthropicLLM, GeminiLLM, OllamaLLM
    from roohai.vad import SileroVAD

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

    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"],
    }, 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"],
    }, 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"],
    }, 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"],
    }, 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"],
    }, 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.",
    })

    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"],
    }, 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.",
    }, 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.",
    }, model="nova-2")

    NvidiaCanarySTT = stt_classes["nvidia-canary"]
    NvidiaParakeetSTT = stt_classes["nvidia-parakeet"]

    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,
    }, model_id="nvidia/canary-qwen-2.5b")

    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,
    }, model_id="nvidia/parakeet-tdt-0.6b-v2")

    # --- TTS variants ---
    SpeechT5TTS = tts_classes["speecht5"]
    BarkTTS = tts_classes["bark"]
    PiperTTS = tts_classes["piper"]
    CartesiaTTS = tts_classes["cartesia"]

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

    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"],
    }, 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"],
    }, model_id="suno/bark")

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

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

    DeepgramTTS = tts_classes["deepgram-tts"]

    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.",
    })

    # --- 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_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

    def load(self) -> None:
        """Load all configured models."""
        if self._stt_name:
            self.stt = registry.create_stt(self._stt_name)
            self.stt.load()
            logger.info("Loaded STT model: %s", self._stt_name)

        if self._tts_name:
            self.tts = registry.create_tts(self._tts_name)
            self.tts.load()
            logger.info("Loaded TTS model: %s", self._tts_name)

        if self._llm_name:
            self.llm = registry.create_llm(self._llm_name)
            self.llm.load()
            logger.info("Loaded LLM model: %s", self._llm_name)

        if self._vad_name:
            self.vad = registry.create_vad(self._vad_name)
            self.vad.load()
            logger.info("Loaded VAD model: %s", self._vad_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."""
        if model_type == "stt":
            if self.stt is not None:
                self.stt.unload()
            self.stt = registry.create_stt(model_name)
            self.stt.load()
            self._stt_name = model_name
        elif model_type == "tts":
            if self.tts is not None:
                self.tts.unload()
            self.tts = registry.create_tts(model_name)
            self.tts.load()
            self._tts_name = model_name
        elif model_type == "llm":
            if self.llm is not None:
                self.llm.unload()
            self.llm = registry.create_llm(model_name)
            self.llm.load()
            self._llm_name = model_name
        elif model_type == "vad":
            if self.vad is not None:
                self.vad.unload()
            self.vad = registry.create_vad(model_name)
            self.vad.load()
            self._vad_name = model_name
        else:
            raise ValueError(f"Unknown model type: {model_type}")
        logger.info("Swapped %s model to: %s", model_type, 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
