from typing import Dict, Any

# ── Model Registry ──────────────────────────────────────────────────────
MODEL_REGISTRY: Dict[str, Dict[str, Any]] = {
    # ── Original models (fixed) ─────────────────────────────────────────
    "HuggingFaceTB/SmolLM2-360M-Instruct": {
        "label": "SmolLM2 360M",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 2048,
    },
    "Qwen/Qwen2.5-0.5B-Instruct": {
        "label": "Qwen2.5 0.5B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 32768,
    },
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0": {
        "label": "TinyLlama 1.1B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 2048,
    },
    "Qwen/Qwen2.5-1.5B-Instruct": {
        "label": "Qwen2.5 1.5B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 32768,
    },
    "microsoft/Phi-3.5-mini-instruct": {
        "label": "Phi-3.5 mini",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 128000,
    },
    "microsoft/Phi-3-mini-4k-instruct": {
        "label": "Phi-3 mini 4K",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 4096,
    },
    "Qwen/Qwen2.5-3B-Instruct": {
        "label": "Qwen2.5 3B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 32768,
    },
    "google/gemma-3-1b-it": {
        "label": "Gemma-3 1B IT",
        "family": "causal-lm",  # fixed typo: was "casual-lm"
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 32768,
    },  # fixed: missing closing brace
    "openbmb/MiniCPM5-2B": {
        "label": "MiniCPM5 2B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 131072,
    },
    "constructai/Qwenite3.5-2B": {
        "label": "Qwenite 3.5 2B · T4 trained reasoning",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 32768,
    },

    # ── Added 1–3B models ───────────────────────────────────────────────
    "HuggingFaceTB/SmolLM2-1.7B-Instruct": {
        "label": "SmolLM2 1.7B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 8192,
    },
    "google/gemma-2-2b-it": {
        "label": "Gemma-2 2B IT",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 8192,
    },
    "meta-llama/Llama-3.2-1B-Instruct": {
        "label": "Llama-3.2 1B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 131072,
    },
    "meta-llama/Llama-3.2-3B-Instruct": {
        "label": "Llama-3.2 3B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 131072,
    },
    "stabilityai/stablelm-2-1_6b-chat": {
        "label": "StableLM-2 1.6B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 4096,
    },
    "stabilityai/stablelm-zephyr-3b": {
        "label": "StableLM Zephyr 3B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 4096,
    },
    "allenai/OLMo-1B-hf": {
        "label": "OLMo 1B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 2048,
    },
    "HuggingFaceTB/SmolLM-1.7B-Instruct": {
        "label": "SmolLM 1.7B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 2048,
    },
    "google/gemma-2b-it": {
        "label": "Gemma 2B IT",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 8192,
    },
    "Qwen/Qwen2-1.5B-Instruct": {
        "label": "Qwen2 1.5B",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 32768,
    },
    "HuggingFaceTB/SmolLM2-1.7B-Instruct-16k": {
        "label": "SmolLM2 1.7B · 16K",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 16384,
    },
    "Qwen/Qwen2.5-1.5B-Instruct-GGUF": {
        "label": "Qwen2.5 1.5B GGUF",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 32768,
    },
    "microsoft/Phi-3-mini-128k-instruct": {
        "label": "Phi-3 mini 128K",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 128000,
    },
    "google/gemma-3-1b-it-GGUF": {
        "label": "Gemma-3 1B IT GGUF",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 32768,
    },
    "meta-llama/Llama-3.2-1B-Instruct-GGUF": {
        "label": "Llama-3.2 1B GGUF",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 131072,
    },
    "meta-llama/Llama-3.2-3B-Instruct-GGUF": {
        "label": "Llama-3.2 3B GGUF",
        "family": "causal-lm",
        "tasks": ["chat", "code", "summarize", "translate", "agent"],
        "context": 131072,
    },
}

DEFAULT_MODEL = "HuggingFaceTB/SmolLM2-360M-Instruct"

# ── Task Registry ───────────────────────────────────────────────────────
TASK_REGISTRY: Dict[str, Dict[str, Any]] = {
    "chat": {
        "label": "💬 Chat",
        "description": "Open-ended conversation",
        "system_prompt": "You are a friendly, helpful conversational assistant.",
    },
    "code": {
        "label": "💻 Code",
        "description": "Generate, explain, or debug code",
        "system_prompt": (
            "You are a senior software engineer. Produce clean, idiomatic, "
            "well-commented code. Wrap code in fenced markdown blocks."
        ),
    },
    "summarize": {
        "label": "📝 Summarize",
        "description": "Condense long text into key points",
        "system_prompt": (
            "Summarise the user's text in 5–7 bullet points. "
            "Be concise and faithful to the source."
        ),
    },
    "translate": {
        "label": "🌍 Translate",
        "description": "Translate text between languages",
        "system_prompt": (
            "Translate the user's text. If they don't specify a target "
            "language, translate to English."
        ),
    },
    "agent": {
        "label": "🤖 Agent",
        "description": "Autonomous tool-using agent",
        "system_prompt": (
            "You are an autonomous agent. Use tools when they help, "
            "then give a final answer."
        ),
    },
}