"""Model suggestion engine -- recommends optimal models based on hardware profile.

Given a :class:`HardwareProfile`, this module recommends which Ollama models
to pull based on available VRAM/RAM, quantisation preferences, and use case.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel

if TYPE_CHECKING:
    from llmhosts.discovery.models import HardwareProfile

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Suggestion model
# ---------------------------------------------------------------------------


class ModelSuggestion(BaseModel):
    """A single model recommendation with rationale."""

    model_name: str  # e.g. "llama3.2:8b"
    display_name: str  # e.g. "Llama 3.2 8B"
    size_gb: float  # download size in GB
    vram_required_gb: float  # VRAM needed to run the model
    description: str
    use_cases: list[str]  # ["general", "coding", "creative"]
    quality_rating: int  # 1-5 stars
    speed_rating: int  # 1-5 stars
    pull_command: str  # "ollama pull llama3.2:8b"
    already_installed: bool = False
    reason: str  # "Fits in your 12GB VRAM, great for coding"


# ---------------------------------------------------------------------------
# Model catalog -- static knowledge of popular Ollama models
# ---------------------------------------------------------------------------

_CatalogEntry = dict[str, Any]

MODEL_CATALOG: dict[str, _CatalogEntry] = {
    # ---- Tiny (<4 GB VRAM) ----
    "phi3:mini": {
        "display_name": "Phi-3 Mini",
        "size_gb": 2.3,
        "vram_required_gb": 3.0,
        "description": "Microsoft's compact 3.8B model. Surprisingly capable for its size, great at reasoning and coding tasks.",
        "use_cases": ["general", "coding"],
        "quality_rating": 3,
        "speed_rating": 5,
        "tier": "tiny",
    },
    "qwen2.5:1.5b": {
        "display_name": "Qwen 2.5 1.5B",
        "size_gb": 1.0,
        "vram_required_gb": 2.0,
        "description": "Alibaba's ultra-compact model. Fast inference, decent for simple tasks and chat.",
        "use_cases": ["general", "chat"],
        "quality_rating": 2,
        "speed_rating": 5,
        "tier": "tiny",
    },
    "gemma2:2b": {
        "display_name": "Gemma 2 2B",
        "size_gb": 1.6,
        "vram_required_gb": 2.5,
        "description": "Google's lightweight model. Good balance of quality and speed for edge deployment.",
        "use_cases": ["general", "chat"],
        "quality_rating": 2,
        "speed_rating": 5,
        "tier": "tiny",
    },
    # ---- Small (4-8 GB VRAM) ----
    "llama3.2:3b": {
        "display_name": "Llama 3.2 3B",
        "size_gb": 2.0,
        "vram_required_gb": 4.0,
        "description": "Meta's compact Llama 3.2. Excellent quality for a 3B model, strong at general tasks.",
        "use_cases": ["general", "chat", "creative"],
        "quality_rating": 3,
        "speed_rating": 5,
        "tier": "small",
    },
    "mistral:7b-q4_K_M": {
        "display_name": "Mistral 7B (Q4)",
        "size_gb": 4.1,
        "vram_required_gb": 5.5,
        "description": "Mistral AI's 7B model, 4-bit quantised. Strong instruction following and reasoning.",
        "use_cases": ["general", "creative", "analysis"],
        "quality_rating": 4,
        "speed_rating": 4,
        "tier": "small",
    },
    "codellama:7b-q4_K_M": {
        "display_name": "Code Llama 7B (Q4)",
        "size_gb": 4.0,
        "vram_required_gb": 5.5,
        "description": "Meta's code-specialised Llama. Excellent for code completion, review, and generation.",
        "use_cases": ["coding"],
        "quality_rating": 4,
        "speed_rating": 4,
        "tier": "small",
    },
    # ---- Medium (8-12 GB VRAM) ----
    "llama3.2:8b": {
        "display_name": "Llama 3.2 8B",
        "size_gb": 4.7,
        "vram_required_gb": 8.0,
        "description": "Meta's workhorse 8B model. Excellent across general, coding, and creative tasks.",
        "use_cases": ["general", "coding", "creative", "analysis"],
        "quality_rating": 4,
        "speed_rating": 4,
        "tier": "medium",
    },
    "mistral:7b": {
        "display_name": "Mistral 7B",
        "size_gb": 4.1,
        "vram_required_gb": 8.0,
        "description": "Mistral AI's full-precision 7B. Top-tier reasoning and instruction following at this size.",
        "use_cases": ["general", "creative", "analysis"],
        "quality_rating": 4,
        "speed_rating": 4,
        "tier": "medium",
    },
    "codellama:7b": {
        "display_name": "Code Llama 7B",
        "size_gb": 3.8,
        "vram_required_gb": 8.0,
        "description": "Full-precision Code Llama 7B. Superior code quality vs quantised version.",
        "use_cases": ["coding"],
        "quality_rating": 4,
        "speed_rating": 4,
        "tier": "medium",
    },
    "qwen2.5:7b": {
        "display_name": "Qwen 2.5 7B",
        "size_gb": 4.4,
        "vram_required_gb": 8.0,
        "description": "Alibaba's Qwen 2.5 7B. Excellent multilingual support and strong coding ability.",
        "use_cases": ["general", "coding", "chat"],
        "quality_rating": 4,
        "speed_rating": 4,
        "tier": "medium",
    },
    "deepseek-coder:6.7b": {
        "display_name": "DeepSeek Coder 6.7B",
        "size_gb": 3.8,
        "vram_required_gb": 8.0,
        "description": "Purpose-built coding model from DeepSeek. Excels at code generation and debugging.",
        "use_cases": ["coding"],
        "quality_rating": 4,
        "speed_rating": 4,
        "tier": "medium",
    },
    # ---- Large (12-24 GB VRAM) ----
    "llama3.1:13b": {
        "display_name": "Llama 3.1 13B",
        "size_gb": 7.4,
        "vram_required_gb": 13.0,
        "description": "Meta's 13B parameter model. Significant quality jump over 8B for complex reasoning.",
        "use_cases": ["general", "coding", "creative", "analysis"],
        "quality_rating": 4,
        "speed_rating": 3,
        "tier": "large",
    },
    "codellama:13b": {
        "display_name": "Code Llama 13B",
        "size_gb": 7.4,
        "vram_required_gb": 13.0,
        "description": "Larger Code Llama. Substantially better at complex multi-file code generation.",
        "use_cases": ["coding"],
        "quality_rating": 5,
        "speed_rating": 3,
        "tier": "large",
    },
    "qwen2.5:14b": {
        "display_name": "Qwen 2.5 14B",
        "size_gb": 8.2,
        "vram_required_gb": 14.0,
        "description": "Alibaba's 14B model. Competitive with much larger models on benchmarks.",
        "use_cases": ["general", "coding", "analysis"],
        "quality_rating": 5,
        "speed_rating": 3,
        "tier": "large",
    },
    "mixtral:8x7b-q4_K_M": {
        "display_name": "Mixtral 8x7B (Q4)",
        "size_gb": 26.0,
        "vram_required_gb": 24.0,
        "description": "Mistral's MoE architecture. Only uses 2 experts per token, so faster than size suggests.",
        "use_cases": ["general", "creative", "analysis", "coding"],
        "quality_rating": 5,
        "speed_rating": 3,
        "tier": "large",
    },
    # ---- XL (24+ GB VRAM) ----
    "llama3.1:70b-q4_K_M": {
        "display_name": "Llama 3.1 70B (Q4)",
        "size_gb": 40.0,
        "vram_required_gb": 40.0,
        "description": "Meta's flagship 70B model, 4-bit quantised. Near-GPT-4 quality for local inference.",
        "use_cases": ["general", "coding", "creative", "analysis"],
        "quality_rating": 5,
        "speed_rating": 2,
        "tier": "xl",
    },
    "mixtral:8x7b": {
        "display_name": "Mixtral 8x7B",
        "size_gb": 26.0,
        "vram_required_gb": 48.0,
        "description": "Full-precision Mixtral MoE. Premium quality, needs serious VRAM.",
        "use_cases": ["general", "creative", "analysis", "coding"],
        "quality_rating": 5,
        "speed_rating": 2,
        "tier": "xl",
    },
    "codellama:34b": {
        "display_name": "Code Llama 34B",
        "size_gb": 19.0,
        "vram_required_gb": 28.0,
        "description": "The largest Code Llama. Best-in-class local code generation and understanding.",
        "use_cases": ["coding"],
        "quality_rating": 5,
        "speed_rating": 2,
        "tier": "xl",
    },
    "qwen2.5:32b": {
        "display_name": "Qwen 2.5 32B",
        "size_gb": 18.5,
        "vram_required_gb": 28.0,
        "description": "Alibaba's 32B model. Excellent at complex reasoning and multilingual tasks.",
        "use_cases": ["general", "coding", "creative", "analysis"],
        "quality_rating": 5,
        "speed_rating": 2,
        "tier": "xl",
    },
}

# VRAM tier boundaries in GB (inclusive lower bound)
_VRAM_TIERS: list[tuple[float, list[str]]] = [
    (0.0, ["tiny"]),
    (4.0, ["tiny", "small"]),
    (8.0, ["tiny", "small", "medium"]),
    (12.0, ["tiny", "small", "medium", "large"]),
    (24.0, ["tiny", "small", "medium", "large", "xl"]),
]

# Use-case priority weights (higher = more relevant)
_USE_CASE_PRIORITY: dict[str, list[str]] = {
    "coding": [
        "codellama:34b",
        "codellama:13b",
        "deepseek-coder:6.7b",
        "codellama:7b",
        "codellama:7b-q4_K_M",
        "qwen2.5:14b",
        "qwen2.5:7b",
        "llama3.2:8b",
    ],
    "general": [
        "llama3.1:70b-q4_K_M",
        "qwen2.5:32b",
        "llama3.1:13b",
        "qwen2.5:14b",
        "mixtral:8x7b",
        "llama3.2:8b",
        "mistral:7b",
        "llama3.2:3b",
    ],
    "creative": [
        "llama3.1:70b-q4_K_M",
        "mixtral:8x7b",
        "mixtral:8x7b-q4_K_M",
        "llama3.1:13b",
        "llama3.2:8b",
        "mistral:7b",
        "mistral:7b-q4_K_M",
    ],
    "chat": [
        "llama3.2:8b",
        "llama3.2:3b",
        "qwen2.5:7b",
        "mistral:7b",
        "qwen2.5:1.5b",
        "gemma2:2b",
        "phi3:mini",
    ],
    "analysis": [
        "llama3.1:70b-q4_K_M",
        "qwen2.5:32b",
        "qwen2.5:14b",
        "llama3.1:13b",
        "mixtral:8x7b-q4_K_M",
        "llama3.2:8b",
        "mistral:7b",
    ],
}


# ---------------------------------------------------------------------------
# Suggester
# ---------------------------------------------------------------------------


class ModelSuggester:
    """Suggests optimal models based on hardware profile.

    Recommends which models to pull via Ollama based on:
    - Available VRAM (or effective VRAM from RAM for CPU-only systems)
    - Use case preferences (coding, general, creative, etc.)
    - Already-installed models (filtered out unless explicitly requested)
    """

    def __init__(self) -> None:
        self._model_catalog: dict[str, _CatalogEntry] = MODEL_CATALOG

    def suggest(
        self,
        hardware: HardwareProfile,
        installed_models: list[str] | None = None,
        use_cases: list[str] | None = None,
    ) -> list[ModelSuggestion]:
        """Generate model suggestions based on hardware.

        Rules
        -----
        - ``<4 GB`` VRAM: phi3:mini, qwen2.5:1.5b (tiny models only)
        - ``4-8 GB`` VRAM: llama3.2:3b, mistral:7b-q4, codellama:7b-q4
        - ``8-12 GB`` VRAM: llama3.2:8b, mistral:7b, codellama:7b, qwen2.5:7b
        - ``12-24 GB`` VRAM: llama3.1:13b, codellama:13b, deepseek-coder:6.7b
        - ``24+ GB`` VRAM: llama3.1:70b-q4, mixtral:8x7b, codellama:34b
        - No GPU: use RAM-based recommendations (halve the effective VRAM)

        Parameters
        ----------
        hardware:
            Detected hardware profile.
        installed_models:
            List of model names already installed (e.g. from ``ollama list``).
            These are marked ``already_installed=True`` and sorted to the end.
        use_cases:
            Optional list of use cases to prioritise.
            Valid values: ``"coding"``, ``"general"``, ``"creative"``, ``"chat"``, ``"analysis"``.
        """
        installed_set = {m.lower().strip() for m in (installed_models or [])}
        effective_vram_gb = self._effective_vram(hardware)
        allowed_tiers = self._allowed_tiers(effective_vram_gb)

        logger.debug(
            "Suggesting models: effective_vram=%.1fGB, tiers=%s, use_cases=%s",
            effective_vram_gb,
            allowed_tiers,
            use_cases,
        )

        suggestions: list[ModelSuggestion] = []
        for model_name, entry in self._model_catalog.items():
            if entry["tier"] not in allowed_tiers:
                continue
            if entry["vram_required_gb"] > effective_vram_gb * 1.1:
                # Allow 10% headroom overshoot, but no more
                continue

            is_installed = self._is_installed(model_name, installed_set)
            reason = self._build_reason(model_name, entry, hardware, effective_vram_gb, is_installed)

            suggestions.append(
                ModelSuggestion(
                    model_name=model_name,
                    display_name=entry["display_name"],
                    size_gb=entry["size_gb"],
                    vram_required_gb=entry["vram_required_gb"],
                    description=entry["description"],
                    use_cases=entry["use_cases"],
                    quality_rating=entry["quality_rating"],
                    speed_rating=entry["speed_rating"],
                    pull_command=f"ollama pull {model_name}",
                    already_installed=is_installed,
                    reason=reason,
                )
            )

        # Sort: prioritise by use case, then quality, then not-installed first
        suggestions = self._sort_suggestions(suggestions, use_cases)
        return suggestions

    def get_pull_commands(self, suggestions: list[ModelSuggestion]) -> list[str]:
        """Generate ``ollama pull`` commands for suggestions that aren't already installed."""
        return [s.pull_command for s in suggestions if not s.already_installed]

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _effective_vram(hardware: HardwareProfile) -> float:
        """Calculate effective VRAM in GB.

        If a GPU is present, use total VRAM.
        Otherwise, use half the system RAM as a rough CPU-inference estimate
        (Ollama offloads to CPU/RAM, but throughput is roughly half).
        """
        if hardware.has_gpu and hardware.total_vram_mb > 0:
            return hardware.total_vram_mb / 1024.0
        # CPU-only: use half of system RAM as effective VRAM
        return hardware.ram_total_gb / 2.0

    @staticmethod
    def _allowed_tiers(effective_vram_gb: float) -> set[str]:
        """Determine which model tiers fit in the available VRAM."""
        tiers: set[str] = set()
        for threshold, tier_list in _VRAM_TIERS:
            if effective_vram_gb >= threshold:
                tiers.update(tier_list)
        return tiers if tiers else {"tiny"}

    @staticmethod
    def _is_installed(model_name: str, installed_set: set[str]) -> bool:
        """Check if a model is already installed (fuzzy match on name prefix)."""
        model_lower = model_name.lower()
        if model_lower in installed_set:
            return True
        # Also match without quantisation suffix (e.g. "mistral:7b" matches "mistral:7b-q4_K_M")
        base_name = model_lower.split("-q")[0].split(":")[0]
        for installed in installed_set:
            installed_base = installed.split("-q")[0].split(":")[0]
            if (
                base_name == installed_base
                and model_lower.split(":")[0] == installed.split(":")[0]
                and (model_lower in installed or installed in model_lower)
            ):
                return True
        return False

    @staticmethod
    def _build_reason(
        model_name: str,
        entry: _CatalogEntry,
        hardware: HardwareProfile,
        effective_vram_gb: float,
        is_installed: bool,
    ) -> str:
        """Generate a human-readable reason for this suggestion."""
        parts: list[str] = []

        vram_req = entry["vram_required_gb"]
        headroom = effective_vram_gb - vram_req
        if headroom >= 4:
            parts.append(f"Fits comfortably in your {effective_vram_gb:.0f}GB VRAM")
        elif headroom >= 0:
            parts.append(f"Fits in your {effective_vram_gb:.0f}GB VRAM")
        else:
            parts.append(f"Tight fit for your {effective_vram_gb:.0f}GB VRAM")

        if not hardware.has_gpu:
            parts.append("runs on CPU/RAM (slower but functional)")

        cases = entry["use_cases"]
        if len(cases) == 1:
            parts.append(f"specialised for {cases[0]}")
        elif len(cases) >= 3:
            parts.append("versatile all-rounder")
        else:
            parts.append(f"great for {' and '.join(cases)}")

        if is_installed:
            parts.append("already installed")

        return ", ".join(parts)

    @staticmethod
    def _sort_suggestions(
        suggestions: list[ModelSuggestion],
        use_cases: list[str] | None,
    ) -> list[ModelSuggestion]:
        """Sort suggestions by relevance.

        Priority:
        1. Use-case match score (higher = better)
        2. Quality rating (higher = better)
        3. Not-installed first
        4. Speed rating as tiebreaker
        """
        use_case_set = set(use_cases) if use_cases else set()

        # Build priority lookup from _USE_CASE_PRIORITY
        priority_scores: dict[str, float] = {}
        if use_case_set:
            for uc in use_case_set:
                uc_models = _USE_CASE_PRIORITY.get(uc, [])
                for rank, model_name in enumerate(uc_models):
                    score = len(uc_models) - rank  # higher rank = higher score
                    priority_scores[model_name] = priority_scores.get(model_name, 0) + score

        def sort_key(s: ModelSuggestion) -> tuple[int, float, int, int]:
            # installed last (0 = not installed, 1 = installed)
            installed_penalty = 1 if s.already_installed else 0

            # Use-case relevance (negative for descending)
            uc_score = 0.0
            if use_case_set:
                # Direct priority score
                uc_score = priority_scores.get(s.model_name, 0)
                # Also boost if model's use_cases overlap with requested
                overlap = len(use_case_set & set(s.use_cases))
                uc_score += overlap * 5

            return (
                installed_penalty,  # not-installed first
                -uc_score,  # higher use-case relevance first
                -s.quality_rating,  # higher quality first
                -s.speed_rating,  # faster first as tiebreaker
            )

        suggestions.sort(key=sort_key)
        return suggestions
