"""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",
        "whisper-tiny": "whisper",
        "whisper-base": "whisper",
        "whisper-small": "whisper",
        "whisper-medium": "whisper",
        "whisper-large": "whisper",
        "wav2vec2": "wav2vec2",
        "wav2vec2-large": "wav2vec2",
        "deepgram": "deepgram",
        "deepgram-nova-2": "deepgram",
        "deepgram-nova-3": "deepgram",
        "nvidia-canary": "nvidia-canary",
        "nvidia-parakeet": "nvidia-parakeet",
    },
    "tts": {
        "speecht5": "speecht5",
        "bark": "bark",
        "bark-small": "bark",
        "bark-large": "bark",
        "piper": "piper",
        "cartesia": "cartesia",
    },
    "llm": {
        "bedrock": "bedrock",
        "bedrock-claude": "bedrock",
        "bedrock-claude-haiku": "bedrock",
        "bedrock-claude-sonnet": "bedrock",
        "bedrock-claude-opus": "bedrock",
        "local": "local",
        "local-tinyllama": "local",
    },
    "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
    from roohai.llm import BedrockLLM, LocalLLM
    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}
    _llm_classes = {"bedrock": BedrockLLM, "local": LocalLLM}
    _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.",
    })

    # --- LLM variants ---
    BedrockLLM = llm_classes["bedrock"]
    LocalLLM = llm_classes["local"]

    registry.register_llm("bedrock-claude-haiku", BedrockLLM, meta={
        "display_name": "Claude Haiku (Bedrock)",
        "description": "Fast, cost-effective Claude model via AWS Bedrock. Great for voice.",
    }, model_id="anthropic.claude-3-haiku-20240307-v1:0")

    registry.register_llm("bedrock-claude-sonnet", BedrockLLM, meta={
        "display_name": "Claude Sonnet (Bedrock)",
        "description": "Balanced Claude model via AWS Bedrock. Better reasoning, moderate latency.",
    }, model_id="anthropic.claude-sonnet-4-20250514-v1:0")

    registry.register_llm("bedrock-claude-opus", BedrockLLM, meta={
        "display_name": "Claude Opus (Bedrock)",
        "description": "Most capable Claude model via AWS Bedrock. Best quality, higher latency.",
    }, model_id="anthropic.claude-opus-4-20250514-v1:0")

    registry.register_llm("local-tinyllama", LocalLLM, meta={
        "display_name": "TinyLlama 1.1B",
        "description": "Tiny local LLM. Fast but limited quality. Good for testing.",
        "estimated_download_mb": 2200,
    }, model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0")

    # --- 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") -> None:
        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._llm_hook = None
        self._llm_stream_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

    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."""
        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)
        if self.llm is None:
            raise RuntimeError("No LLM model loaded")
        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.

        Uses the stream hook if available, falls back to batch hook, then to the 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
        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

    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._llm_hook = None
        instance._llm_stream_hook = None
        return instance
