"""
Core data models and enums for PAGANS.

This module contains model family detection and model-name mappings used to
route optimization through OpenRouter while targeting provider-specific models.
"""

from dataclasses import dataclass
from enum import Enum

from .constants import (
    FAMILY_ANTHROPIC,
    FAMILY_GOOGLE,
    FAMILY_OPENAI,
    FAMILY_XAI,
)


class Provider(Enum):
    """Enum representing available inference providers for PAGANS."""

    OPENROUTER = "openrouter"
    OPENAI = "openai"
    GOOGLE = "google"
    ANTHROPIC = "anthropic"
    ZAI = "zai"


class ModelFamily(Enum):
    """Enum representing supported LLM model families."""

    OPENAI = FAMILY_OPENAI
    ANTHROPIC = FAMILY_ANTHROPIC
    GOOGLE = FAMILY_GOOGLE
    XAI = FAMILY_XAI


@dataclass
class OptimizationResult:
    """Result of a prompt optimization operation."""

    original: str
    optimized: str
    target_model: str
    target_family: ModelFamily
    optimization_notes: str | None = None
    tokens_used: int | None = None
    optimization_time: float | None = None


@dataclass
class OptimizationRequest:
    """Request for prompt optimization."""

    prompt: str
    target_model: str
    optimization_notes: str | None = None


# Short names users can pass to PAGANS.
SHORT_MODEL_NAMES: dict[str, ModelFamily] = {
    # OpenAI (latest + backward-compatible aliases)
    "gpt-5.2": ModelFamily.OPENAI,
    "gpt-5.2-mini": ModelFamily.OPENAI,
    "gpt-5.2-nano": ModelFamily.OPENAI,
    "gpt-5.1": ModelFamily.OPENAI,
    "gpt-5.4": ModelFamily.OPENAI,
    "gpt-5.4-pro": ModelFamily.OPENAI,
    "gpt-5.4-mini": ModelFamily.OPENAI,
    "gpt-5.4-nano": ModelFamily.OPENAI,
    "gpt-5-mini": ModelFamily.OPENAI,
    "gpt-5-nano": ModelFamily.OPENAI,
    "gpt-5": ModelFamily.OPENAI,
    "gpt-4.1": ModelFamily.OPENAI,
    "gpt-4o": ModelFamily.OPENAI,
    # Anthropic (latest + backward-compatible aliases)
    "claude-opus-4-20250514": ModelFamily.ANTHROPIC,
    "claude-sonnet-4-20250514": ModelFamily.ANTHROPIC,
    "claude-opus-4-1-20250805": ModelFamily.ANTHROPIC,
    "claude-sonnet-4-5-20250929": ModelFamily.ANTHROPIC,
    "claude-haiku-4-5-20251001": ModelFamily.ANTHROPIC,
    "claude-3-7-sonnet-20250219": ModelFamily.ANTHROPIC,
    "claude-opus-4.6": ModelFamily.ANTHROPIC,
    "claude-sonnet-4.6": ModelFamily.ANTHROPIC,
    "claude-haiku-4.5": ModelFamily.ANTHROPIC,
    "claude-opus-4.1": ModelFamily.ANTHROPIC,
    "claude-opus-4": ModelFamily.ANTHROPIC,
    "claude-sonnet-4": ModelFamily.ANTHROPIC,
    "claude-3.5-sonnet": ModelFamily.ANTHROPIC,
    # Google Gemini (latest text models + stable line)
    "gemini-2.5-flash-lite": ModelFamily.GOOGLE,
    "gemini-2.5-flash-lite-preview-09-2025": ModelFamily.GOOGLE,
    "gemini-2.0-flash": ModelFamily.GOOGLE,
    "gemini-3.1-pro-preview": ModelFamily.GOOGLE,
    "gemini-3-flash-preview": ModelFamily.GOOGLE,
    "gemini-3.1-flash-lite-preview": ModelFamily.GOOGLE,
    "gemini-2.5-pro": ModelFamily.GOOGLE,
    "gemini-2.5-flash": ModelFamily.GOOGLE,
    # xAI Grok (text models)
    "grok-4-1-fast-reasoning": ModelFamily.XAI,
    "grok-4-1-fast-non-reasoning": ModelFamily.XAI,
    "grok-4-fast-reasoning": ModelFamily.XAI,
    "grok-4-fast-non-reasoning": ModelFamily.XAI,
    "grok-4.20-beta": ModelFamily.XAI,
    "grok-4.20-multi-agent-beta": ModelFamily.XAI,
    "grok-4": ModelFamily.XAI,
    "grok-4-fast": ModelFamily.XAI,
    "grok-code-fast-1": ModelFamily.XAI,
}


