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


DEFAULT_MODELS: dict[str, str] = {
    "transformers": "Qwen/Qwen2.5-1.5B-Instruct",
    "ollama": "llama3.2:3b",
    "openai": "gpt-4o-mini",
    "inference": "HuggingFaceH4/zephyr-7b-beta",
}

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


@dataclass
class ModelPreset:
    model_id: str
    description: str
    backend: str
    max_new_tokens: int = 512
    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] = {
    # ── OpenAI ──────────────────────────────────────────────────────────────
    "gpt4o-mini": ModelPreset(
        model_id="gpt-4o-mini",
        description="GPT-4o mini — fast, cheap, very capable. Needs OPENAI_API_KEY.",
        backend="openai",
        max_new_tokens=1024,
    ),
    "gpt4o": ModelPreset(
        model_id="gpt-4o",
        description="GPT-4o — best OpenAI model. Needs OPENAI_API_KEY.",
        backend="openai",
        max_new_tokens=2048,
    ),
    "gpt35": ModelPreset(
        model_id="gpt-3.5-turbo",
        description="GPT-3.5 Turbo — fast and affordable. Needs OPENAI_API_KEY.",
        backend="openai",
        max_new_tokens=1024,
    ),
    # ── Local (Transformers) ─────────────────────────────────────────────────
    "qwen": ModelPreset(
        model_id="Qwen/Qwen2.5-1.5B-Instruct",
        description="Qwen 2.5 1.5B — strong instruction-following, runs on CPU.",
        backend="transformers",
        max_new_tokens=512,
    ),
    "qwen-small": ModelPreset(
        model_id="Qwen/Qwen2.5-0.5B-Instruct",
        description="Qwen 2.5 0.5B — tiny and fast.",
        backend="transformers",
        max_new_tokens=256,
    ),
    "tinyllama": ModelPreset(
        model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        description="TinyLlama 1.1B — lightweight and fast.",
        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,
    ),
    # ── Ollama ───────────────────────────────────────────────────────────────
    "llama3": ModelPreset(
        model_id="llama3.2:3b",
        description="LLaMA 3.2 3B via Ollama — great quality, fully local.",
        backend="ollama",
        max_new_tokens=1024,
    ),
    "llama3-small": ModelPreset(
        model_id="llama3.2:1b",
        description="LLaMA 3.2 1B via Ollama — fast and light.",
        backend="ollama",
        max_new_tokens=512,
    ),
    "phi3": ModelPreset(
        model_id="phi3:mini",
        description="Microsoft Phi-3 mini via Ollama — very capable.",
        backend="ollama",
        max_new_tokens=1024,
    ),
    "mistral": ModelPreset(
        model_id="mistral",
        description="Mistral 7B via Ollama — strong all-rounder.",
        backend="ollama",
        max_new_tokens=1024,
    ),
}


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()}
