"""Antigravity model registry for managing model configurations and aliases."""

from __future__ import annotations

from ..shared import ModelCapabilities, ProviderType
from .base import CAPABILITY_FIELD_NAMES, CapabilityModelRegistry


class AntigravityModelRegistry(CapabilityModelRegistry):
    """Capability registry backed by ``conf/antigravity_models.json``.

    The Antigravity provider accesses Claude, Gemini, and other models through
    Google's unified gateway API. Models are split across two quota pools:

    * **antigravity**: Claude models and Gemini 3 variants
    * **gemini-cli**: Gemini 2.5 and preview models using Gemini CLI quota
    """

    # Extra fields beyond standard ModelCapabilities
    EXTRA_KEYS = {"quota_pool"}

    def __init__(self, config_path: str | None = None) -> None:
        super().__init__(
            env_var_name="ANTIGRAVITY_MODELS_CONFIG_PATH",
            default_filename="antigravity_models.json",
            provider=ProviderType.ANTIGRAVITY,
            friendly_prefix="Antigravity ({model})",
            config_path=config_path,
        )

    def _extra_keys(self) -> set[str]:
        """Include quota_pool as a valid entry key."""
        return self.EXTRA_KEYS

    def _finalise_entry(self, entry: dict) -> tuple[ModelCapabilities, dict]:
        """Build ModelCapabilities and extract extras like quota_pool."""
        # Extract extras before building capabilities
        extras = {}
        quota_pool = entry.pop("quota_pool", None)
        if quota_pool:
            extras["quota_pool"] = quota_pool

        # Build capabilities from remaining fields
        filtered = {k: v for k, v in entry.items() if k in CAPABILITY_FIELD_NAMES}
        filtered.setdefault("provider", ProviderType.ANTIGRAVITY)
        filtered.setdefault("friendly_name", f"Antigravity ({entry.get('model_name', 'unknown')})")
        capability = ModelCapabilities(**filtered)

        return capability, extras

    def get_quota_pool(self, model_name: str) -> str:
        """Return the quota pool for a model ('antigravity' or 'gemini-cli')."""
        extras = self.get_entry(model_name)
        if extras and "quota_pool" in extras:
            return extras["quota_pool"]
        return "antigravity"  # Default to antigravity quota pool