from typing import Any, Dict, List, Optional
from unittest.mock import patch

import pytest

from swarm import Agent, DelimiterChunk, Swarm
from swarm.context import RunContext
from swarm.runners.agent import AgentRunResult, AgentSignal


def _make_config() -> Dict[str, str]:
    return {
        "openai_api_key": "openai-key",
        "openai_organization": "org-123",
        "openai_project": "proj-123",
        "anthropic_api_key": "anthropic-key",
        "google_api_key": "google-key",
        "vertex_project": "vertex-project",
        "vertex_location": "europe-west1",
    }


class _FakeContext:
    def __init__(self) -> None:
        self._conversation_id = "conv-1"
        self._context_variables: Dict[str, Any] = {}
        self.workflow_context: Dict[str, Any] = {}
        self._designer_fullname = "Swarm Test"

    @property
    def conversation_id(self) -> str:
        return self._conversation_id

    @property
    def current_context_variables(self) -> Dict[str, Any]:
        return self._context_variables

    def get_history_for_completion(self) -> List[Any]:
        return []

    @property
    def designer_fullname(self) -> str:
        return self._designer_fullname


class _FakeStreamManager:
    async def publish_status(self, status: str) -> None:
        _ = status

    async def publish_error(self, message: str, code: int = 500) -> None:
        _ = (message, code)

    async def publish_message(self, content: Any, **kwargs: Any) -> None:
        _ = (content, kwargs)


class _FakeMessageStore:
    async def persist_messages(self, **kwargs: Any) -> None:
        _ = kwargs


class _FakeRetryPolicy:
    @property
    def max_retries(self) -> int:
        return 3

    @property
    def backoff(self) -> str:
        return "exponential"

    @property
    def base_seconds(self) -> int:
        return 5

    @property
    def jitter(self) -> bool:
        return True

    @property
    def retry_on(self) -> List[str]:
        return ["timeout"]


class _RecordingDialogueRunner:
    def __init__(self) -> None:
        self.calls: List[Dict[str, Any]] = []

    def run_streamed(self, **kwargs: Any):
        self.calls.append(kwargs)

        async def _chunks():
            yield DelimiterChunk(delim="complete")

        return _chunks()


class _RecordingAgentRunner:
    def __init__(self, result: Optional[AgentRunResult] = None) -> None:
        self.calls: List[Dict[str, Any]] = []
        self.result = result or AgentRunResult(
            signal=AgentSignal.COMPLETED,
            turns_executed=1,
        )

    async def run_batch(self, **kwargs: Any) -> AgentRunResult:
        self.calls.append(kwargs)
        return self.result


class TestSwarm:
    def test_init_loads_config_from_env_by_default(self) -> None:
        with patch(
            "swarm.client.ConfigLoader.load", return_value=_make_config()
        ) as load:
            client = Swarm()

        load.assert_called_once_with()
        assert client._router.factory.config == _make_config()

    def test_init_copies_explicit_config(self) -> None:
        config = _make_config()

        client = Swarm(config=config)
        config["openai_api_key"] = "changed"

        assert client._router.factory.config["openai_api_key"] == "openai-key"

    @pytest.mark.asyncio
    async def test_run_streamed_builds_run_context_and_delegates(self) -> None:
        agent = Agent(
            name="Assistant",
            model="openai/gpt-4o",
            instructions="You are helpful.",
        )
        client = Swarm(config=_make_config())
        dialogue_runner = _RecordingDialogueRunner()
        client._dialogue_runner = dialogue_runner
        deps = _FakeContext()
        stream_manager = _FakeStreamManager()
        message_store = _FakeMessageStore()
        retry_policy = _FakeRetryPolicy()

        chunks = [
            chunk
            async for chunk in client.run_streamed(
                agent,
                deps=deps,
                stream_manager=stream_manager,
                message_store=message_store,
                retry_policy=retry_policy,
                max_turns=7,
                is_new_conversation=True,
            )
        ]

        assert len(chunks) == 1
        call = dialogue_runner.calls[0]
        run_context = call["run_context"]
        assert isinstance(run_context, RunContext)
        assert run_context.deps is deps
        assert run_context.stream_manager is stream_manager
        assert run_context.message_store is message_store
        assert run_context.retry_policy is retry_policy
        assert run_context.router is client._router
        assert call["agent"] is agent
        assert call["max_turns"] == 7
        assert call["is_new_conversation"] is True
        assert isinstance(chunks[0], DelimiterChunk)
        assert chunks[0].delim == "complete"

    @pytest.mark.asyncio
    async def test_run_batch_builds_run_context_and_delegates(self) -> None:
        agent = Agent(
            name="Assistant",
            model="openai/gpt-4o",
            instructions="You are helpful.",
        )
        expected_result = AgentRunResult(
            signal=AgentSignal.COMPLETED,
            turns_executed=2,
        )
        client = Swarm(config=_make_config())
        agent_runner = _RecordingAgentRunner(result=expected_result)
        client._agent_runner = agent_runner
        deps = _FakeContext()
        stream_manager = _FakeStreamManager()
        message_store = _FakeMessageStore()
        retry_policy = _FakeRetryPolicy()

        result = await client.run_batch(
            agent,
            deps=deps,
            stream_manager=stream_manager,
            message_store=message_store,
            retry_policy=retry_policy,
            max_turns=9,
        )

        call = agent_runner.calls[0]
        run_context = call["run_context"]
        assert result == expected_result
        assert isinstance(run_context, RunContext)
        assert run_context.deps is deps
        assert run_context.stream_manager is stream_manager
        assert run_context.message_store is message_store
        assert run_context.retry_policy is retry_policy
        assert run_context.router is client._router
        assert call["agent"] is agent
        assert call["max_turns"] == 9
