"""Beginner-facing facade for common Swarm usage."""

from collections.abc import Mapping
from typing import AsyncIterator, Awaitable, Callable, TypeVar, cast

from swarm.agent import Agent
from swarm.chunk import ChunkType
from swarm.config.loader import ConfigLoader
from swarm.config.protocols import (
    MessageStoreProtocol,
    RetryPolicyProtocol,
    StreamManagerProtocol,
)
from swarm.context import ContextProtocol, RunContext
from swarm.providers.factory import ProviderFactory
from swarm.providers.router import ModelRouter
from swarm.runners.agent import AgentRunner, AgentRunResult
from swarm.runners.dialogue import DialogueRunner

DepsT = TypeVar("DepsT", bound=ContextProtocol)


class Swarm:
    """Facade that composes provider configuration for common package consumers."""

    def __init__(self, config: Mapping[str, str] | None = None) -> None:
        resolved_config = ConfigLoader.load() if config is None else dict(config)
        self._router = ModelRouter(factory=ProviderFactory(config=resolved_config))
        self._dialogue_runner = DialogueRunner()
        self._agent_runner = AgentRunner()

    def _build_run_context(
        self,
        *,
        deps: DepsT,
        stream_manager: StreamManagerProtocol,
        message_store: MessageStoreProtocol,
        retry_policy: RetryPolicyProtocol | None = None,
    ) -> RunContext[DepsT]:
        return RunContext(
            deps=deps,
            stream_manager=stream_manager,
            message_store=message_store,
            router=self._router,
            retry_policy=retry_policy,
        )

    def run_streamed(
        self,
        agent: Agent,
        *,
        deps: DepsT,
        stream_manager: StreamManagerProtocol,
        message_store: MessageStoreProtocol,
        retry_policy: RetryPolicyProtocol | None = None,
        max_turns: int = 25,
        is_new_conversation: bool = False,
        check_stop_fn: Callable[[], Awaitable[bool]] | None = None,
        heartbeat_fn: Callable[[], Awaitable[None]] | None = None,
    ) -> AsyncIterator[ChunkType]:
        run_context = self._build_run_context(
            deps=deps,
            stream_manager=stream_manager,
            message_store=message_store,
            retry_policy=retry_policy,
        )
        return self._dialogue_runner.run_streamed(
            agent=agent,
            run_context=cast(RunContext[ContextProtocol], run_context),
            max_turns=max_turns,
            is_new_conversation=is_new_conversation,
            check_stop_fn=check_stop_fn,
            heartbeat_fn=heartbeat_fn,
        )

    async def run_batch(
        self,
        agent: Agent,
        *,
        deps: DepsT,
        stream_manager: StreamManagerProtocol,
        message_store: MessageStoreProtocol,
        retry_policy: RetryPolicyProtocol | None = None,
        max_turns: int,
    ) -> AgentRunResult:
        run_context = self._build_run_context(
            deps=deps,
            stream_manager=stream_manager,
            message_store=message_store,
            retry_policy=retry_policy,
        )
        return await self._agent_runner.run_batch(
            agent=agent,
            run_context=cast(RunContext[ContextProtocol], run_context),
            max_turns=max_turns,
        )
