"""Tests for HandoffOrchestrator."""

from __future__ import annotations

from collections.abc import AsyncGenerator
from typing import Any, Optional, Union

import pytest

from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse, Usage
from agentbyte.cancellation_token import CancellationToken
from agentbyte.context import AgentContext
from agentbyte.messages import AssistantMessage, Message, UserMessage
from agentbyte.orchestration.handoff import HandoffOrchestrator
from agentbyte.termination.max_message import MaxMessageTermination
from agentbyte.termination.text_mention import TextMentionTermination


# ---------------------------------------------------------------------------
# Shared mock agent (same pattern used across orchestration tests)
# ---------------------------------------------------------------------------


class MockAgent(BaseAgent):
    def __init__(self, name: str, responses: list[str]):
        self.name = name
        self.description = f"Description of {name}"
        self.responses = responses
        self._resp_idx = 0
        self.tools = []

    async def run(
        self,
        task: Optional[Union[UserMessage, list[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
    ) -> AgentResponse:
        del cancellation_token
        resp_text = self.responses[self._resp_idx % len(self.responses)]
        self._resp_idx += 1
        ctx = context or AgentContext()
        if task and not ctx.messages:
            if isinstance(task, UserMessage):
                ctx.messages.append(task)
            else:
                ctx.messages.extend(task)
        msg = AssistantMessage(content=resp_text, source=self.name)
        ctx.messages.append(msg)
        return AgentResponse(source=self.name, context=ctx, finish_reason="stop")

    async def run_stream(
        self,
        task: Optional[Union[UserMessage, list[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[Union[Message, Any, AgentResponse], None]:
        del cancellation_token, verbose, stream_tokens
        resp_text = self.responses[self._resp_idx % len(self.responses)]
        self._resp_idx += 1

        ctx = context or AgentContext()
        if task:
            if isinstance(task, UserMessage):
                yield task
                if not ctx.messages:
                    ctx.messages.append(task)
            else:
                for message in task:
                    yield message
                    if not ctx.messages:
                        ctx.messages.append(message)

        msg = AssistantMessage(content=resp_text, source=self.name)
        yield msg
        ctx.messages.append(msg)
        yield AgentResponse(
            source=self.name,
            context=ctx,
            usage=Usage(tokens_input=5, tokens_output=5),
            finish_reason="stop",
        )


# ---------------------------------------------------------------------------
# Construction tests
# ---------------------------------------------------------------------------


def test_default_entry_agent():
    """First agent in the list becomes the entry agent by default."""
    a = MockAgent("alpha", ["hi"])
    b = MockAgent("beta", ["hi"])
    orch = HandoffOrchestrator(agents=[a, b], termination=MaxMessageTermination(4))
    assert orch._entry_agent is a


def test_explicit_entry_agent():
    a = MockAgent("alpha", ["hi"])
    b = MockAgent("beta", ["hi"])
    orch = HandoffOrchestrator(
        agents=[a, b],
        termination=MaxMessageTermination(4),
        entry_agent_name="beta",
    )
    assert orch._entry_agent is b


def test_unknown_entry_agent_raises():
    a = MockAgent("alpha", ["hi"])
    with pytest.raises(ValueError, match="entry_agent_name"):
        HandoffOrchestrator(
            agents=[a],
            termination=MaxMessageTermination(4),
            entry_agent_name="nonexistent",
        )


def test_duplicate_agent_names_raises():
    a = MockAgent("alpha", ["hi"])
    b = MockAgent("alpha", ["hi"])
    with pytest.raises(ValueError, match="unique"):
        HandoffOrchestrator(agents=[a, b], termination=MaxMessageTermination(4))


# ---------------------------------------------------------------------------
# _extract_handoff_target
# ---------------------------------------------------------------------------


def _make_orch_with_messages(messages: list[Message]) -> HandoffOrchestrator:
    a = MockAgent("researcher", [""])
    b = MockAgent("writer", [""])
    orch = HandoffOrchestrator(agents=[a, b], termination=MaxMessageTermination(10))
    orch.shared_messages = messages
    return orch


def test_extract_handoff_target_found():
    msg = AssistantMessage(
        content="I've done my research. HANDOFF: writer", source="researcher"
    )
    orch = _make_orch_with_messages([msg])
    assert orch._extract_handoff_target() == "writer"


def test_extract_handoff_target_case_insensitive_prefix():
    msg = AssistantMessage(
        content="Passing over now. handoff: researcher", source="writer"
    )
    orch = _make_orch_with_messages([msg])
    assert orch._extract_handoff_target() == "researcher"


def test_extract_handoff_target_not_found():
    msg = AssistantMessage(content="Just a normal response.", source="researcher")
    orch = _make_orch_with_messages([msg])
    assert orch._extract_handoff_target() is None


def test_extract_handoff_target_hyphenated_name():
    msg = AssistantMessage(
        content="Done. HANDOFF: fact-checker", source="researcher"
    )
    a = MockAgent("researcher", [""])
    b = MockAgent("fact-checker", [""])
    orch = HandoffOrchestrator(agents=[a, b], termination=MaxMessageTermination(10))
    orch.shared_messages = [msg]
    assert orch._extract_handoff_target() == "fact-checker"


def test_extract_handoff_uses_last_assistant_message():
    """Only the most recent assistant message is inspected."""
    old = AssistantMessage(content="HANDOFF: writer", source="researcher")
    new = AssistantMessage(content="All done, no handoff.", source="writer")
    orch = _make_orch_with_messages([old, new])
    assert orch._extract_handoff_target() is None


# ---------------------------------------------------------------------------
# _build_colleagues_block
# ---------------------------------------------------------------------------


def test_build_colleagues_block_excludes_self():
    a = MockAgent("researcher", [""])
    b = MockAgent("writer", [""])
    orch = HandoffOrchestrator(agents=[a, b], termination=MaxMessageTermination(10))
    block = orch._build_colleagues_block(a)
    assert "writer" in block
    assert "researcher" not in block


def test_build_colleagues_block_single_agent_empty():
    a = MockAgent("solo", [""])
    orch = HandoffOrchestrator(agents=[a], termination=MaxMessageTermination(10))
    assert orch._build_colleagues_block(a) == ""


# ---------------------------------------------------------------------------
# Full orchestration runs
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_entry_agent_runs_first():
    """The entry agent always executes on the very first iteration."""
    a = MockAgent("alpha", ["HANDOFF: beta"])
    b = MockAgent("beta", ["done APPROVED"])
    orch = HandoffOrchestrator(
        agents=[a, b],
        termination=TextMentionTermination("APPROVED"),
        entry_agent_name="alpha",
    )
    result = await orch.run("start")
    sources = [m.source for m in result.messages if hasattr(m, "role")]
    assert sources[0] == "user"
    # alpha should appear before beta
    alpha_idx = next(i for i, s in enumerate(sources) if s == "alpha")
    beta_idx = next(i for i, s in enumerate(sources) if s == "beta")
    assert alpha_idx < beta_idx


@pytest.mark.asyncio
async def test_handoff_signal_routes_to_correct_agent():
    """When alpha signals HANDOFF: beta, beta should execute next."""
    alpha = MockAgent("alpha", ["research done. HANDOFF: beta"])
    beta = MockAgent("beta", ["article written APPROVED"])
    orch = HandoffOrchestrator(
        agents=[alpha, beta],
        termination=TextMentionTermination("APPROVED"),
    )
    result = await orch.run("write about AI")
    sources = [m.source for m in result.messages if getattr(m, "role", "") == "assistant"]
    assert sources == ["alpha", "beta"]


@pytest.mark.asyncio
async def test_fallback_round_robin_when_no_signal():
    """Without a handoff signal the fallback cycles through agents."""
    alpha = MockAgent("alpha", ["no signal"])
    beta = MockAgent("beta", ["no signal"])
    orch = HandoffOrchestrator(
        agents=[alpha, beta],
        termination=MaxMessageTermination(5),
        fallback_strategy="round_robin",
    )
    result = await orch.run("go")
    assistant_sources = [
        m.source for m in result.messages if getattr(m, "role", "") == "assistant"
    ]
    # With round-robin fallback and 5-message cap we expect alternating agents
    assert len(assistant_sources) >= 2
    assert "alpha" in assistant_sources
    assert "beta" in assistant_sources


@pytest.mark.asyncio
async def test_fallback_entry_strategy():
    """With fallback_strategy='entry', no-signal responses return to entry agent."""
    alpha = MockAgent("alpha", ["no signal"])
    beta = MockAgent("beta", ["no signal"])
    orch = HandoffOrchestrator(
        agents=[alpha, beta],
        termination=MaxMessageTermination(5),
        entry_agent_name="alpha",
        fallback_strategy="entry",
    )
    result = await orch.run("go")
    assistant_sources = [
        m.source for m in result.messages if getattr(m, "role", "") == "assistant"
    ]
    # alpha is always chosen on fallback
    assert all(s == "alpha" for s in assistant_sources)


@pytest.mark.asyncio
async def test_unknown_handoff_target_falls_back():
    """An unresolvable handoff target should not crash — fallback applies."""
    alpha = MockAgent("alpha", ["HANDOFF: ghost_agent"])
    beta = MockAgent("beta", ["done APPROVED"])
    orch = HandoffOrchestrator(
        agents=[alpha, beta],
        termination=TextMentionTermination("APPROVED"),
        fallback_strategy="round_robin",
    )
    result = await orch.run("start")
    # Should complete without error
    assert result.stop_message is not None


@pytest.mark.asyncio
async def test_handoff_counts_tracked():
    """Metadata should record how many times each agent was handed off to."""
    alpha = MockAgent("alpha", ["HANDOFF: beta"])
    beta = MockAgent("beta", ["done APPROVED"])
    orch = HandoffOrchestrator(
        agents=[alpha, beta],
        termination=TextMentionTermination("APPROVED"),
    )
    result = await orch.run("start")
    assert result.pattern_metadata["handoff_counts"].get("beta", 0) >= 1


@pytest.mark.asyncio
async def test_colleagues_injected_in_context():
    """When inject_colleagues_context=True the context contains the colleagues block."""
    received_contexts: list[str] = []

    class CapturingAgent(BaseAgent):
        def __init__(self, name: str, response: str):
            self.name = name
            self.description = f"Desc {name}"
            self.tools = []
            self._response = response

        async def run(self, task=None, context=None, cancellation_token=None):
            if isinstance(task, UserMessage):
                received_contexts.append(task.content)
            ctx = AgentContext()
            msg = AssistantMessage(content=self._response, source=self.name)
            ctx.messages.append(msg)
            return AgentResponse(source=self.name, context=ctx, finish_reason="stop")

        async def run_stream(
            self, task=None, context=None, cancellation_token=None,
            verbose=False, stream_tokens=False,
        ):
            if isinstance(task, UserMessage):
                received_contexts.append(task.content)
            ctx = AgentContext()
            if task:
                if isinstance(task, UserMessage):
                    yield task
                    ctx.messages.append(task)
                else:
                    for m in task:
                        yield m
                        ctx.messages.append(m)
            msg = AssistantMessage(content=self._response, source=self.name)
            yield msg
            ctx.messages.append(msg)
            yield AgentResponse(
                source=self.name, context=ctx,
                usage=Usage(tokens_input=1, tokens_output=1), finish_reason="stop"
            )

    alpha = CapturingAgent("alpha", "done APPROVED")
    beta = CapturingAgent("beta", "ok")
    orch = HandoffOrchestrator(
        agents=[alpha, beta],
        termination=TextMentionTermination("APPROVED"),
        inject_colleagues_context=True,
    )
    await orch.run("start")
    assert any("beta" in ctx for ctx in received_contexts), (
        "Expected colleagues block to mention 'beta' in alpha's context"
    )


@pytest.mark.asyncio
async def test_no_colleagues_injection_when_disabled():
    """When inject_colleagues_context=False agents receive plain history context."""
    received_contexts: list[str] = []

    class CapturingAgent(BaseAgent):
        def __init__(self, name: str, response: str):
            self.name = name
            self.description = f"Desc {name}"
            self.tools = []
            self._response = response

        async def run(self, task=None, context=None, cancellation_token=None):
            ctx = AgentContext()
            msg = AssistantMessage(content=self._response, source=self.name)
            ctx.messages.append(msg)
            return AgentResponse(source=self.name, context=ctx, finish_reason="stop")

        async def run_stream(
            self, task=None, context=None, cancellation_token=None,
            verbose=False, stream_tokens=False,
        ):
            if isinstance(task, UserMessage):
                received_contexts.append(task.content)
            ctx = AgentContext()
            if task:
                if isinstance(task, UserMessage):
                    yield task
                    ctx.messages.append(task)
                else:
                    for m in task:
                        yield m
                        ctx.messages.append(m)
            msg = AssistantMessage(content=self._response, source=self.name)
            yield msg
            ctx.messages.append(msg)
            yield AgentResponse(
                source=self.name, context=ctx,
                usage=Usage(tokens_input=1, tokens_output=1), finish_reason="stop"
            )

    alpha = CapturingAgent("alpha", "done APPROVED")
    beta = CapturingAgent("beta", "ok")
    orch = HandoffOrchestrator(
        agents=[alpha, beta],
        termination=TextMentionTermination("APPROVED"),
        inject_colleagues_context=False,
    )
    await orch.run("start")
    # Colleagues block should NOT appear
    assert all("Colleagues" not in ctx for ctx in received_contexts)


@pytest.mark.asyncio
async def test_pattern_metadata_fields():
    """pattern_metadata should include handoff-specific keys."""
    alpha = MockAgent("alpha", ["done APPROVED"])
    orch = HandoffOrchestrator(
        agents=[alpha],
        termination=TextMentionTermination("APPROVED"),
    )
    result = await orch.run("go")
    meta = result.pattern_metadata
    assert "entry_agent" in meta
    assert "fallback_strategy" in meta
    assert "handoff_signal_prefix" in meta
    assert "handoff_counts" in meta


@pytest.mark.asyncio
async def test_max_iterations_respected():
    """Orchestration must stop at max_iterations even without a termination signal."""
    alpha = MockAgent("alpha", ["no signal"])
    orch = HandoffOrchestrator(
        agents=[alpha],
        termination=MaxMessageTermination(100),
        max_iterations=3,
    )
    result = await orch.run("go")
    assert result.stop_message is not None
    assert "Maximum iterations" in result.stop_message.content