# OpenRouter routed model ids.
OPENROUTER_MODEL_MAPPINGS: dict[str, str] = {
    # OpenAI
    "gpt-5.2": "openai/gpt-5.2",
    "gpt-5.2-mini": "openai/gpt-5.2-mini",
    "gpt-5.2-nano": "openai/gpt-5.2-nano",
    "gpt-5.1": "openai/gpt-5.1",
    "gpt-5.4": "openai/gpt-5.4",
    "gpt-5.4-pro": "openai/gpt-5.4-pro",
    "gpt-5.4-mini": "openai/gpt-5.4-mini",
    "gpt-5.4-nano": "openai/gpt-5.4-nano",
    "gpt-5-mini": "openai/gpt-5.4-mini",
    "gpt-5-nano": "openai/gpt-5.4-nano",
    "gpt-5": "openai/gpt-5",
    "gpt-4.1": "openai/gpt-4.1",
    "gpt-4o": "openai/gpt-4o",
    # Anthropic
    "claude-opus-4-20250514": "anthropic/claude-opus-4-20250514",
    "claude-sonnet-4-20250514": "anthropic/claude-sonnet-4-20250514",
    "claude-opus-4-1-20250805": "anthropic/claude-opus-4-1-20250805",
    "claude-sonnet-4-5-20250929": "anthropic/claude-sonnet-4-5-20250929",
    "claude-haiku-4-5-20251001": "anthropic/claude-haiku-4-5-20251001",
    "claude-3-7-sonnet-20250219": "anthropic/claude-3-7-sonnet-20250219",
    "claude-opus-4.6": "anthropic/claude-opus-4.6",
    "claude-sonnet-4.6": "anthropic/claude-sonnet-4.6",
    "claude-haiku-4.5": "anthropic/claude-haiku-4.5",
    "claude-opus-4.1": "anthropic/claude-opus-4.1",
    "claude-opus-4": "anthropic/claude-opus-4",
    "claude-sonnet-4": "anthropic/claude-sonnet-4",
    "claude-3.5-sonnet": "anthropic/claude-3.5-sonnet",
    # Google
    "gemini-2.5-flash-lite": "google/gemini-2.5-flash-lite",
    "gemini-2.5-flash-lite-preview-09-2025": "google/gemini-2.5-flash-lite-preview-09-2025",
    "gemini-2.0-flash": "google/gemini-2.0-flash",
    "gemini-3.1-pro-preview": "google/gemini-3.1-pro-preview",
    "gemini-3-flash-preview": "google/gemini-3-flash-preview",
    "gemini-3.1-flash-lite-preview": "google/gemini-3.1-flash-lite-preview",
    "gemini-2.5-pro": "google/gemini-2.5-pro",
    "gemini-2.5-flash": "google/gemini-2.5-flash",
    # xAI
    "grok-4-1-fast-reasoning": "x-ai/grok-4-1-fast-reasoning",
    "grok-4-1-fast-non-reasoning": "x-ai/grok-4-1-fast-non-reasoning",
    "grok-4-fast-reasoning": "x-ai/grok-4-fast-reasoning",
    "grok-4-fast-non-reasoning": "x-ai/grok-4-fast-non-reasoning",
    "grok-4.20-beta": "x-ai/grok-4.20-beta",
    "grok-4.20-multi-agent-beta": "x-ai/grok-4.20-multi-agent-beta",
    "grok-4": "x-ai/grok-4",
    "grok-4-fast": "x-ai/grok-4-fast",
    "grok-code-fast-1": "x-ai/grok-code-fast-1",
}


def _strip_provider_prefix(model_id: str, prefix: str) -> str:
    if model_id.startswith(prefix):
        return model_id[len(prefix) :]
    return model_id


OPENAI_MODEL_MAPPINGS: dict[str, str] = {
    short_name: _strip_provider_prefix(provider_id, "openai/")
    for short_name, provider_id in OPENROUTER_MODEL_MAPPINGS.items()
    if provider_id.startswith("openai/")
}

