"""Shared test helpers for WebUI coverage."""

from __future__ import annotations

from collections.abc import AsyncGenerator

from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse, ToolApprovalEvent
from agentbyte.context import AgentContext
from agentbyte.messages import AssistantMessage, ToolCallRequest, Usage, UserMessage
from agentbyte.orchestration.base import BaseOrchestrator
from agentbyte.termination import MaxMessageTermination
from agentbyte.types import StopMessage


class DummyModelClient:
    """Tiny stand-in model client."""

    def __init__(self, model: str = "dummy-model") -> None:
        self.model = model


class DummyAgent(BaseAgent):
    """Simple agent used for WebUI tests."""

    def __init__(self, require_approval: bool = False) -> None:
        super().__init__(
            name="dummy-agent",
            description="Dummy agent",
            instructions="Be helpful.",
            model_client=DummyModelClient(),
        )
        self.require_approval = require_approval

    async def run(
        self,
        task=None,
        context: AgentContext | None = None,
        cancellation_token=None,
    ) -> AgentResponse:
        final_response: AgentResponse | None = None
        async for item in self.run_stream(
            task=task,
            context=context,
            cancellation_token=cancellation_token,
            verbose=True,
        ):
            if isinstance(item, AgentResponse):
                final_response = item
        assert final_response is not None
        return final_response

    async def run_stream(
        self,
        task=None,
        context: AgentContext | None = None,
        cancellation_token=None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[object, None]:
        working_context = context or AgentContext()

        if task is not None:
            for message in self._convert_task_to_messages(task):
                working_context.add_message(message)

        if self.require_approval and working_context.get_approval_response("call_1") is None:
            request = working_context.add_approval_request(
                ToolCallRequest(
                    tool_name="dangerous_tool",
                    parameters={"value": 5},
                    call_id="call_1",
                ),
                "dangerous_tool",
            )
            yield ToolApprovalEvent(source=self.name, approval_request=request)
            yield AgentResponse(
                context=working_context,
                source=self.name,
                finish_reason="approval_needed",
                usage=Usage(llm_calls=1, tokens_input=10, tokens_output=2),
            )
            return

        if self.require_approval:
            working_context.consume_approval_response("call_1")

        last_user = next(
            (
                message
                for message in reversed(working_context.messages)
                if isinstance(message, UserMessage)
            ),
            None,
        )
        reply = AssistantMessage(
            content=f"echo: {last_user.content if last_user else 'done'}",
            source=self.name,
            usage=Usage(tokens_input=8, tokens_output=4, llm_calls=1),
        )
        working_context.add_message(reply)
        yield reply
        yield AgentResponse(
            context=working_context,
            source=self.name,
            finish_reason="stop",
            usage=Usage(llm_calls=1, tokens_input=8, tokens_output=4),
        )


class DummyOrchestrator(BaseOrchestrator):
    """Minimal orchestrator for execution tests."""

    def __init__(self, agent: BaseAgent) -> None:
        super().__init__(
            agents=[agent],
            termination=MaxMessageTermination(max_messages=10),
            max_iterations=3,
            name="dummy-orchestrator",
            description="Dummy orchestrator",
        )
        self._should_stop = False

    async def select_next_agent(self) -> BaseAgent:
        return self.agents[0]

    async def prepare_context_for_agent(self, agent: BaseAgent):
        return [UserMessage(content="orchestrated task", source="orchestrator")]

    async def update_shared_state(self, result: AgentResponse) -> None:
        for message in result.messages:
            if message not in self.shared_messages:
                self.shared_messages.append(message)
        self._should_stop = True

    def _get_pattern_stop_message(self) -> StopMessage | None:
        if self._should_stop:
            return StopMessage(content="Pattern complete", source="DummyOrchestrator")
        return None


__all__ = ["DummyAgent", "DummyOrchestrator", "DummyModelClient"]
