"""Model selection suggestion via LLM."""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

# Well-known model families with their characteristics
MODEL_CATALOG: list[dict] = [
    {
        "id": "meta-llama/Llama-3.1-8B",
        "family": "llama",
        "params": "8B",
        "min_vram_gb": 6,
        "strengths": ["general", "instruction-following", "multilingual"],
    },
    {
        "id": "meta-llama/Llama-3.1-70B",
        "family": "llama",
        "params": "70B",
        "min_vram_gb": 40,
        "strengths": ["general", "reasoning", "multilingual"],
    },
    {
        "id": "mistralai/Mistral-7B-v0.3",
        "family": "mistral",
        "params": "7B",
        "min_vram_gb": 6,
        "strengths": ["general", "code", "fast-inference"],
    },
    {
        "id": "mistralai/Mixtral-8x7B-v0.1",
        "family": "mistral",
        "params": "46.7B (MoE)",
        "min_vram_gb": 30,
        "strengths": ["general", "reasoning", "MoE"],
    },
    {
        "id": "Qwen/Qwen2.5-7B",
        "family": "qwen",
        "params": "7B",
        "min_vram_gb": 6,
        "strengths": ["general", "multilingual", "code", "math"],
    },
    {
        "id": "Qwen/Qwen2.5-72B",
        "family": "qwen",
        "params": "72B",
        "min_vram_gb": 40,
        "strengths": ["general", "reasoning", "multilingual"],
    },
    {
        "id": "google/gemma-2-9b",
        "family": "gemma",
        "params": "9B",
        "min_vram_gb": 8,
        "strengths": ["general", "lightweight", "efficient"],
    },
    {
        "id": "microsoft/phi-3-mini-4k-instruct",
        "family": "phi",
        "params": "3.8B",
        "min_vram_gb": 4,
        "strengths": ["reasoning", "code", "small-footprint"],
    },
    {
        "id": "codellama/CodeLlama-7b-hf",
        "family": "codellama",
        "params": "7B",
        "min_vram_gb": 6,
        "strengths": ["code", "code-generation", "code-completion"],
    },
    {
        "id": "codellama/CodeLlama-34b-hf",
        "family": "codellama",
        "params": "34B",
        "min_vram_gb": 20,
        "strengths": ["code", "code-generation"],
    },
]

SUGGEST_MODEL_PROMPT = """You are an expert ML engineer. Based on the hardware profile and
dataset characteristics below, suggest the best model for fine-tuning.

Hardware Profile:
{hardware_json}

Dataset Profile:
{dataset_json}

Available models that fit in the user's VRAM:
{candidate_models}

Task context (inferred from dataset format and columns):
- Format: {dataset_format}
- Columns: {dataset_columns}
- Avg tokens per sample: {avg_tokens}

Respond with a JSON object:
{{
  "recommended_model": "<model_id>",
  "reason": "<1-2 sentence explanation>",
  "alternatives": ["<model_id_1>", "<model_id_2>"]
}}

Only recommend models from the available list. Prioritize:
1. Models that fit comfortably in VRAM
2. Models suited for the task (code -> CodeLlama, general -> Llama/Qwen)
3. Smaller models for small datasets (<1K rows)"""


@dataclass
class ModelSuggestion:
    """Result of model selection."""

    recommended_model: str
    reason: str
    alternatives: list[str] = field(default_factory=list)
    source: str = "rule_engine"  # "rule_engine" or "llm"


def _get_available_vram(hardware_dict: dict) -> float:
    """Get total available VRAM from hardware profile dict."""
    gpus = hardware_dict.get("gpus", [])
    if not gpus:
        return 0.0
    return sum(g.get("vram_gb", 0.0) for g in gpus)


def _filter_candidates(vram_gb: float, quantization: str = "4bit") -> list[dict]:
    """Filter models that fit in available VRAM with given quantization."""
    quant_factor = {"4bit": 0.25, "8bit": 0.5, "none": 1.0}.get(quantization, 0.25)
    candidates = []
    for model in MODEL_CATALOG:
        effective_vram = model["min_vram_gb"] * quant_factor
        if effective_vram <= vram_gb * 0.85:  # 15% safety margin
            candidates.append(model)
    return candidates