GOOGLE_MODEL_MAPPINGS: dict[str, str] = {
    short_name: _strip_provider_prefix(provider_id, "google/")
    for short_name, provider_id in OPENROUTER_MODEL_MAPPINGS.items()
    if provider_id.startswith("google/")
}

ANTHROPIC_MODEL_MAPPINGS: dict[str, str] = {
    short_name: _strip_provider_prefix(provider_id, "anthropic/")
    for short_name, provider_id in OPENROUTER_MODEL_MAPPINGS.items()
    if provider_id.startswith("anthropic/")
}

# Z.ai is used as an optimizer provider. Model aliases are currently user-specified.
ZAI_MODEL_MAPPINGS: dict[str, str] = {}

PROVIDER_MODEL_MAPPINGS: dict[Provider, dict[str, str]] = {
    Provider.OPENROUTER: OPENROUTER_MODEL_MAPPINGS,
    Provider.OPENAI: OPENAI_MODEL_MAPPINGS,
    Provider.GOOGLE: GOOGLE_MODEL_MAPPINGS,
    Provider.ANTHROPIC: ANTHROPIC_MODEL_MAPPINGS,
    Provider.ZAI: ZAI_MODEL_MAPPINGS,
}


MODEL_MAPPINGS: dict[ModelFamily, list[str]] = {
    ModelFamily.OPENAI: [
        "gpt-5.2",
        "gpt-5.2-mini",
        "gpt-5.2-nano",
        "gpt-5.1",
        "gpt-5.4",
        "gpt-5.4-pro",
        "gpt-5.4-mini",
        "gpt-5.4-nano",
        "gpt-5-mini",
        "gpt-5-nano",
        "gpt-5",
        "gpt-4.1",
        "gpt-4o",
        "openai/gpt-5.2",
        "openai/gpt-5.2-mini",
        "openai/gpt-5.2-nano",
        "openai/gpt-5.1",
        "openai/gpt-5.4",
        "openai/gpt-5.4-pro",
        "openai/gpt-5.4-mini",
        "openai/gpt-5.4-nano",
        "openai/gpt-5",
        "openai/gpt-4.1",
        "openai/gpt-4o",
    ],
    ModelFamily.ANTHROPIC: [
        "claude-opus-4-20250514",
        "claude-sonnet-4-20250514",
        "claude-opus-4-1-20250805",
        "claude-sonnet-4-5-20250929",
        "claude-haiku-4-5-20251001",
        "claude-3-7-sonnet-20250219",
        "claude-opus-4.6",
        "claude-sonnet-4.6",
        "claude-haiku-4.5",
        "claude-opus-4.1",
        "claude-opus-4",
        "claude-sonnet-4",
        "claude-3.5-sonnet",
        "anthropic/claude-opus-4-20250514",
        "anthropic/claude-sonnet-4-20250514",
        "anthropic/claude-opus-4-1-20250805",
        "anthropic/claude-sonnet-4-5-20250929",
        "anthropic/claude-haiku-4-5-20251001",
        "anthropic/claude-3-7-sonnet-20250219",
        "anthropic/claude-opus-4.6",
        "anthropic/claude-sonnet-4.6",
        "anthropic/claude-haiku-4.5",
        "anthropic/claude-opus-4.1",
        "anthropic/claude-opus-4",
        "anthropic/claude-sonnet-4",
        "anthropic/claude-3.5-sonnet",
    ],
    ModelFamily.GOOGLE: [
        "gemini-2.5-flash-lite",
        "gemini-2.5-flash-lite-preview-09-2025",
        "gemini-2.0-flash",
        "gemini-3.1-pro-preview",
        "gemini-3-flash-preview",
        "gemini-3.1-flash-lite-preview",
        "gemini-2.5-pro",
        "gemini-2.5-flash",
        "google/gemini-2.5-flash-lite",
        "google/gemini-2.5-flash-lite-preview-09-2025",
        "google/gemini-2.0-flash",
        "google/gemini-3.1-pro-preview",
        "google/gemini-3-flash-preview",
        "google/gemini-3.1-flash-lite-preview",
        "google/gemini-2.5-pro",
        "google/gemini-2.5-flash",
    ],
    ModelFamily.XAI: [
        "grok-4-1-fast-reasoning",
        "grok-4-1-fast-non-reasoning",
        "grok-4-fast-reasoning",
        "grok-4-fast-non-reasoning",
        "grok-4.20-beta",
        "grok-4.20-multi-agent-beta",
        "grok-4",
        "grok-4-fast",
        "grok-code-fast-1",
        "x-ai/grok-4-1-fast-reasoning",
        "x-ai/grok-4-1-fast-non-reasoning",
        "x-ai/grok-4-fast-reasoning",
        "x-ai/grok-4-fast-non-reasoning",
        "x-ai/grok-4.20-beta",
        "x-ai/grok-4.20-multi-agent-beta",
        "x-ai/grok-4",
        "x-ai/grok-4-fast",
        "x-ai/grok-code-fast-1",
    ],
}


