#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Swarm Providers: Model registry and provider detection for PULSE swarm.

Functions:
  - _detect_providers(): Available providers (CLI, API, Ollama)
  - _resolve_access(): Best access method for a provider
  - _find_provider_for_model(): Reverse lookup model->provider
  - _ollama_models(): Dynamic model list from local Ollama
"""
from __future__ import annotations

import json
import os
import shutil
from urllib.request import urlopen
from urllib.error import URLError

# ---------------------------------------------------------------------------
# Model registry — one list per company (no duplication)
# ---------------------------------------------------------------------------
MODELS: dict[str, list[tuple[str, str]]] = {
    "claude": [
        ("claude-haiku-4-5-20251001", "fast"),
        ("claude-sonnet-4-6", "standard"),
        ("claude-opus-4-6", "premium"),
    ],
    "gemini": [
        ("gemini-2.5-flash", "fast"),
        ("gemini-2.5-pro", "standard"),
        ("gemini-3-flash-preview", "fast"),
        ("gemini-3-pro-preview", "premium"),
        ("gemini-3.1-pro-preview", "premium"),
    ],
    "openai": [
        ("o4-mini", "fast"),
        ("gpt-5-mini", "fast"),
        ("gpt-5-nano", "fast"),
        ("gpt-5.2-codex", "standard"),
        ("gpt-5.3-codex", "premium"),
        ("gpt-5.4", "premium"),
        ("gpt-5.4-pro", "premium"),
        ("gpt-5.1-codex-max", "premium"),
        ("o3", "premium"),
    ],
    "ollama": [  # static cloud models (always shown) + dynamically detected local
        ("qwen3-coder-next:cloud", "premium"),
        ("qwen3.5:cloud", "premium"),
        ("devstral-2:cloud", "premium"),
        ("glm-5:cloud", "premium"),
        ("minimax-m2.5:cloud", "premium"),
        ("kimi-k2.5:cloud", "premium"),
    ],
}

# One entry per company — auto-detect CLI vs API at launch time
_PROVIDER_META = {
    "claude": {
        "label": "Claude (Anthropic)",
        "cli": "claude",
        "api_env": "ANTHROPIC_API_KEY",
        "worker_provider": "anthropic",
    },
    "gemini": {
        "label": "Gemini (Google)",
        "cli": "gemini",
        "api_env": "GEMINI_API_KEY",
        "worker_provider": "google",
    },
    "openai": {
        "label": "OpenAI / Codex",
        "cli": "codex",
        "api_env": "OPENAI_API_KEY",
        "worker_provider": "openai",
    },
    "ollama": {
        "label": "Ollama (local + cloud, free)",
        "worker_provider": "ollama",
    },
}

# Tier mapping for known Ollama models (prefix match, checked in order)
_OLLAMA_TIER_MAP: list[tuple[str, str]] = [
    # Premium (cloud frontier + large local)
    ("qwen3-coder-next", "premium"),
    ("qwen3.5", "premium"),
    ("devstral-2", "premium"),
    ("glm-5", "premium"),
    ("minimax-m2.5", "premium"),
    ("kimi-k2.5", "premium"),
    ("qwen3-coder", "premium"),
    ("gpt-oss:120b", "premium"),
    ("deepseek-r1:671b", "premium"),
    ("llama3.3:70b", "premium"),
    ("llama3.1:70b", "premium"),
    ("qwen2.5:72b", "premium"),
    # Standard (mid-size local + mid-tier cloud)
    ("qwen3-next", "standard"),
    ("nemotron-3-nano", "standard"),
    ("qwen2.5-coder:32b", "standard"),
    ("deepseek-coder-v2", "standard"),
    ("deepseek-r1:32b", "standard"),
    ("codellama:34b", "standard"),
    ("codellama:70b", "standard"),
    ("deepseek-r1:14b", "standard"),
    # Fast (small cloud + small local)
    ("devstral-small-2", "fast"),
    ("glm-4.7-flash", "fast"),
    ("ministral-3", "fast"),
    ("rnj-1", "fast"),
    # Everything else → fast
]


def _ollama_tier(name: str) -> str:
    """Detect tier for an Ollama model by prefix match."""
    lower = name.lower()
    for prefix, tier in _OLLAMA_TIER_MAP:
        if lower.startswith(prefix):
            return tier
    return "fast"


# ---------------------------------------------------------------------------
# Provider detection
# ---------------------------------------------------------------------------

def _ollama_running() -> bool:
    try:
        urlopen("http://localhost:11434/api/tags", timeout=2)
        return True
    except (URLError, OSError):
        return False


def _ollama_models() -> list[tuple[str, str]]:
    """Return cloud models first, then local models (no dupes)."""
    all_models = list(MODELS.get("ollama", []))
    known_names = {name for name, _ in all_models}
    try:
        resp = urlopen("http://localhost:11434/api/tags", timeout=3)
        data = json.loads(resp.read())
        for m in data.get("models", []):
            name = m.get("name")
            if name and name not in known_names:
                all_models.append((name, _ollama_tier(name)))
                known_names.add(name)
    except Exception:
        pass
    # Group: cloud first, then local
    cloud = [(n, t) for n, t in all_models if ":cloud" in n]
    local = [(n, t) for n, t in all_models if ":cloud" not in n]
    return cloud + local


def _detect_providers() -> list[tuple[str, str]]:
    """Return [(provider_key, label)] for available providers."""
    available = []
    for key, meta in _PROVIDER_META.items():
        if key == "ollama":
            if _ollama_running():
                available.append((key, meta["label"]))
            continue
        has_cli = bool(meta.get("cli") and shutil.which(meta["cli"]))
        has_api = bool(meta.get("api_env") and os.environ.get(meta["api_env"]))
        if has_cli or has_api:
            methods = []
            if has_cli:
                methods.append("CLI")
            if has_api:
                methods.append("API")
            available.append((key, f"{meta['label']} ({' + '.join(methods)})"))
    return available


def _resolve_access(provider_key: str) -> str:
    """Determine best access method: 'cli', 'api', or 'local'. Prefers CLI."""
    meta = _PROVIDER_META.get(provider_key, {})
    if provider_key == "ollama":
        return "local"
    if meta.get("cli") and shutil.which(meta["cli"]):
        return "cli"
    if meta.get("api_env") and os.environ.get(meta["api_env"]):
        return "api"
    return "cli"


def _find_provider_for_model(model: str) -> str | None:
    """Given a model name, find which provider serves it."""
    if ":" in model:
        return "ollama" if _ollama_running() else None
    for key, model_list in MODELS.items():
        for m_name, _ in model_list:
            if m_name == model:
                if any(pk == key for pk, _ in _detect_providers()):
                    return key
    return None
