"""
Pre-configured model presets for common use cases.
"""

from dataclasses import dataclass, field
from typing import Optional


@dataclass
class ModelPreset:
    model_id: str
    description: str
    backend: str
    max_new_tokens: int = 512
    temperature: float = 0.7
    system_prompt: str = ""
    extra_kwargs: dict = field(default_factory=dict)


PRESETS: dict[str, ModelPreset] = {
    "tiny": ModelPreset(
        model_id="facebook/blenderbot-400M-distill",
        description="Small, fast BlenderBot model (~400M params). Good for testing.",
        backend="transformers",
        max_new_tokens=256,
        temperature=0.7,
    ),
    "dialogue": ModelPreset(
        model_id="microsoft/DialoGPT-medium",
        description="Microsoft DialoGPT medium (~345M params). Conversational.",
        backend="transformers",
        max_new_tokens=256,
        temperature=0.8,
    ),
    "zephyr": ModelPreset(
        model_id="HuggingFaceH4/zephyr-7b-beta",
        description="Zephyr 7B – a strong instruction-tuned model.",
        backend="transformers",
        max_new_tokens=512,
        temperature=0.7,
        system_prompt="You are a helpful, respectful assistant.",
    ),
    "mistral": ModelPreset(
        model_id="mistralai/Mistral-7B-Instruct-v0.2",
        description="Mistral 7B Instruct – fast and capable.",
        backend="transformers",
        max_new_tokens=512,
        temperature=0.7,
        system_prompt="You are a helpful AI assistant.",
    ),
    "ollama-llama3": ModelPreset(
        model_id="llama3",
        description="Meta LLaMA 3 via Ollama (must have Ollama running locally).",
        backend="ollama",
        max_new_tokens=512,
        temperature=0.7,
        system_prompt="You are a helpful AI assistant.",
    ),
    "ollama-mistral": ModelPreset(
        model_id="mistral",
        description="Mistral via Ollama (must have Ollama running locally).",
        backend="ollama",
        max_new_tokens=512,
        temperature=0.7,
    ),
    "ollama-gemma": ModelPreset(
        model_id="gemma",
        description="Google Gemma via Ollama (must have Ollama running locally).",
        backend="ollama",
        max_new_tokens=512,
        temperature=0.7,
    ),
}


def get_preset(name: str) -> Optional[ModelPreset]:
    """Return a preset by name, or None if not found."""
    return PRESETS.get(name)


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