"""API provider registry — register and retrieve LLM provider stream functions."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable


@dataclass
class ApiProviderInternal:
    """Internal representation of a registered API provider."""

    api: str
    stream: Callable[..., Any]
    stream_simple: Callable[..., Any]


@dataclass
class _RegisteredProvider:
    provider: ApiProviderInternal
    source_id: str | None = None


_registry: dict[str, _RegisteredProvider] = {}


def register_api_provider(
    api: str,
    stream: Callable[..., Any],
    stream_simple: Callable[..., Any],
    source_id: str | None = None,
) -> None:
    """Register an API provider with stream functions."""
    _registry[api] = _RegisteredProvider(
        provider=ApiProviderInternal(api=api, stream=stream, stream_simple=stream_simple),
        source_id=source_id,
    )


def get_api_provider(api: str) -> ApiProviderInternal | None:
    """Get a registered API provider by API type."""
    entry = _registry.get(api)
    return entry.provider if entry else None


def get_api_providers() -> list[ApiProviderInternal]:
    """Get all registered API providers."""
    return [entry.provider for entry in _registry.values()]


def unregister_api_providers(source_id: str) -> None:
    """Unregister all API providers with the given source ID."""
    to_remove = [api for api, entry in _registry.items() if entry.source_id == source_id]
    for api in to_remove:
        del _registry[api]


def clear_api_providers() -> None:
    """Clear all registered API providers."""
    _registry.clear()


# Auto-register built-ins so GET /ai/api-providers matches TypeScript default behavior.
# This import has the side-effect of registering stream adapters for known APIs.
#
# Note: import must be at the bottom to avoid circular import issues.
from pi_agents.ai.providers.register_builtins import register_builtin_providers as _register_builtin_providers

_register_builtin_providers()
