"""Unified inference engine manager.

Supports switching between Ollama and vLLM backends.
"""

from __future__ import annotations

import json
import logging
import time
from enum import Enum
from typing import TYPE_CHECKING

import httpx
import toml  # type: ignore[import-untyped]
from pydantic import BaseModel, Field

from llmhosts.config import _config_path, load_config
from llmhosts.constants import VLLM_DEFAULT_HOST
from llmhosts.discovery.ollama import OllamaDiscovery
from llmhosts.discovery.vllm import VLLMDiscovery

if TYPE_CHECKING:
    from llmhosts.config import LLMHostsConfig

logger = logging.getLogger(__name__)


class EngineType(str, Enum):
    """Supported inference engine types."""

    OLLAMA = "ollama"
    VLLM = "vllm"
    LLAMA_CPP = "llama.cpp"
    LMSTUDIO = "lmstudio"
    LOCALAI = "localai"
    TGI = "tgi"
    GENERIC = "generic"  # Any OpenAI-compatible server at a user-specified URL
    AUTO = "auto"  # Discover all available, prefer best


class InferenceEngine(BaseModel):
    """Represents a discovered inference engine."""

    engine_type: EngineType
    host: str
    available: bool = False
    models: list[str] = Field(default_factory=list)
    throughput_tps: float | None = None  # tokens/sec if measured


class EngineManager:
    """Manages Ollama and vLLM backends.

    - Discovers both engines
    - Tracks which is active
    - Provides unified model listing
    - Measures throughput for comparison
    """

    def __init__(self, config: LLMHostsConfig | None = None) -> None:
        self._config = config or load_config()
        self._ollama_host = self._config.ollama.host
        self._vllm_host = self._config.engine.vllm_host
        self._engine_type = self._get_engine_type_from_config()

    def _get_engine_type_from_config(self) -> EngineType:
        """Get engine type from config."""
        raw = self._config.engine.type
        try:
            return EngineType(str(raw).lower())
        except ValueError:
            return EngineType.OLLAMA

    async def discover_all(self) -> list[InferenceEngine]:
        """Discover all available engines."""
        results: list[InferenceEngine] = []

        # Discover Ollama
        async with OllamaDiscovery(host=self._ollama_host) as ollama:
            available = await ollama.is_available()
            models: list[str] = []
            if available:
                ollama_models = await ollama.list_models()
                models = [m.name for m in ollama_models]
            results.append(
                InferenceEngine(
                    engine_type=EngineType.OLLAMA,
                    host=self._ollama_host,
                    available=available,
                    models=models,
                )
            )

        # Discover vLLM
        async with VLLMDiscovery(host=self._vllm_host) as vllm:
            available = await vllm.is_available()
            models = []
            throughput: float | None = None
            if available:
                vllm_models = await vllm.list_models()
                models = [m.id for m in vllm_models]
                metrics = await vllm.get_metrics()
                if metrics.throughput_tps > 0:
                    throughput = metrics.throughput_tps
            results.append(
                InferenceEngine(
                    engine_type=EngineType.VLLM,
                    host=self._vllm_host,
                    available=available,
                    models=models,
                    throughput_tps=throughput,
                )
            )

        return results

    async def get_active_engine(self) -> InferenceEngine | None:
        """Get the currently active/preferred engine."""
        engines = await self.discover_all()
        engine_type = self._get_engine_type_from_config()

        if engine_type == EngineType.AUTO:
            # Prefer vLLM if available
            vllm_eng = next((e for e in engines if e.engine_type == EngineType.VLLM and e.available), None)
            if vllm_eng:
                return vllm_eng
            return next((e for e in engines if e.engine_type == EngineType.OLLAMA and e.available), None)

        return next((e for e in engines if e.engine_type == engine_type), None)

    async def switch_engine(self, engine_type: EngineType) -> bool:
        """Switch active engine. Persists to config file."""
        path = _config_path()
        path.parent.mkdir(parents=True, exist_ok=True)

        data: dict = {}
        if path.is_file():
            try:
                data = toml.load(path)
            except Exception as exc:
                logger.warning("Could not load config for engine switch: %s", exc)

        data.setdefault("engine", {})
        data["engine"]["type"] = engine_type.value
        data["engine"]["vllm_host"] = data["engine"].get("vllm_host", VLLM_DEFAULT_HOST)

        try:
            path.write_text(toml.dumps(data), encoding="utf-8")
            self._engine_type = engine_type
            return True
        except OSError as exc:
            logger.error("Failed to save engine config: %s", exc)
            return False

    async def compare_throughput(self, model: str) -> dict[str, float]:
        """Compare throughput between Ollama and vLLM for a model.

        Returns: {"ollama": tok/s, "vllm": tok/s}
        """
        result: dict[str, float] = {"ollama": 0.0, "vllm": 0.0}
        messages = [{"role": "user", "content": "Say 'hi' in one word."}]

        # Benchmark Ollama
        async with OllamaDiscovery(host=self._ollama_host) as ollama:
            if await ollama.is_available():
                try:
                    client = httpx.AsyncClient(
                        base_url=self._ollama_host,
                        timeout=httpx.Timeout(30.0),
                    )
                    payload = {"model": model, "messages": messages, "stream": True}
                    start = time.perf_counter()
                    token_count = 0
                    async with client.stream("POST", "/api/chat", json=payload) as resp:
                        if resp.status_code == 200:
                            async for line in resp.aiter_lines():
                                if line.strip():
                                    try:
                                        obj = json.loads(line)
                                        if obj.get("message", {}).get("content"):
                                            token_count += 1
                                    except json.JSONDecodeError:
                                        pass
                    elapsed = max(time.perf_counter() - start, 0.001)
                    result["ollama"] = token_count / elapsed
                    await client.aclose()
                except Exception as exc:
                    logger.debug("Ollama throughput benchmark failed: %s", exc)

        # Benchmark vLLM
        async with VLLMDiscovery(host=self._vllm_host) as vllm:
            if await vllm.is_available():
                tokens, elapsed = await vllm.chat_for_benchmark(model, messages)
                result["vllm"] = tokens / elapsed if elapsed > 0 else 0.0

        return result