def suggest_model_rule_based(
    hardware_dict: dict,
    dataset_dict: dict,
) -> ModelSuggestion:
    """Suggest a model using rule-based heuristics (no LLM needed)."""
    vram_gb = _get_available_vram(hardware_dict)
    num_rows = dataset_dict.get("num_rows", 1000)
    columns = dataset_dict.get("columns", [])

    # Determine task type from dataset
    col_str = " ".join(c.lower() for c in columns)
    is_code_task = any(kw in col_str for kw in ["code", "programming", "function", "snippet"])

    # Filter by VRAM
    candidates = _filter_candidates(vram_gb, "4bit")
    if not candidates:
        # Even 4bit doesn't fit — suggest smallest model
        return ModelSuggestion(
            recommended_model="microsoft/phi-3-mini-4k-instruct",
            reason=(
                "Very limited VRAM; Phi-3 Mini (3.8B) is the smallest viable model for fine-tuning."
            ),
            alternatives=[],
            source="rule_engine",
        )

    # Prefer code models for code tasks
    if is_code_task:
        code_models = [m for m in candidates if "code" in m.get("strengths", [])]
        if code_models:
            best = code_models[0]
            alts = [m["id"] for m in code_models[1:2]] + [
                m["id"] for m in candidates if m not in code_models
            ][:1]
            return ModelSuggestion(
                recommended_model=best["id"],
                reason=(
                    f"Code task detected; {best['id']} ({best['params']}) "
                    f"is optimized for code and fits in {vram_gb:.0f}GB VRAM."
                ),
                alternatives=alts,
                source="rule_engine",
            )

    # For small datasets, prefer smaller models
    if num_rows < 1000:
        small_models = [m for m in candidates if m["min_vram_gb"] <= 8]
        if small_models:
            candidates = small_models

    # Default: pick the largest model that fits
    candidates.sort(key=lambda m: m["min_vram_gb"], reverse=True)
    best = candidates[0]
    alts = [m["id"] for m in candidates[1:3]]

    return ModelSuggestion(
        recommended_model=best["id"],
        reason=(
            f"{best['id']} ({best['params']}) is the best fit for "
            f"{vram_gb:.0f}GB VRAM with {num_rows} training samples."
        ),
        alternatives=alts,
        source="rule_engine",
    )


def suggest_model_with_llm(
    hardware_dict: dict,
    dataset_dict: dict,
    llm_config: dict | None = None,
) -> ModelSuggestion:
    """Suggest a model using LLM enhancement over rule-based selection.

    Falls back to rule-based if LLM is unavailable.
    """
    # Start with rule-based suggestion as fallback
    rule_suggestion = suggest_model_rule_based(hardware_dict, dataset_dict)

    llm_config = llm_config or {}
    try:
        from bilgeai.llm import get_llm_client

        client = get_llm_client(
            provider=llm_config.get("provider", "anthropic"),
            model=llm_config.get("model", "claude-sonnet-4-20250514"),
            api_key_env=llm_config.get("api_key_env", "ANTHROPIC_API_KEY"),
            base_url=llm_config.get("base_url"),
        )

        vram_gb = _get_available_vram(hardware_dict)
        candidates = _filter_candidates(vram_gb, "4bit")
        candidate_str = "\n".join(
            f"- {m['id']} ({m['params']}, min {m['min_vram_gb']}GB, "
            f"strengths: {', '.join(m['strengths'])})"
            for m in candidates
        )

        prompt = SUGGEST_MODEL_PROMPT.format(
            hardware_json=json.dumps(
                {k: v for k, v in hardware_dict.items() if k != "gpus"},
                indent=2,
            ),
            dataset_json=json.dumps(
                {
                    k: v
                    for k, v in dataset_dict.items()
                    if k not in ("quality_issues", "pii_findings")
                },
                indent=2,
            ),
            candidate_models=candidate_str,
            dataset_format=dataset_dict.get("format", "unknown"),
            dataset_columns=", ".join(dataset_dict.get("columns", [])),
            avg_tokens=dataset_dict.get("avg_tokens", "unknown"),
        )

        response = client.generate(prompt)
        result = _parse_suggestion(response.content)

        if result:
            return ModelSuggestion(
                recommended_model=result.get(
                    "recommended_model", rule_suggestion.recommended_model
                ),
                reason=result.get("reason", rule_suggestion.reason),
                alternatives=result.get("alternatives", []),
                source=client.source_name,
            )

    except Exception:
        logger.debug("LLM model suggestion failed, using rule engine", exc_info=True)

    return rule_suggestion


def _parse_suggestion(content: str) -> dict | None:
    """Parse LLM response to extract model suggestion JSON."""
    content = content.strip()

    # Handle markdown code fences
    if "```" in content:
        parts = content.split("```")
        for part in parts:
            part = part.strip()
            if part.startswith("json"):
                part = part[4:].strip()
            try:
                return json.loads(part)
            except (json.JSONDecodeError, ValueError):
                continue

    try:
        return json.loads(content)
    except (json.JSONDecodeError, ValueError):
        return None