MODEL_NAME_MAPPINGS: dict[str, ModelFamily] = {}
for family, models in MODEL_MAPPINGS.items():
    for model in models:
        MODEL_NAME_MAPPINGS[model] = family


FAMILY_MODEL_MAPPINGS: dict[str, list[str]] = {
    family.value: models for family, models in MODEL_MAPPINGS.items()
}


def detect_model_family(model_name: str) -> ModelFamily:
    """Detect the model family from a model name."""
    normalized = model_name.lower().strip()

    for family, models in MODEL_MAPPINGS.items():
        for model in models:
            if model.lower() == normalized:
                return family

    for family, models in MODEL_MAPPINGS.items():
        for model in models:
            lower_model = model.lower()
            if normalized in lower_model or lower_model in normalized:
                return family

    msg = f"Unknown model family for: {normalized}"
    raise ValueError(msg)


def get_supported_models() -> dict[ModelFamily, list[str]]:
    """Get all supported models organized by family."""
    return MODEL_MAPPINGS.copy()


def is_supported_model(model_name: str) -> bool:
    """Check if a model is supported."""
    try:
        detect_model_family(model_name)
        return True
    except ValueError:
        return False


def get_model_family_models(family: ModelFamily) -> list[str]:
    """Get all models for a specific family."""
    return MODEL_MAPPINGS.get(family, [])


def get_family_models(family_name: str) -> list[str]:
    """Get all models for a family by name."""
    try:
        family = ModelFamily(family_name)
        return get_model_family_models(family)
    except ValueError:
        return []


def resolve_model_and_provider(model_name: str) -> tuple[str, ModelFamily]:
    """Resolve a short model name to OpenRouter model id and family."""
    normalized = model_name.strip()

    if normalized in SHORT_MODEL_NAMES:
        model_family = SHORT_MODEL_NAMES[normalized]
        if normalized in OPENROUTER_MODEL_MAPPINGS:
            return OPENROUTER_MODEL_MAPPINGS[normalized], model_family

        msg = f"Model '{normalized}' is not available on OpenRouter"
        raise ValueError(msg)

    try:
        return normalized, detect_model_family(normalized)
    except ValueError as e:
        raise ValueError(f"Unknown model: {normalized}") from e


def get_supported_models_by_provider(
    provider: Provider | str,
) -> dict[str, ModelFamily]:
    """Return short model names supported by a provider."""
    if isinstance(provider, str):
        provider = Provider(provider.lower())

    provider_mappings = PROVIDER_MODEL_MAPPINGS.get(provider, {})

    return {short_name: SHORT_MODEL_NAMES[short_name] for short_name in provider_mappings}


def get_all_supported_providers() -> list[Provider]:
    """Get all providers supported by PAGANS runtime."""
    return list(PROVIDER_MODEL_MAPPINGS.keys())


def get_provider_model_name(short_name: str, provider: Provider) -> str:
    """Resolve a short name to provider-specific model id."""
    provider_mappings = PROVIDER_MODEL_MAPPINGS.get(provider)
    if provider_mappings is None:
        msg = f"Provider not supported: {provider.value}"
        raise ValueError(msg)

    if short_name in provider_mappings:
        return provider_mappings[short_name]

    # Pass-through already provider-prefixed names.
    if provider == Provider.OPENROUTER and "/" in short_name:
        return short_name
    if provider != Provider.OPENROUTER:
        provider_prefix = f"{provider.value}/"
        if short_name.startswith(provider_prefix):
            return _strip_provider_prefix(short_name, provider_prefix)
        if "/" not in short_name and short_name:
            # Native providers usually accept plain model ids.
            return short_name

    provider_name = provider.value.capitalize()
    msg = f"Model '{short_name}' is not available on {provider_name}"
    raise ValueError(msg)
