"""Agent class — stateful agent with tool execution and event streaming."""

from __future__ import annotations

import asyncio
from typing import Any, Callable

from pi_agents.ai.types import (
    AssistantMessage,
    ImageContent,
    Message,
    Model,
    TextContent,
    UserMessage,
)

from pi_agents.agent.agent_loop import run_agent_loop, run_agent_loop_continue
from pi_agents.agent.types import (
    AgentEvent,
    AgentMessage,
    AgentState,
    AgentTool,
    BeforeToolCallResult,
    AfterToolCallResult,
    ToolExecutionMode,
)


class Agent:
    """Main agent class that orchestrates LLM interactions."""

    def __init__(self, opts: dict[str, Any] | None = None):
        opts = opts or {}
        initial_state = opts.get("initial_state", opts.get("initialState", {}))

        self._state = AgentState(
            system_prompt=initial_state.get("systemPrompt", initial_state.get("system_prompt", "")),
            model=initial_state.get("model"),
            thinking_level=initial_state.get("thinkingLevel", initial_state.get("thinking_level", "off")),
            tools=initial_state.get("tools", []),
            messages=initial_state.get("messages", []),
        )

        self._convert_to_llm = opts.get("convert_to_llm", opts.get("convertToLlm", self._default_convert))
        self._transform_context = opts.get("transform_context", opts.get("transformContext"))
        self._steering_mode = opts.get("steering_mode", opts.get("steeringMode", "one-at-a-time"))
        self._follow_up_mode = opts.get("follow_up_mode", opts.get("followUpMode", "one-at-a-time"))
        self._stream_fn = opts.get("stream_fn", opts.get("streamFn"))
        self._session_id = opts.get("session_id", opts.get("sessionId"))
        self._get_api_key = opts.get("get_api_key", opts.get("getApiKey"))
        self._thinking_budgets = opts.get("thinking_budgets", opts.get("thinkingBudgets"))
        self._transport = opts.get("transport")
        self._max_retry_delay_ms = opts.get("max_retry_delay_ms", opts.get("maxRetryDelayMs"))
        self._tool_execution: ToolExecutionMode = opts.get("tool_execution", opts.get("toolExecution", "parallel"))
        self._before_tool_call = opts.get("before_tool_call", opts.get("beforeToolCall"))
        self._after_tool_call = opts.get("after_tool_call", opts.get("afterToolCall"))

        self._subscribers: list[Callable[[AgentEvent], None]] = []
        self._steering_queue: list[AgentMessage] = []
        self._follow_up_queue: list[AgentMessage] = []
        self._abort_controller: asyncio.Event | None = None
        self._idle_event = asyncio.Event()
        self._idle_event.set()

    @staticmethod
    def _default_convert(messages: list[AgentMessage]) -> list[Message]:
        """Default message converter — filters to user/assistant/toolResult roles."""
        result = []
        for msg in messages:
            if hasattr(msg, "role"):
                if msg.role in ("user", "assistant", "toolResult"):
                    result.append(msg)
            elif isinstance(msg, dict) and msg.get("role") in ("user", "assistant", "toolResult"):
                result.append(msg)
        return result

    @property
    def state(self) -> AgentState:
        return self._state

    @property
    def session_id(self) -> str | None:
        return self._session_id

    @session_id.setter
    def session_id(self, value: str | None):
        self._session_id = value

    @property
    def tool_execution(self) -> ToolExecutionMode:
        return self._tool_execution

    async def prompt(
        self,
        message: str | AgentMessage | list[AgentMessage],
        images: list[ImageContent] | None = None,
    ) -> None:
        """Send a prompt and run the agent loop until completion."""
        # Build prompt messages
        prompts: list[AgentMessage] = []
        if isinstance(message, str):
            content: Any = message
            if images:
                content = [TextContent(type="text", text=message)] + list(images)
            prompts.append(UserMessage(role="user", content=content, timestamp=0))
        elif isinstance(message, list):
            prompts = message
        else:
            prompts = [message]

        self._state.is_streaming = True
        self._state.error = None
        self._idle_event.clear()

        try:
            context = {
                "system_prompt": self._state.system_prompt,
                "messages": list(self._state.messages),
                "tools": self._state.tools,
            }

            config = self._build_config()

            result = await run_agent_loop(
                prompts=prompts,
                context=context,
                config=config,
                emit=self._emit,
                stream_fn=self._stream_fn,
            )

            self._state.messages = list(context.get("messages", [])) + [m for m in result if m not in context.get("messages", [])]
            # Simplify: just use result messages
            if result:
                self._state.messages = result

        except Exception as e:
            self._state.error = str(e)
        finally:
            self._state.is_streaming = False
            self._state.stream_message = None
            self._state.pending_tool_calls.clear()
            self._idle_event.set()

    async def continue_(self) -> None:
        """Resume from current context."""
        self._state.is_streaming = True
        self._state.error = None
        self._idle_event.clear()

        try:
            context = {
                "system_prompt": self._state.system_prompt,
                "messages": list(self._state.messages),
                "tools": self._state.tools,
            }

            config = self._build_config()

            result = await run_agent_loop_continue(
                context=context,
                config=config,
                emit=self._emit,
                stream_fn=self._stream_fn,
            )

            if result:
                self._state.messages.extend(result)

        except Exception as e:
            self._state.error = str(e)
        finally:
            self._state.is_streaming = False
            self._state.stream_message = None
            self._state.pending_tool_calls.clear()
            self._idle_event.set()

    def abort(self) -> None:
        """Abort the current agent loop."""
        if self._abort_controller:
            self._abort_controller.set()

    async def wait_for_idle(self) -> None:
        """Wait until agent is not streaming."""
        await self._idle_event.wait()

    def reset(self) -> None:
        """Reset agent state."""
        self._state.messages.clear()
        self._state.is_streaming = False
        self._state.stream_message = None
        self._state.pending_tool_calls.clear()
        self._state.error = None
        self._steering_queue.clear()
        self._follow_up_queue.clear()

    def subscribe(self, fn: Callable[[AgentEvent], None]) -> Callable[[], None]:
        """Subscribe to events. Returns unsubscribe function."""
        self._subscribers.append(fn)
        return lambda: self._subscribers.remove(fn) if fn in self._subscribers else None

    # State setters
    def set_system_prompt(self, v: str) -> None:
        self._state.system_prompt = v

    def set_model(self, m: Model) -> None:
        self._state.model = m

    def set_thinking_level(self, level: str) -> None:
        self._state.thinking_level = level

    def set_tools(self, tools: list[AgentTool]) -> None:
        self._state.tools = tools

    def set_tool_execution(self, mode: ToolExecutionMode) -> None:
        self._tool_execution = mode

    def set_steering_mode(self, mode: str) -> None:
        self._steering_mode = mode

    def set_follow_up_mode(self, mode: str) -> None:
        self._follow_up_mode = mode

    def get_steering_mode(self) -> str:
        return self._steering_mode

    def get_follow_up_mode(self) -> str:
        return self._follow_up_mode

    # Message management
    def replace_messages(self, messages: list[AgentMessage]) -> None:
        self._state.messages = list(messages)

    def append_message(self, message: AgentMessage) -> None:
        self._state.messages.append(message)

    def clear_messages(self) -> None:
        self._state.messages.clear()

    # Queue management
    def steer(self, message: AgentMessage) -> None:
        self._steering_queue.append(message)

    def follow_up(self, message: AgentMessage) -> None:
        self._follow_up_queue.append(message)

    def clear_steering_queue(self) -> None:
        self._steering_queue.clear()

    def clear_follow_up_queue(self) -> None:
        self._follow_up_queue.clear()

    def clear_all_queues(self) -> None:
        self._steering_queue.clear()
        self._follow_up_queue.clear()

    def has_queued_messages(self) -> bool:
        return len(self._steering_queue) > 0 or len(self._follow_up_queue) > 0

    def _build_config(self) -> dict[str, Any]:
        """Build agent loop config."""
        return {
            "model": self._state.model,
            "convert_to_llm": self._convert_to_llm,
            "transform_context": self._transform_context,
            "get_api_key": self._get_api_key,
            "get_steering_messages": self._get_steering_messages,
            "get_follow_up_messages": self._get_follow_up_messages,
            "tool_execution": self._tool_execution,
            "before_tool_call": self._before_tool_call,
            "after_tool_call": self._after_tool_call,
            "thinking_level": self._state.thinking_level,
        }

    async def _get_steering_messages(self) -> list[AgentMessage]:
        """Get steering messages from queue."""
        if not self._steering_queue:
            return []
        if self._steering_mode == "all":
            msgs = list(self._steering_queue)
            self._steering_queue.clear()
            return msgs
        else:
            return [self._steering_queue.pop(0)]

    async def _get_follow_up_messages(self) -> list[AgentMessage]:
        """Get follow-up messages from queue."""
        if not self._follow_up_queue:
            return []
        if self._follow_up_mode == "all":
            msgs = list(self._follow_up_queue)
            self._follow_up_queue.clear()
            return msgs
        else:
            return [self._follow_up_queue.pop(0)]

    def _emit(self, event: AgentEvent) -> None:
        """Emit an event to all subscribers."""
        for sub in self._subscribers:
            try:
                sub(event)
            except Exception:
                pass
