"""Autonomous agent runner with batched turn execution."""

import copy
import logging
from dataclasses import dataclass
from enum import StrEnum
from typing import (
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Protocol,
    Sequence,
    Union,
    cast,
)
from uuid import uuid4

from openai.types.chat import ChatCompletion
from openai.types.chat.chat_completion_message_tool_call import (
    ChatCompletionMessageToolCall,
)

from swarm.agent import Agent
from swarm.chunk import TokenUsage
from swarm.config.protocols import MessageStoreProtocol
from swarm.context import ContextProtocol, RunContext
from swarm.guardrails import (
    OutputGuardrailContext,
    OutputGuardrailTripped,
    _extract_last_assistant_text,
    run_input_guardrails,
    run_output_guardrails,
)
from swarm.hooks import fire_hook_on_llm_end, fire_hook_on_llm_start
from swarm.message import SwarmAssistantMessage, SwarmMessage
from swarm.providers.types import CompletionRequest, ToolCallDelta
from swarm.tool import (
    ToolExecutionResponse,
    execute_tools,
    prepare_tools_for_completion,
)

logger = logging.getLogger(__name__)


class AgentSignal(StrEnum):
    """Terminal/autonomous signals produced by a batch run."""

    COMPLETED = "completed"
    PAUSED = "paused"
    MAX_TURNS = "max_turns"


@dataclass(slots=True)
class AgentRunResult:
    """Result of an autonomous batch run."""

    signal: AgentSignal
    turns_executed: int
    completion_summary: Optional[str] = None
    pause_reason: Optional[str] = None


@dataclass(slots=True)
class _BatchState:
    """Mutable state for autonomous batch execution."""

    message_history: List[Union[Dict[str, Any], SwarmMessage]]
    context_variables: Dict[str, Any]
    system_prompt: str


class ExecuteToolsFn(Protocol):
    """Callable protocol for tool execution — matches ``execute_tools`` signature."""

    async def __call__(
        self,
        tool_calls: List[ChatCompletionMessageToolCall],
        tools: List[Callable[..., Any]],
        run_context: RunContext[ContextProtocol],
        current_agent_name: Optional[str] = None,
        *,
        hooks: Optional[List[Any]] = None,
        agent: Optional[Agent] = None,
    ) -> ToolExecutionResponse: ...


