"""Agent preparation engine — checks deps, installs packages, downloads models."""

from __future__ import annotations

import asyncio
import importlib.util
import json
import logging
import re
from typing import Any, Callable, Awaitable

from fastapi import WebSocket

logger = logging.getLogger(__name__)

SendProgress = Callable[[dict], Awaitable[None]]

# Known pip-name -> import-name mappings
_IMPORT_MAP = {
    "deepgram-sdk": "deepgram",
    "silero-vad": "silero_vad",
    "piper-tts": "piper",
    "huggingface-hub": "huggingface_hub",
    "nemo_toolkit": "nemo",
}


async def _ws_send(ws: WebSocket, msg: dict) -> None:
    await ws.send_text(json.dumps(msg))


def _pip_to_module(dep: str) -> str:
    """Convert pip dependency spec to importable module name."""
    name = re.split(r"[>=<!\[;]", dep)[0].strip()
    return _IMPORT_MAP.get(name, name.replace("-", "_"))


def check_dependencies(pip_deps: list[str]) -> dict[str, bool]:
    """Check which pip dependencies are already installed."""
    results = {}
    for dep in pip_deps:
        module_name = _pip_to_module(dep)
        results[dep] = importlib.util.find_spec(module_name) is not None
    return results


async def install_dependencies(pip_deps: list[str], send_progress: SendProgress) -> None:
    """Install missing pip dependencies via subprocess."""
    missing = [dep for dep, found in check_dependencies(pip_deps).items() if not found]
    if not missing:
        await send_progress({
            "status": "progress", "progress": 100,
            "message": "All dependencies already installed",
        })
        return

    for i, dep in enumerate(missing):
        pct = int((i / len(missing)) * 100)
        await send_progress({
            "status": "progress", "progress": pct,
            "message": f"Installing {dep}...",
        })
        proc = await asyncio.create_subprocess_exec(
            "pip", "install", dep,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        _, stderr = await proc.communicate()
        if proc.returncode != 0:
            raise RuntimeError(f"Failed to install {dep}: {stderr.decode()}")

    await send_progress({
        "status": "progress", "progress": 100,
        "message": "All dependencies installed",
    })


async def download_models(hf_models: list[str], send_progress: SendProgress) -> None:
    """Download HuggingFace models via snapshot_download in a thread."""
    if not hf_models:
        return

    for i, model_id in enumerate(hf_models):
        pct = int((i / len(hf_models)) * 100)
        await send_progress({
            "status": "progress", "progress": pct,
            "message": f"Downloading {model_id}...",
        })
        try:
            from huggingface_hub import snapshot_download
            await asyncio.to_thread(snapshot_download, model_id)
        except Exception as exc:
            raise RuntimeError(f"Failed to download {model_id}: {exc}") from exc

    await send_progress({
        "status": "progress", "progress": 100,
        "message": "All models downloaded",
    })


async def validate_credentials(
    model_class: str, config: dict[str, Any], send_progress: SendProgress
) -> None:
    """Validate cloud credentials by attempting to instantiate a client."""
    await send_progress({
        "status": "progress", "progress": 50,
        "message": f"Validating credentials for {model_class}...",
    })

    if model_class == "bedrock":
        try:
            from roohai.llm.bedrock import BedrockLLM
            client = BedrockLLM(**{k: v for k, v in config.items() if k not in ("type", "class")})
            await asyncio.to_thread(client._ensure_client)
        except Exception as exc:
            raise RuntimeError(f"Bedrock credential validation failed: {exc}") from exc
    elif model_class == "deepgram":
        api_key = config.get("api_key")
        if not api_key:
            raise RuntimeError("Deepgram requires an API key")
    elif model_class == "cartesia":
        api_key = config.get("api_key")
        if not api_key:
            raise RuntimeError("Cartesia requires an API key")

    await send_progress({
        "status": "progress", "progress": 100,
        "message": f"Credentials valid for {model_class}",
    })


def _get_meta_for_class(class_name: str) -> dict:
    """Look up META from the model class via registry or direct import."""
    from roohai.pipeline import register_all_models
    from roohai.registry import registry

    register_all_models()
    for store in (registry._stt, registry._tts, registry._llm, registry._vad):
        if class_name in store:
            cls, _kwargs, meta_overrides = store[class_name]
            base_meta = getattr(cls, "META", {})
            return {**base_meta, **meta_overrides}
    return {}


async def prepare_agent(agent_config: dict, websocket: WebSocket) -> None:
    """Orchestrate full agent preparation with WebSocket progress updates."""
    send = lambda msg: _ws_send(websocket, msg)
    models = agent_config.get("models", {})

    try:
        # Collect deps and models from class META
        pip_deps: list[str] = []
        hf_models: list[str] = []
        cloud_models: list[tuple[str, dict]] = []

        for name, model_cfg in models.items():
            cls_name = model_cfg.get("class", "")
            meta = _get_meta_for_class(cls_name)

            pip_deps.extend(meta.get("pip_dependencies", []))
            hf_models.extend(meta.get("hf_models", []))

            if meta.get("kind") == "cloud":
                cloud_models.append((cls_name, model_cfg))

        # Install missing dependencies
        if pip_deps:
            await install_dependencies(pip_deps, send)

        # Download HF models
        if hf_models:
            await download_models(hf_models, send)

        # Validate cloud credentials
        for cls, cfg in cloud_models:
            await validate_credentials(cls, cfg, send)

        await send({"status": "complete", "message": "Voice agent ready!"})

    except RuntimeError as exc:
        await send({"status": "error", "message": str(exc)})
    except Exception as exc:
        logger.exception("Agent preparation failed")
        await send({"status": "error", "message": f"Unexpected error: {exc}"})
