"""Registry resolver for the inference gateway.

Resolves agent stage bindings to complete InferenceTarget instances
using TTL-cached lookups against the registry tables. Auth secrets
are NOT resolved during caching — only at invocation time.

Requirements: 3.5, 3.9
"""
from __future__ import annotations

import logging
import time
from typing import Any, TypeVar
from uuid import UUID

from services.shared.inference.errors import InferenceError, InferenceErrorCategory
from services.shared.inference.models import InferenceTarget, ProviderCapabilities

logger = logging.getLogger(__name__)

T = TypeVar("T")

# Default cache TTL in seconds
DEFAULT_CACHE_TTL_SECONDS = 60.0


# ---------------------------------------------------------------------------
# RegistryCache — internal TTL dict
# ---------------------------------------------------------------------------


class RegistryCache:
    """TTL-based in-memory cache for registry lookups.

    Keys follow the format:
      - "binding:{agent_id}:{stage}"
      - "endpoint:{endpoint_id}"
      - "deployment:{deployment_id}"

    Expired entries are lazily evicted on access.
    """

    def __init__(self, ttl_seconds: float = DEFAULT_CACHE_TTL_SECONDS) -> None:
        self._ttl_seconds = ttl_seconds
        self._store: dict[str, tuple[float, Any]] = {}

    @property
    def ttl_seconds(self) -> float:
        return self._ttl_seconds

    def get(self, key: str) -> Any | None:
        """Retrieve a cached value if it exists and hasn't expired."""
        entry = self._store.get(key)
        if entry is None:
            return None

        stored_at, value = entry
        if (time.monotonic() - stored_at) > self._ttl_seconds:
            del self._store[key]
            return None

        return value

    def set(self, key: str, value: Any) -> None:
        """Store a value with the current timestamp."""
        self._store[key] = (time.monotonic(), value)

    def invalidate(self, key_pattern: str) -> None:
        """Invalidate all entries whose key starts with the given pattern.

        Supports prefix-based invalidation:
          - invalidate("endpoint:abc-123") removes that specific endpoint
          - invalidate("endpoint:") removes ALL endpoint entries
          - invalidate("binding:agent-1:") removes all bindings for agent-1
        """
        keys_to_remove = [k for k in self._store if k.startswith(key_pattern)]
        for k in keys_to_remove:
            del self._store[k]

    def clear(self) -> None:
        """Clear all cached entries."""
        self._store.clear()

    def __len__(self) -> int:
        return len(self._store)

    def __contains__(self, key: str) -> bool:
        """Check if a non-expired entry exists for the key."""
        return self.get(key) is not None


# ---------------------------------------------------------------------------
# Database query protocol
# ---------------------------------------------------------------------------


class RegistryDB:
    """Protocol for registry database queries.

    In production this would be backed by asyncpg. For testing,
    a simple dict-based mock implements the same interface.
    """

    async def get_active_binding(
        self, agent_id: UUID, stage: str
    ) -> dict[str, Any] | None:
        """Get the active binding for an agent+stage.

        Returns a dict with keys: id, agent_id, stage, model_deployment_id,
        route_order, routing_config, is_active, revision.
        Returns None if no active binding exists.
        """
        raise NotImplementedError

    async def get_model_deployment(
        self, deployment_id: UUID
    ) -> dict[str, Any] | None:
        """Get a model deployment by ID.

        Returns a dict with keys: id, endpoint_id, served_model_name,
        display_name, capabilities, context_window, max_output_tokens,
        quantization, runtime_metadata, enabled, revision.
        """
        raise NotImplementedError

    async def get_inference_endpoint(
        self, endpoint_id: UUID
    ) -> dict[str, Any] | None:
        """Get an inference endpoint by ID.

        Returns a dict with keys: id, name, protocol, base_url,
        auth_secret_ref, auth_scheme, default_headers, health_path,
        enabled, revision.
        """
        raise NotImplementedError


# ---------------------------------------------------------------------------
# RegistryResolver
# ---------------------------------------------------------------------------