class AgentRunner:
    """Execute autonomous agent turns without dialogue-only behavior."""

    def __init__(
        self,
        execute_tools_fn: ExecuteToolsFn = execute_tools,
    ) -> None:
        self._execute_tools_fn = execute_tools_fn

    async def run_batch(
        self,
        agent: Agent,
        run_context: RunContext[ContextProtocol],
        max_turns: int,
    ) -> AgentRunResult:
        """Run autonomous turns up to max_turns and return a typed result."""
        if max_turns <= 0:
            return AgentRunResult(signal=AgentSignal.MAX_TURNS, turns_executed=0)

        context = run_context.deps
        router = run_context.router
        state = self._initialize_state(agent, context)

        # Run input guardrails before any LLM call
        await run_input_guardrails(
            agent, list(state.message_history), state.context_variables
        )

        # Main loop — collect result then break so output guardrails run once
        result = AgentRunResult(signal=AgentSignal.MAX_TURNS, turns_executed=max_turns)

        for turn_index in range(max_turns):
            await run_context.publish_status(
                f"Autonomous turn {turn_index + 1}/{max_turns}"
            )

            request = self._build_completion_request(agent, state)

            # LLM hooks with try/finally
            await fire_hook_on_llm_start(
                agent, list(state.message_history), state.context_variables
            )
            completion: Optional[ChatCompletion] = None
            llm_error: BaseException | None = None
            try:
                completion = await router.create_completion_async(request)
            except BaseException as e:
                llm_error = e
                raise
            finally:
                await fire_hook_on_llm_end(
                    agent,
                    completion if llm_error is None else None,
                    error=llm_error,
                )

            assistant_message = self._build_assistant_message(completion, agent)

            state.message_history.append(assistant_message)
            await self._persist_messages(
                messages=[assistant_message],
                context=context,
                context_variables=state.context_variables,
                agent=agent,
                message_store=run_context.message_store,
            )

            tool_calls = self._extract_tool_calls(assistant_message)
            turns_executed = turn_index + 1
            if not tool_calls:
                result = AgentRunResult(
                    signal=AgentSignal.COMPLETED,
                    turns_executed=turns_executed,
                )
                break

            tool_response = await self._execute_tools_fn(
                tool_calls,
                agent.tools,
                run_context,
                agent.name,
                hooks=agent.hooks,
                agent=agent,
            )

            if tool_response.messages:
                state.message_history.extend(tool_response.messages)
                await self._persist_messages(
                    messages=tool_response.messages,
                    context=context,
                    context_variables=state.context_variables,
                    agent=agent,
                    message_store=run_context.message_store,
                )

            self._sync_context_after_tools(state, run_context, tool_response)

            if tool_response.metadata.agent_paused:
                result = AgentRunResult(
                    signal=AgentSignal.PAUSED,
                    turns_executed=turns_executed,
                    pause_reason=tool_response.metadata.pause_reason,
                )
                break
            if tool_response.metadata.agent_completed:
                result = AgentRunResult(
                    signal=AgentSignal.COMPLETED,
                    turns_executed=turns_executed,
                    completion_summary=tool_response.metadata.completion_summary,
                )
                break

        # Run output guardrails (blocking — result not yet returned)
        if getattr(agent, "output_guardrails", None):
            agent_output = _extract_last_assistant_text(state.message_history)
            output_ctx = OutputGuardrailContext(
                agent_name=agent.name,
                agent_output=agent_output,
                context_variables=state.context_variables,
            )
            guardrail_results = await run_output_guardrails(agent, output_ctx)
            for gr in guardrail_results:
                if gr.tripwire_triggered:
                    raise OutputGuardrailTripped(gr.guardrail_name, gr)

        return result

    def _initialize_state(self, agent: Agent, context: ContextProtocol) -> _BatchState:
        """Prepare message history, context variables, and prompt for a batch run."""
        system_prompt = (
            agent.instructions(context)
            if callable(agent.instructions)
            else agent.instructions
        )
        return _BatchState(
            message_history=list(context.get_history_for_completion()),
            context_variables=copy.deepcopy(context.current_context_variables),
            system_prompt=system_prompt,
        )

    def _build_completion_request(
        self,
        agent: Agent,
        state: _BatchState,
    ) -> CompletionRequest:
        """Build a non-streaming completion request for autonomous execution."""
        models = agent.model if isinstance(agent.model, list) else [agent.model]
        include_tools = bool(agent.tools)
        return CompletionRequest(
            model=models,
            stream=False,
            messages=[
                {"role": "system", "content": state.system_prompt},
                *[
                    message if isinstance(message, dict) else message.to_openai_dict()
                    for message in state.message_history
                ],
            ],
            tools=prepare_tools_for_completion(agent.tools) if include_tools else None,
            **(
                agent.model_settings.to_completion_params(has_tools=include_tools)
                if agent.model_settings
                else {}
            ),
        )

    def _build_assistant_message(
        self,
        completion: ChatCompletion,
        agent: Agent,
    ) -> SwarmAssistantMessage:
        """Convert a completion payload into a persisted assistant message."""
        if not completion.choices:
            raise ValueError("Completion response did not include any choices")

        message = completion.choices[0].message
        if message is None:
            raise ValueError("Completion response choice was missing a message")

        tool_calls = self._convert_tool_calls(
            cast(Optional[List[ChatCompletionMessageToolCall]], message.tool_calls)
        )
        fallback_model = self._resolve_designer_model(agent)
        usage = self._extract_usage(completion, fallback_model)

        content_payload: Any = message.content
        if content_payload is None:
            content_payload = []

        return SwarmAssistantMessage(
            id=str(uuid4()),
            content=content_payload,
            sender=agent.id,
            role="assistant",
            tool_calls=tool_calls,
            usage=usage,
        )

    @staticmethod
    def _convert_tool_calls(
        tool_calls: Optional[List[ChatCompletionMessageToolCall]],
    ) -> Optional[List[ToolCallDelta]]:
        """Convert OpenAI tool calls to Swarm ToolCallDelta records."""
        if not tool_calls:
            return None

        return [
            ToolCallDelta(
                id=tool_call.id,
                type=tool_call.type,
                index=index,
                function={
                    "name": tool_call.function.name,
                    "arguments": tool_call.function.arguments,
                },
            )
            for index, tool_call in enumerate(tool_calls)
        ]

    @staticmethod
    def _extract_tool_calls(
        assistant_message: SwarmAssistantMessage,
    ) -> List[ChatCompletionMessageToolCall]:
        """Convert persisted tool deltas back to complete tool call payloads."""
        if not assistant_message.tool_calls:
            return []
        return [
            tool_call.to_completion_tool_call()
            for tool_call in assistant_message.tool_calls
        ]

    @staticmethod
    def _resolve_designer_model(agent: Agent) -> str:
        """Resolve a single model string for persistence and usage fallback."""
        if isinstance(agent.model, list) and agent.model:
            return agent.model[0]
        if isinstance(agent.model, str):
            return agent.model
        return "unknown"

    @staticmethod
    def _extract_usage(
        completion: ChatCompletion,
        fallback_model: str,
    ) -> Optional[TokenUsage]:
        """Extract token usage from non-stream completion payload when available."""
        usage = completion.usage
        if usage is None:
            return None

        model = completion.model if completion.model else fallback_model
        return TokenUsage(
            model=model,
            input_tokens=usage.prompt_tokens or 0,
            output_tokens=usage.completion_tokens or 0,
        )

    @staticmethod
    def _sync_context_after_tools(
        state: _BatchState,
        run_context: RunContext[ContextProtocol],
        tool_response: ToolExecutionResponse,
    ) -> None:
        """Apply tool-returned context variables to batch and shared workflow context."""
        state.context_variables = tool_response.context_variables.copy()
        run_context.deps.workflow_context.update(tool_response.context_variables)

    async def _persist_messages(
        self,
        messages: Sequence[SwarmMessage],
        context: ContextProtocol,
        context_variables: Dict[str, Any],
        agent: Agent,
        message_store: MessageStoreProtocol,
    ) -> None:
        """Persist messages through MessageStoreProtocol with optional billing context."""
        raw_user_id = getattr(context, "user_id", None)
        await message_store.persist_messages(
            messages=messages,
            conversation_id=str(context.conversation_id),
            context_variables=context_variables,
            designer_model=self._resolve_designer_model(agent),
            designer_fullname=context.designer_fullname,
            user_id=str(raw_user_id) if raw_user_id is not None else None,
            audience=getattr(context, "audience", None),
            node_key=getattr(context, "node_key", None),
        )
