"""Handoff orchestration pattern.

Agents drive routing by embedding a handoff signal in their response.
The orchestrator reads the signal, resolves the target agent, and
passes control there.  When no signal is present a configurable
fallback strategy (round-robin or fixed entry agent) keeps execution
moving forward.
"""

from __future__ import annotations

from collections.abc import Sequence
import re
from typing import Any, Literal

from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse
from agentbyte.messages import Message, UserMessage
from agentbyte.termination import BaseTermination
from .base import BaseOrchestrator
from .policies import AgentHistoryPolicy, HistoryContextPolicy


class HandoffOrchestrator(BaseOrchestrator):
    """
    Handoff orchestration pattern.

    Instead of the orchestrator always choosing who speaks next, each
    agent decides who should handle the task from here by embedding a
    handoff signal at the end of its response::

        HANDOFF: researcher

    The orchestrator parses this signal, resolves the named agent, and
    passes control to them.  When an agent produces no handoff signal
    the fallback strategy is applied.

    Fallback strategies:

    * ``"round_robin"`` — cycle through agents in order (default)
    * ``"entry"``       — always route back to the entry agent

    Context injection:

    Each agent automatically receives a "colleagues" block listing the
    other agents they can hand off to, so they know valid target names
    without any hard-coding in their instructions.

    Example::

        researcher = Agent(
            name="researcher",
            description="Gathers facts and data on a topic",
            instructions=(
                "You research topics thoroughly.  When you have gathered"
                " enough information, hand off to the writer.\\n"
                "To hand off write: HANDOFF: writer"
            ),
            model_client=...,
        )

        writer = Agent(
            name="writer",
            description="Writes clear articles from research notes",
            instructions=(
                "You write well-structured articles.  After writing, hand"
                " off to the reviewer.\\nTo hand off write: HANDOFF: reviewer"
            ),
            model_client=...,
        )

        reviewer = Agent(
            name="reviewer",
            description="Reviews and approves or requests revisions",
            instructions=(
                "Review the article.  If it meets quality standards say"
                " APPROVED, otherwise hand back to the writer with feedback.\\n"
                "To hand off write: HANDOFF: writer"
            ),
            model_client=...,
        )

        orchestrator = HandoffOrchestrator(
            agents=[researcher, writer, reviewer],
            termination=TextMentionTermination("APPROVED"),
            entry_agent_name="researcher",
        )

        result = await orchestrator.run("Write an article about quantum computing")
    """

    #: Pattern that agents embed to request a handoff.
    #: Matches ``HANDOFF: <agent_name>`` anywhere in a response (case-insensitive).
    DEFAULT_SIGNAL_PREFIX: str = "HANDOFF:"

    def __init__(
        self,
        agents: Sequence[BaseAgent],
        termination: BaseTermination,
        entry_agent_name: str | None = None,
        fallback_strategy: Literal["round_robin", "entry"] = "round_robin",
        handoff_signal_prefix: str = DEFAULT_SIGNAL_PREFIX,
        inject_colleagues_context: bool = True,
        max_iterations: int = 50,
        history_policy: HistoryContextPolicy | None = None,
    ):
        """
        Initialise the handoff orchestrator.

        Args:
            agents: Agents available for selection via handoffs.
            termination: When to stop the orchestration run.
            entry_agent_name: Name of the first agent to execute.
                Defaults to the first agent in the list.
            fallback_strategy: What to do when an agent produces no
                handoff signal. ``"round_robin"`` cycles through agents;
                ``"entry"`` always routes back to the entry agent.
            handoff_signal_prefix: Token that starts a handoff directive
                in an agent response (default ``"HANDOFF:"``).
            inject_colleagues_context: When ``True`` a "Colleagues you
                can hand off to" block is prepended to each agent's
                context so they know the available targets.
            max_iterations: Safety cap on orchestration iterations.
            history_policy: History window/format configuration.
        """
        super().__init__(
            agents,
            termination,
            max_iterations,
            history_policy=history_policy,
        )
        self.fallback_strategy = fallback_strategy
        self.handoff_signal_prefix = handoff_signal_prefix.upper()
        self.inject_colleagues_context = inject_colleagues_context

        # Resolve entry agent
        if entry_agent_name is not None:
            entry = self._find_agent_by_name(entry_agent_name)
            if entry is None:
                raise ValueError(
                    f"entry_agent_name '{entry_agent_name}' not found in agents"
                )
            self._entry_agent = entry
        else:
            self._entry_agent = self.agents[0]

        # Runtime state
        self._next_agent: BaseAgent | None = None
        self._fallback_index: int = 0
        self._last_handoff_target: str | None = None
        self._handoff_counts: dict[str, int] = {}

    # ------------------------------------------------------------------
    # BaseOrchestrator abstract methods
    # ------------------------------------------------------------------

    async def select_next_agent(self) -> BaseAgent:
        """
        Select the next agent.

        On the very first iteration the entry agent is chosen.
        On subsequent iterations, the last assistant message is scanned
        for a handoff signal.  The signal names the next agent directly.
        When no signal is found the fallback strategy is applied.
        """
        if self.iteration_count == 0:
            self._next_agent = self._entry_agent
            return self._entry_agent

        target = self._extract_handoff_target()
        if target is not None:
            agent = self._find_agent_by_name(target)
            if agent is not None:
                self._last_handoff_target = agent.name
                self._handoff_counts[agent.name] = (
                    self._handoff_counts.get(agent.name, 0) + 1
                )
                self._next_agent = agent
                return agent

            # Signal present but unresolvable — treat as no handoff
            self._last_handoff_target = None

        return self._select_fallback()

    async def prepare_context_for_agent(
        self, agent: BaseAgent
    ) -> UserMessage | list[Message]:
        """Build execution context from history policy, optionally injecting colleagues."""
        context_payload = self._build_agent_context_from_history(agent)

        if not self.inject_colleagues_context:
            return context_payload

        colleagues_block = self._build_colleagues_block(agent)
        if not colleagues_block:
            return context_payload

        # Prepend the colleagues block to the content of a UserMessage,
        # or prepend as a new UserMessage when the payload is a list.
        if isinstance(context_payload, UserMessage):
            merged_content = f"{colleagues_block}\n\n{context_payload.content}"
            return UserMessage(content=merged_content, source="orchestrator")

        # List[Message] — insert a UserMessage at position 0
        colleagues_msg = UserMessage(
            content=colleagues_block, source="orchestrator"
        )
        return [colleagues_msg, *context_payload]

    async def update_shared_state(self, result: AgentResponse) -> None:
        """Append new agent messages to shared orchestration history."""
        if self._last_context_was_text:
            new_messages = result.messages[1:] if len(result.messages) > 1 else []
        else:
            new_messages = result.messages[self._last_context_size :]
        self.shared_messages.extend(new_messages)

    # ------------------------------------------------------------------
    # Metadata and reset
    # ------------------------------------------------------------------

    def _get_pattern_metadata(self) -> dict[str, Any]:
        """Return handoff-specific metadata in addition to base fields."""
        base_metadata = super()._get_pattern_metadata()
        base_metadata.update(
            {
                "entry_agent": self._entry_agent.name,
                "fallback_strategy": self.fallback_strategy,
                "handoff_signal_prefix": self.handoff_signal_prefix,
                "handoff_counts": dict(self._handoff_counts),
                "last_handoff_target": self._last_handoff_target,
            }
        )
        return base_metadata

    def _get_selection_reason(self, selected_agent: BaseAgent) -> str:
        """Return a human-readable selection reason for verbose events."""
        if self.iteration_count == 0:
            return "Entry agent for handoff orchestration"
        if self._last_handoff_target == selected_agent.name:
            return f"Handoff signal → '{selected_agent.name}'"
        return f"Fallback ({self.fallback_strategy}) → '{selected_agent.name}'"

    def _reset_for_run(self) -> None:
        """Reset handoff runtime state before a new run."""
        super()._reset_for_run()
        self._next_agent = None
        self._fallback_index = 0
        self._last_handoff_target = None
        self._handoff_counts = {}

    def _default_history_policy(self) -> HistoryContextPolicy:
        """Full conversation text by default for handoff agents."""
        return HistoryContextPolicy(
            default=AgentHistoryPolicy(window="full", format="text")
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _extract_handoff_target(self) -> str | None:
        """
        Scan the most recent assistant message for a handoff signal.

        Looks for ``<HANDOFF_PREFIX> <agent_name>`` anywhere in the
        message content (case-insensitive prefix match, then the first
        word after it is taken as the agent name).

        Returns:
            The raw target string from the message, or ``None`` when no
            signal is present.
        """
        return self._extract_handoff_target_from_messages(self.shared_messages)

    def _extract_handoff_target_from_messages(
        self,
        messages: Sequence[Message],
    ) -> str | None:
        """Scan the most recent assistant/tool message for a handoff signal."""
        for message in reversed(messages):
            role = getattr(message, "role", "unknown")
            if role in ("assistant", "tool"):
                content_upper = message.content.upper()
                prefix = self.handoff_signal_prefix
                idx = content_upper.rfind(prefix)
                if idx == -1:
                    return None

                remainder = message.content[idx + len(prefix) :].strip()
                # Take the first token (strip punctuation) as the target name
                match = re.match(r"[\w\-]+", remainder)
                if match:
                    return match.group(0)
                return None

        return None

    def _select_fallback(self) -> BaseAgent:
        """Apply the configured fallback strategy when no handoff signal is found."""
        self._last_handoff_target = None

        if self.fallback_strategy == "entry":
            self._next_agent = self._entry_agent
            return self._entry_agent

        # round_robin fallback
        agent = self.agents[self._fallback_index]
        self._fallback_index = (self._fallback_index + 1) % len(self.agents)
        self._next_agent = agent
        return agent

    def _find_agent_by_name(self, name: str) -> BaseAgent | None:
        """Resolve agent by exact name first, then case-insensitive match."""
        for agent in self.agents:
            if agent.name == name:
                return agent
        target = name.casefold()
        for agent in self.agents:
            if agent.name.casefold() == target:
                return agent
        return None

    def _build_colleagues_block(self, current_agent: BaseAgent) -> str:
        """
        Build a short summary of the other agents that ``current_agent``
        can hand off to.

        Returns an empty string when there are no other agents.
        """
        others = [a for a in self.agents if a.name != current_agent.name]
        if not others:
            return ""

        lines = [
            "Colleagues you can hand off to (write HANDOFF: <name> to transfer control):"
        ]
        for agent in others:
            lines.append(f"  - {agent.name}: {agent.description}")
        return "\n".join(lines)
