from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional


DEFAULT_MODELS: dict[str, str] = {
    "transformers": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    "ollama": "llama3.2:1b",
    "inference": "HuggingFaceH4/zephyr-7b-beta",
}

OLLAMA_PREFERRED = ["llama3.2:1b", "llama3.2", "llama3", "phi3:mini", "mistral", "gemma"]


@dataclass
class ModelPreset:
    model_id: str
    description: str
    backend: str
    max_new_tokens: int = 256
    temperature: float = 0.7
    system_prompt: str = "You are a helpful AI assistant. Be concise and clear."
    extra_kwargs: dict = field(default_factory=dict)


PRESETS: dict[str, ModelPreset] = {
    "tinyllama": ModelPreset(
        model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        description="TinyLlama 1.1B — fast, instruction-tuned, runs on any CPU.",
        backend="transformers",
        max_new_tokens=256,
    ),
    "qwen": ModelPreset(
        model_id="Qwen/Qwen2.5-0.5B-Instruct",
        description="Qwen 2.5 0.5B — tiny but surprisingly capable.",
        backend="transformers",
        max_new_tokens=256,
    ),
    "phi2": ModelPreset(
        model_id="microsoft/phi-2",
        description="Microsoft Phi-2 2.7B — excellent reasoning for its size.",
        backend="transformers",
        max_new_tokens=512,
    ),
    "zephyr": ModelPreset(
        model_id="HuggingFaceH4/zephyr-7b-beta",
        description="Zephyr 7B — strong instruction-following (needs GPU).",
        backend="transformers",
        max_new_tokens=512,
    ),
    "mistral": ModelPreset(
        model_id="mistralai/Mistral-7B-Instruct-v0.2",
        description="Mistral 7B Instruct — fast and capable (needs GPU).",
        backend="transformers",
        max_new_tokens=512,
    ),
    "llama3": ModelPreset(
        model_id="llama3.2:1b",
        description="LLaMA 3.2 1B via Ollama — great quality, no GPU needed.",
        backend="ollama",
        max_new_tokens=512,
    ),
    "llama3-3b": ModelPreset(
        model_id="llama3.2:3b",
        description="LLaMA 3.2 3B via Ollama — excellent quality.",
        backend="ollama",
        max_new_tokens=512,
    ),
    "phi3": ModelPreset(
        model_id="phi3:mini",
        description="Microsoft Phi-3 mini via Ollama — very capable small model.",
        backend="ollama",
        max_new_tokens=512,
    ),
}


def get_preset(name: str) -> Optional[ModelPreset]:
    return PRESETS.get(name)


def list_presets() -> dict[str, str]:
    return {name: preset.description for name, preset in PRESETS.items()}
