"""LLMHosts Core inference engine wrapper."""

from __future__ import annotations

import json
import logging
from typing import Any, cast

logger = logging.getLogger(__name__)

LLMHOSTS_CORE_AVAILABLE: bool = False
_LLMHostsEngineNative: type | None = None

try:
    from _llmhosts import core as _llmhosts_core_mod

    _LLMHostsEngineNative = _llmhosts_core_mod.LLMHostsEngine
    LLMHOSTS_CORE_AVAILABLE = True
    logger.debug("LLMHosts Core engine loaded (_llmhosts.core %s)", _llmhosts_core_mod.LLMHOSTS_CORE_VERSION)
except ImportError:
    # Fallback: try standalone module (development mode with individual maturin develop)
    try:
        import _llmhosts_core

        _LLMHostsEngineNative = _llmhosts_core.LLMHostsEngine
        LLMHOSTS_CORE_AVAILABLE = True
        logger.debug("LLMHosts Core engine loaded (standalone _llmhosts_core %s)", _llmhosts_core.LLMHOSTS_CORE_VERSION)
    except ImportError:
        logger.debug("LLMHosts Core engine not available")


class LLMHostsInferenceEngine:
    def __init__(self) -> None:
        self._native = _LLMHostsEngineNative() if _LLMHostsEngineNative else None
        self._tokenizer: Any = None

    @property
    def is_loaded(self) -> bool:
        return self._native is not None and self._native.is_loaded

    def load_model(self, path: str) -> dict[str, Any]:
        if self._native is None:
            raise RuntimeError("LLMHosts Core engine not available (build with: cd llmhosts_core && maturin develop)")
        info: dict[str, Any] = json.loads(self._native.load_model(path))
        self._tokenizer = self._load_tokenizer(path)
        return info

    def _load_tokenizer(self, model_path: str) -> Any:
        import os

        try:
            from tokenizers import Tokenizer
        except ImportError:
            logger.debug("tokenizers package not available (install with: pip install llmhosts[smart])")
            return None

        model_dir = os.path.dirname(model_path)
        model_name = os.path.basename(model_path).lower()

        # Search multiple locations for tokenizer.json
        search_paths = [
            os.path.join(model_dir, "tokenizer.json"),
            os.path.join(model_dir, "..", "tokenizer.json"),
        ]

        for path in search_paths:
            if os.path.exists(path):
                try:
                    logger.info("Loading tokenizer from %s", path)
                    return Tokenizer.from_file(path)
                except Exception as exc:
                    logger.warning("Failed to load tokenizer from %s: %s", path, exc)

        # Auto-download from HuggingFace if model name is recognizable
        hf_repo = self._guess_hf_repo(model_name)
        if hf_repo:
            try:
                from huggingface_hub import hf_hub_download

                logger.info("Downloading tokenizer.json from %s", hf_repo)
                path = hf_hub_download(  # nosec B615 — revision pinned to main
                    repo_id=hf_repo, filename="tokenizer.json", local_dir=model_dir, revision="main"
                )
                return Tokenizer.from_file(path)
            except Exception as exc:
                logger.warning("Failed to download tokenizer from HuggingFace (%s): %s", hf_repo, exc)

        logger.warning("No tokenizer found for %s", model_path)
        return None

    @staticmethod
    def _guess_hf_repo(model_name: str) -> str | None:
        """Guess the HuggingFace repo ID from a GGUF filename."""
        patterns: dict[str, str] = {
            "qwen2.5-0.5b-instruct": "Qwen/Qwen2.5-0.5B-Instruct",
            "qwen2.5-1.5b-instruct": "Qwen/Qwen2.5-1.5B-Instruct",
            "qwen2.5-3b-instruct": "Qwen/Qwen2.5-3B-Instruct",
            "qwen2.5-7b-instruct": "Qwen/Qwen2.5-7B-Instruct",
            "qwen2.5-14b-instruct": "Qwen/Qwen2.5-14B-Instruct",
            "llama-3": "meta-llama/Llama-3.1-8B-Instruct",
            "mistral": "mistralai/Mistral-7B-Instruct-v0.3",
        }
        for pattern, repo in patterns.items():
            if pattern in model_name:
                return repo
        return None

    def generate(
        self, prompt: str, max_tokens: int = 256, temperature: float = 0.0, top_p: float = 0.9
    ) -> dict[str, Any]:
        if self._native is None:
            raise RuntimeError("LLMHosts Core engine not available")
        if self._tokenizer is None:
            raise RuntimeError(
                "No tokenizer loaded. Place a tokenizer.json next to the model file, "
                "or install the tokenizers package: pip install llmhosts[smart]"
            )
        encoded = self._tokenizer.encode(prompt)
        token_ids = encoded.ids
        generated = json.loads(self._native.generate_tokens(token_ids, max_tokens, temperature, top_p))
        text = self._tokenizer.decode(generated)
        return {"text": text, "tokens": generated, "token_count": len(generated)}

    def generate_tokens(
        self, token_ids: list[int], max_tokens: int = 256, temperature: float = 0.0, top_p: float = 0.9
    ) -> list[int]:
        if self._native is None:
            raise RuntimeError("LLMHosts Core engine not available")
        result: list[int] = json.loads(self._native.generate_tokens(token_ids, max_tokens, temperature, top_p))
        return result

    def unload(self) -> None:
        if self._native:
            self._native.unload()
        self._tokenizer = None

    @property
    def backend_name(self) -> str:
        """Return the active compute backend name (e.g. 'cpu', 'cuda:0', 'wgpu:Metal')."""
        if self._native is not None and self._native.is_loaded:
            return cast("str", self._native.backend_name())
        return "none"

    @staticmethod
    def available_backends() -> list[str]:
        """Return list of compute backends compiled into this build."""
        if _LLMHostsEngineNative is not None:
            return cast("list[str]", _LLMHostsEngineNative.available_backends())  # type: ignore[attr-defined]
        return ["cpu"]

    def status(self) -> dict[str, Any]:
        return {
            "available": LLMHOSTS_CORE_AVAILABLE,
            "loaded": self.is_loaded,
            "backend": self.backend_name,
            "available_backends": self.available_backends(),
        }