class RegistryResolver:
    """Resolves agent stage bindings to InferenceTarget instances.

    Resolution order:
      1. Find active binding for (agent_id, stage)
      2. Get model_deployment from binding
      3. Get inference_endpoint from deployment
      4. Build InferenceTarget

    Fail-closed: any missing or disabled resource raises InferenceError
    with CAPABILITY_UNAVAILABLE. Never returns a fallback target.

    Auth secrets are NOT resolved during caching — auth_secret_ref is
    preserved as-is in the target. Resolution happens at invocation time
    by the client layer.
    """

    def __init__(
        self,
        db: RegistryDB,
        *,
        cache_ttl_seconds: float = DEFAULT_CACHE_TTL_SECONDS,
    ) -> None:
        self._db = db
        self._cache = RegistryCache(ttl_seconds=cache_ttl_seconds)

    @property
    def cache(self) -> RegistryCache:
        """Access the internal cache (for testing/inspection)."""
        return self._cache

    async def resolve_target(self, agent_id: UUID, stage: str) -> InferenceTarget:
        """Resolve the active binding for an agent+stage into an InferenceTarget.

        Uses TTL-cached lookups. If any component is missing or disabled,
        raises InferenceError(CAPABILITY_UNAVAILABLE).

        Auth secret is NOT resolved here — only at invocation time.
        """
        # Check cache first for the full resolved target
        cache_key = f"binding:{agent_id}:{stage}"
        cached_target = self._cache.get(cache_key)
        if cached_target is not None:
            return cached_target

        # Step 1: Find active binding
        binding = await self._db.get_active_binding(agent_id, stage)
        if binding is None:
            raise InferenceError(
                InferenceErrorCategory.CAPABILITY_UNAVAILABLE,
                f"No active binding for agent={agent_id} stage={stage}",
            )

        if not binding.get("is_active", False):
            raise InferenceError(
                InferenceErrorCategory.CAPABILITY_UNAVAILABLE,
                f"Binding for agent={agent_id} stage={stage} is inactive",
            )

        deployment_id = binding.get("model_deployment_id")
        if deployment_id is None:
            raise InferenceError(
                InferenceErrorCategory.CAPABILITY_UNAVAILABLE,
                f"Binding for agent={agent_id} stage={stage} has no deployment",
            )

        # Step 2: Get model deployment
        deployment = await self._resolve_deployment(deployment_id)

        # Step 3: Get inference endpoint
        endpoint_id = deployment["endpoint_id"]
        endpoint = await self._resolve_endpoint(endpoint_id)

        # Step 4: Build InferenceTarget
        capabilities_data = deployment.get("capabilities", {})
        capabilities = ProviderCapabilities(
            chat_completions=capabilities_data.get("chat_completions", False),
            responses_api=capabilities_data.get("responses_api", False),
            json_schema=capabilities_data.get("json_schema", False),
            json_object=capabilities_data.get("json_object", False),
            seed=capabilities_data.get("seed", False),
            usage=capabilities_data.get("usage", False),
            max_completion_tokens=capabilities_data.get("max_completion_tokens", False),
            reasoning_toggle=capabilities_data.get("reasoning_toggle", False),
            model_listing=capabilities_data.get("model_listing", False),
        )

        extra_headers = endpoint.get("default_headers", {})
        if not isinstance(extra_headers, dict):
            extra_headers = {}

        runtime_metadata = deployment.get("runtime_metadata", {})
        extra_body = runtime_metadata.get("extra_body", {}) if isinstance(runtime_metadata, dict) else {}

        target = InferenceTarget(
            endpoint_id=endpoint_id,
            deployment_id=deployment_id,
            protocol=endpoint["protocol"],
            base_url=endpoint["base_url"],
            model=deployment["served_model_name"],
            capabilities=capabilities,
            auth_secret_ref=endpoint.get("auth_secret_ref"),
            auth_scheme=endpoint.get("auth_scheme", "bearer"),
            extra_headers=extra_headers,
            extra_body=extra_body,
            context_window=deployment.get("context_window") or 0,
            max_output_tokens=deployment.get("max_output_tokens"),
        )

        # Cache the resolved target
        self._cache.set(cache_key, target)

        return target

    async def _resolve_deployment(self, deployment_id: UUID) -> dict[str, Any]:
        """Resolve a model deployment, using cache if available."""
        dep_cache_key = f"deployment:{deployment_id}"
        cached = self._cache.get(dep_cache_key)
        if cached is not None:
            return cached

        deployment = await self._db.get_model_deployment(deployment_id)
        if deployment is None:
            raise InferenceError(
                InferenceErrorCategory.CAPABILITY_UNAVAILABLE,
                f"Model deployment {deployment_id} not found",
            )

        if not deployment.get("enabled", False):
            raise InferenceError(
                InferenceErrorCategory.CAPABILITY_UNAVAILABLE,
                f"Model deployment {deployment_id} is disabled",
            )

        self._cache.set(dep_cache_key, deployment)
        return deployment

    async def _resolve_endpoint(self, endpoint_id: UUID) -> dict[str, Any]:
        """Resolve an inference endpoint, using cache if available."""
        ep_cache_key = f"endpoint:{endpoint_id}"
        cached = self._cache.get(ep_cache_key)
        if cached is not None:
            return cached

        endpoint = await self._db.get_inference_endpoint(endpoint_id)
        if endpoint is None:
            raise InferenceError(
                InferenceErrorCategory.CAPABILITY_UNAVAILABLE,
                f"Inference endpoint {endpoint_id} not found",
            )

        if not endpoint.get("enabled", False):
            raise InferenceError(
                InferenceErrorCategory.CAPABILITY_UNAVAILABLE,
                f"Inference endpoint {endpoint_id} is disabled",
            )

        self._cache.set(ep_cache_key, endpoint)
        return endpoint

    def invalidate(self, endpoint_id: UUID) -> None:
        """Evict cache entries related to an endpoint.

        Called on revision changes or probe failures to force
        re-resolution on the next request.
        """
        self._cache.invalidate(f"endpoint:{endpoint_id}")
        # Also clear all binding caches since they may reference this endpoint
        # We clear all bindings because we can't efficiently know which bindings
        # use this endpoint without scanning
        self._cache.invalidate("binding:")
        logger.info("Invalidated cache for endpoint %s", endpoint_id)

    def invalidate_all(self) -> None:
        """Full cache clear."""
        self._cache.clear()
        logger.info("Full registry cache invalidated")
