"""AI-driven orchestration pattern."""

from __future__ import annotations

import json
from typing import Any, Optional, Sequence

from pydantic import BaseModel, ConfigDict, Field

from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse, Usage
from agentbyte.llm.base import BaseChatCompletionClient
from agentbyte.messages import Message, UserMessage
from agentbyte.termination import BaseTermination

from .base import BaseOrchestrator
from .policies import (
    AgentHistoryPolicy,
    FinalResultPolicy,
    HistoryContextPolicy,
    SelectorContextPolicy,
)


class AgentSelection(BaseModel):
    """Structured output model for AI-driven agent selection."""

    model_config = ConfigDict(frozen=True)

    selected_agent: str = Field(..., description="Name of the agent to execute next")
    reasoning: str = Field(
        ...,
        description="Why this agent was selected given current conversation needs",
    )


class AIOrchestrator(BaseOrchestrator):
    """
    AI-driven orchestration pattern.

    Uses an LLM to select the next agent based on conversation state and
    agent capabilities, while reusing the same base orchestration loop.

    If AI selection fails, falls back to round-robin selection to maintain
    forward progress.
    """

    def __init__(
        self,
        agents: Sequence[BaseAgent],
        termination: BaseTermination,
        model_client: BaseChatCompletionClient,
        max_iterations: int = 50,
        history_policy: HistoryContextPolicy | None = None,
        selector_policy: SelectorContextPolicy | None = None,
        final_result_policy: FinalResultPolicy | None = None,
    ):
        super().__init__(
            agents,
            termination,
            max_iterations,
            history_policy=history_policy,
            final_result_policy=final_result_policy,
        )
        self.model_client = model_client
        self.selector_policy = selector_policy or SelectorContextPolicy()

        self._fallback_agent_index = 0
        self._last_selection_reason = ""
        self._selector_usage_stats: list[Usage] = []

    async def select_next_agent(self) -> BaseAgent:
        """Select next agent using structured LLM output with safe fallback."""
        try:
            selection = await self._select_agent_with_llm()
            selected_agent = self._find_agent_by_name(selection.selected_agent)
            if selected_agent is None:
                return self._select_next_by_fallback(
                    f"Selector returned unknown agent '{selection.selected_agent}'"
                )

            self._last_selection_reason = selection.reasoning.strip() or (
                f"LLM selected '{selected_agent.name}'"
            )
            return selected_agent
        except Exception as exc:
            return self._select_next_by_fallback(
                f"Selector failed ({type(exc).__name__}: {exc})"
            )

    async def prepare_context_for_agent(
        self, agent: BaseAgent
    ) -> UserMessage | list[Message]:
        """Build worker-agent execution context from the history policy."""
        return self._build_agent_context_from_history(agent)

    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)

    def _get_pattern_metadata(self) -> dict[str, Any]:
        """Get AI-driven pattern metadata."""
        base_metadata = super()._get_pattern_metadata()
        base_metadata.update(
            {
                "selector_policy": self.selector_policy.model_dump(),
                "selector_calls": len(self._selector_usage_stats),
                "selector_model": self.model_client.model,
            }
        )
        return base_metadata

    def _get_selection_reason(self, selected_agent: BaseAgent) -> str:
        """Get selection reason for verbose selection events."""
        if self._last_selection_reason:
            return self._last_selection_reason
        return super()._get_selection_reason(selected_agent)

    def _get_additional_usage_stats(self) -> list[Usage]:
        """Include selector LLM calls in orchestration usage aggregation."""
        return list(self._selector_usage_stats)

    def _reset_for_run(self) -> None:
        """Reset AI-driven runtime state for new run."""
        super()._reset_for_run()
        self._fallback_agent_index = 0
        self._last_selection_reason = ""
        self._selector_usage_stats = []

    async def _select_agent_with_llm(self) -> AgentSelection:
        """Request structured agent selection from model client."""
        prompt = self._build_selection_prompt()
        result = await self.model_client.create(
            messages=[UserMessage(content=prompt, source="orchestrator")],
            output_format=AgentSelection,
        )

        self._selector_usage_stats.append(result.usage)

        structured_output = result.structured_output
        if isinstance(structured_output, AgentSelection):
            return structured_output
        if isinstance(structured_output, dict):
            return AgentSelection.model_validate(structured_output)

        raise ValueError("Selector response did not include AgentSelection output")

    def _build_selection_prompt(self) -> str:
        """Build LLM prompt for choosing the next best agent."""
        agent_names = ", ".join(agent.name for agent in self.agents)
        capabilities = self.get_agent_capabilities_summary()
        conversation_context = self._format_conversation_for_selection()

        return (
            "You are coordinating a multi-agent team.\n"
            "Select exactly one agent name from the available agents.\n\n"
            f"Available agent names: {agent_names}\n"
            "Agent capabilities:\n"
            f"{capabilities}\n\n"
            "Recent conversation context:\n"
            f"{conversation_context}\n\n"
            "Return structured output with:\n"
            "- selected_agent: exact agent name\n"
            "- reasoning: brief rationale\n"
        )

    def _format_conversation_for_selection(self) -> str:
        """Format conversation history for the selector model."""
        messages = self._select_selector_messages()
        if not messages:
            return "No prior conversation yet."

        if self.selector_policy.format == "summary":
            return self._format_selector_summary(messages)
        if self.selector_policy.format == "markdown":
            return self._format_selector_markdown(messages)
        return self._format_selector_json(messages)

    def _select_selector_messages(self) -> list[Message]:
        """Slice selector-visible messages using selector policy."""
        if (
            self.selector_policy.window == "recent"
            and self.selector_policy.message_limit is not None
        ):
            return list(self.shared_messages[-self.selector_policy.message_limit :])
        return list(self.shared_messages)

    def _format_selector_summary(self, messages: list[Message]) -> str:
        """Format selector context as compact bullet lines."""
        lines: list[str] = []
        for message in messages:
            role = getattr(message, "role", "unknown")
            content = message.content.replace("\n", " ").strip()
            if len(content) > 500:
                content = f"{content[:500]}..."
            lines.append(f"- {message.source} ({role}): {content}")
        return "\n".join(lines)

    def _format_selector_markdown(self, messages: list[Message]) -> str:
        """Format selector context as a markdown transcript."""
        lines: list[str] = ["# Recent Conversation", ""]
        for message in messages:
            lines.extend(
                [
                    f"## {message.source} ({getattr(message, 'role', 'unknown')})",
                    message.content,
                    "",
                ]
            )
        return "\n".join(lines).strip()

    def _format_selector_json(self, messages: list[Message]) -> str:
        """Format selector context as JSON."""
        payload = {
            "conversation": [
                {
                    "source": message.source,
                    "role": getattr(message, "role", "unknown"),
                    "content": message.content,
                }
                for message in messages
            ]
        }
        return json.dumps(payload, indent=2)

    def _find_agent_by_name(self, agent_name: str) -> Optional[BaseAgent]:
        """Resolve agent by exact name first, then case-insensitive match."""
        for agent in self.agents:
            if agent.name == agent_name:
                return agent

        target = agent_name.casefold()
        for agent in self.agents:
            if agent.name.casefold() == target:
                return agent

        return None

    def _select_next_by_fallback(self, reason: str) -> BaseAgent:
        """Round-robin fallback selection used when AI selection fails."""
        agent = self.agents[self._fallback_agent_index]
        self._fallback_agent_index = (self._fallback_agent_index + 1) % len(self.agents)
        self._last_selection_reason = (
            f"{reason}. Fallback round-robin selected '{agent.name}'"
        )
        return agent

    def _default_history_policy(self) -> HistoryContextPolicy:
        """Use full conversation text by default for worker-agent execution."""
        return HistoryContextPolicy(
            default=AgentHistoryPolicy(window="full", format="text")
        )
