"""Stream routing — routes requests to the correct provider implementation."""

from __future__ import annotations

from typing import Any

from pi_agents.ai.api_registry import get_api_provider
from pi_agents.ai.types import AssistantMessage, Context, Model, SimpleStreamOptions, StreamOptions
from pi_agents.ai.utils.event_stream import AssistantMessageEventStream

# Ensure built-in providers are registered
import pi_agents.ai.providers.register_builtins  # noqa: F401


def stream(
    model: Model,
    context: Context,
    options: dict[str, Any] | None = None,
) -> AssistantMessageEventStream:
    """Stream responses from an AI model using full provider-specific options.

    Routes to the correct provider based on the model's API type.
    """
    provider = get_api_provider(model.api)
    if not provider:
        raise ValueError(f"No provider registered for API: {model.api}")
    return provider.stream(model, context, options)


async def complete(
    model: Model,
    context: Context,
    options: dict[str, Any] | None = None,
) -> AssistantMessage:
    """Get a complete response from an AI model.

    Internally calls stream() and awaits the result.
    """
    event_stream = stream(model, context, options)
    return await event_stream.result()


def stream_simple(
    model: Model,
    context: Context,
    options: SimpleStreamOptions | None = None,
) -> AssistantMessageEventStream:
    """Stream responses using simplified, provider-agnostic options."""
    provider = get_api_provider(model.api)
    if not provider:
        raise ValueError(f"No provider registered for API: {model.api}")
    return provider.stream_simple(model, context, options)


async def complete_simple(
    model: Model,
    context: Context,
    options: SimpleStreamOptions | None = None,
) -> AssistantMessage:
    """Get a complete response using simplified options."""
    event_stream = stream_simple(model, context, options)
    return await event_stream.result()
