from __future__ import annotations

import asyncio
import copy
from dataclasses import dataclass
from typing import Any, Callable, Mapping

from .agent_loop import agent_loop, agent_loop_continue
from .types import (
    AgentMessage,
    AgentOptions,
    AgentState,
    DEFAULT_MODEL,
    Message,
    Model,
    ThinkingBudgets,
    ThinkingLevel,
    Transport,
    clone_message,
    clone_messages,
    get_value,
    message_content,
    message_role,
    now_ms,
)


@dataclass(slots=True)
class AbortSignal:
    aborted: bool = False


class AbortController:
    def __init__(self) -> None:
        self.signal = AbortSignal()

    def abort(self) -> None:
        self.signal.aborted = True


class Agent:
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        options: dict[str, Any] = {}
        if args:
            if len(args) > 1:
                raise TypeError("Agent accepts at most one positional options argument")
            if not isinstance(args[0], Mapping):
                raise TypeError("Agent positional argument must be a mapping")
            options.update(dict(args[0]))
        options.update(kwargs)

        self._listeners: list[Callable[[dict[str, Any]], None]] = []
        self._steering_queue: list[AgentMessage] = []
        self._follow_up_queue: list[AgentMessage] = []
        self._abort_controller: AbortController | None = None

        self._stream_fn = get_value(options, "stream_fn", "streamFn", default=None)
        self._convert_to_llm = get_value(options, "convert_to_llm", "convertToLlm", default=self._default_convert_to_llm)
        self._transform_context = get_value(options, "transform_context", "transformContext", default=None)
        self._get_api_key = get_value(options, "get_api_key", "getApiKey", default=None)

        self._steering_mode = get_value(options, "steering_mode", "steeringMode", default="one-at-a-time")
        self._follow_up_mode = get_value(options, "follow_up_mode", "followUpMode", default="one-at-a-time")

        self.thinking_budgets: ThinkingBudgets = get_value(
            options,
            "thinking_budgets",
            "thinkingBudgets",
            default={"minimal": 0, "low": 0, "medium": 0, "high": 0},
        )
        self.transport: Transport = get_value(options, "transport", default="sse")
        self.max_retry_delay_ms: int = int(get_value(options, "max_retry_delay_ms", "maxRetryDelayMs", default=1000) or 1000)

        self._session_id: str | None = get_value(options, "session_id", "sessionId", default=None)

        initial_state = get_value(options, "initial_state", "initialState", default=None)
        self._state = self._build_initial_state(initial_state)

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

    def subscribe(self, listener: Callable[[dict[str, Any]], None]) -> Callable[[], None]:
        self._listeners.append(listener)

        def unsubscribe() -> None:
            if listener in self._listeners:
                self._listeners.remove(listener)

        return unsubscribe

    def _emit(self, event: dict[str, Any]) -> None:
        for listener in list(self._listeners):
            try:
                listener(event)
            except Exception:
                continue

    def set_system_prompt(self, system_prompt: str) -> None:
        self._state.system_prompt = system_prompt

    def setSystemPrompt(self, system_prompt: str) -> None:
        self.set_system_prompt(system_prompt)

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

    def setModel(self, model: Model) -> None:
        self.set_model(model)

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

    def setThinkingLevel(self, level: ThinkingLevel) -> None:
        self.set_thinking_level(level)

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

    def setTools(self, tools: list[Any]) -> None:
        self.set_tools(tools)

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

    def replaceMessages(self, messages: list[AgentMessage]) -> None:
        self.replace_messages(messages)

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

    def appendMessage(self, message: AgentMessage) -> None:
        self.append_message(message)

    def clear_messages(self) -> None:
        self._state.messages = []

    def clearMessages(self) -> None:
        self.clear_messages()

    def steer(self, message: AgentMessage) -> None:
        self._steering_queue.append(clone_message(message))

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

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

    def abort(self) -> None:
        if self._abort_controller is not None:
            self._abort_controller.abort()

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

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

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

    @sessionId.setter
    def sessionId(self, value: str | None) -> None:
        self._session_id = value

    async def prompt(self, input_value: str | AgentMessage, images: list[dict[str, Any]] | None = None) -> list[AgentMessage]:
        self._ensure_not_streaming()
        user_message = self._make_user_message(input_value, images=images)
        return await self._run_loop(prompts=[user_message], is_continue=False)

    async def continue_(self) -> list[AgentMessage]:
        self._ensure_not_streaming()

        queued_steering = self._dequeue(self._steering_queue, self._steering_mode)
        if queued_steering:
            return await self._run_loop(prompts=queued_steering, is_continue=False)

        queued_follow_up = self._dequeue(self._follow_up_queue, self._follow_up_mode)
        if queued_follow_up:
            return await self._run_loop(prompts=queued_follow_up, is_continue=False)

        return await self._run_loop(prompts=[], is_continue=True)

    async def resume(self) -> list[AgentMessage]:
        return await self.continue_()

    async def run_continue(self) -> list[AgentMessage]:
        return await self.continue_()

    async def wait_for_idle(self) -> None:
        while self._state.is_streaming:
            await asyncio.sleep(0.005)

    async def _run_loop(self, prompts: list[AgentMessage], *, is_continue: bool) -> list[AgentMessage]:
        self._state.is_streaming = True
        self._state.error = None
        self._state.stream_message = None
        self._state.pending_tool_calls.clear()

        self._abort_controller = AbortController()
        skip_initial_steering_poll = len(prompts) > 0
        steering_polled = False

        async def get_steering_messages() -> list[AgentMessage]:
            nonlocal steering_polled
            if skip_initial_steering_poll and not steering_polled:
                steering_polled = True
                return []
            steering_polled = True
            return self._dequeue(self._steering_queue, self._steering_mode)

        async def get_follow_up_messages() -> list[AgentMessage]:
            return self._dequeue(self._follow_up_queue, self._follow_up_mode)

        context: dict[str, Any] = {
            "systemPrompt": self._state.system_prompt,
            "messages": clone_messages(self._state.messages),
            "tools": list(self._state.tools),
        }
        config = self._build_loop_config(
            get_steering_messages=get_steering_messages,
            get_follow_up_messages=get_follow_up_messages,
        )

        try:
            if is_continue:
                stream = agent_loop_continue(
                    context,
                    config,
                    signal=self._abort_controller.signal,
                    stream_fn=self._stream_fn,
                )
            else:
                stream = agent_loop(
                    prompts,
                    context,
                    config,
                    signal=self._abort_controller.signal,
                    stream_fn=self._stream_fn,
                )

            async for event in stream:
                self._handle_stream_event(event)
                self._emit(event)

            new_messages = await stream.result()
            if not isinstance(new_messages, list):
                new_messages = list(new_messages)

            self._state.messages.extend(clone_messages(new_messages))
            return new_messages
        except Exception as exc:
            self._state.error = str(exc)
            raise
        finally:
            self._state.is_streaming = False
            self._state.stream_message = None
            self._state.pending_tool_calls.clear()
            self._abort_controller = None

    def _handle_stream_event(self, event: Mapping[str, Any]) -> None:
        event_type = str(get_value(event, "type", default=""))

        if event_type in {"message_start", "message_update"}:
            message = get_value(event, "message", default=None)
            if message_role(message) == "assistant":
                self._state.stream_message = clone_message(message)
                self._sync_pending_tool_calls(message)

        if event_type == "message_end":
            message = get_value(event, "message", default=None)
            if message_role(message) == "assistant":
                self._state.stream_message = None

        if event_type == "tool_execution_start":
            tool_call_id = get_value(event, "toolCallId", "tool_call_id", default=None)
            if isinstance(tool_call_id, str):
                self._state.pending_tool_calls.add(tool_call_id)

        if event_type == "tool_execution_end":
            tool_call_id = get_value(event, "toolCallId", "tool_call_id", default=None)
            if isinstance(tool_call_id, str):
                self._state.pending_tool_calls.discard(tool_call_id)

    def _sync_pending_tool_calls(self, message: AgentMessage) -> None:
        content = message_content(message)
        if not isinstance(content, list):
            return

        for block in content:
            if isinstance(block, Mapping) and block.get("type") == "toolCall":
                tool_call_id = block.get("id")
                if isinstance(tool_call_id, str):
                    self._state.pending_tool_calls.add(tool_call_id)

    def _build_loop_config(
        self,
        *,
        get_steering_messages: Callable[[], Any] | None = None,
        get_follow_up_messages: Callable[[], Any] | None = None,
    ) -> dict[str, Any]:
        config: dict[str, Any] = {
            "model": self._state.model,
            "convertToLlm": self._convert_to_llm,
            "getSteeringMessages": get_steering_messages or self._get_steering_messages,
            "getFollowUpMessages": get_follow_up_messages or self._get_follow_up_messages,
            "reasoning": self._state.thinking_level,
            "transport": self.transport,
            "thinkingBudgets": self.thinking_budgets,
            "maxRetryDelayMs": self.max_retry_delay_ms,
        }

        if callable(self._transform_context):
            config["transformContext"] = self._transform_context

        if callable(self._get_api_key):
            config["getApiKey"] = self._get_api_key

        if self._session_id is not None:
            config["sessionId"] = self._session_id
            config["session_id"] = self._session_id

        return config

    async def _get_steering_messages(self) -> list[AgentMessage]:
        return self._dequeue(self._steering_queue, self._steering_mode)

    async def _get_follow_up_messages(self) -> list[AgentMessage]:
        return self._dequeue(self._follow_up_queue, self._follow_up_mode)

    @staticmethod
    def _dequeue(queue: list[AgentMessage], mode: str) -> list[AgentMessage]:
        if not queue:
            return []

        if mode == "all":
            values = queue[:]
            queue.clear()
            return values

        return [queue.pop(0)]

    @staticmethod
    def _default_convert_to_llm(messages: list[AgentMessage]) -> list[Message]:
        converted: list[Message] = []
        for message in messages:
            role = message_role(message)
            if role in {"user", "assistant", "toolResult"}:
                converted.append(copy.deepcopy(message))
        return converted

    def _build_initial_state(self, initial_state: Mapping[str, Any] | AgentState | None) -> AgentState:
        if isinstance(initial_state, AgentState):
            return copy.deepcopy(initial_state)

        state = AgentState()
        if initial_state is None:
            return state

        state.system_prompt = str(get_value(initial_state, "system_prompt", "systemPrompt", default="") or "")

        model_value = get_value(initial_state, "model", default=None)
        if isinstance(model_value, Mapping):
            state.model = copy.deepcopy(dict(model_value))
        else:
            state.model = copy.deepcopy(DEFAULT_MODEL)

        state.thinking_level = get_value(initial_state, "thinking_level", "thinkingLevel", default="off")

        tools_value = get_value(initial_state, "tools", default=[])
        state.tools = list(tools_value) if isinstance(tools_value, list) else []

        messages_value = get_value(initial_state, "messages", default=[])
        state.messages = clone_messages(messages_value if isinstance(messages_value, list) else [])

        state.is_streaming = bool(get_value(initial_state, "is_streaming", "isStreaming", default=False))
        state.stream_message = clone_message(get_value(initial_state, "stream_message", "streamMessage", default=None))
        pending_value = get_value(initial_state, "pending_tool_calls", "pendingToolCalls", default=set())
        if isinstance(pending_value, set):
            state.pending_tool_calls = set(str(item) for item in pending_value)
        elif isinstance(pending_value, list):
            state.pending_tool_calls = set(str(item) for item in pending_value)
        else:
            state.pending_tool_calls = set()
        state.error = get_value(initial_state, "error", default=None)

        return state

    def _ensure_not_streaming(self) -> None:
        if self._state.is_streaming:
            raise RuntimeError("Agent is already processing a request")

    @staticmethod
    def _make_user_message(input_value: str | AgentMessage, images: list[dict[str, Any]] | None = None) -> AgentMessage:
        if isinstance(input_value, Mapping) and message_role(input_value):
            message = clone_message(input_value)
            if get_value(message, "timestamp", default=None) is None:
                message["timestamp"] = now_ms()
            return message

        if isinstance(input_value, str):
            if images:
                content: list[dict[str, Any]] = [{"type": "text", "text": input_value}, *images]
            else:
                content = input_value
            return {"role": "user", "content": content, "timestamp": now_ms()}

        raise TypeError("prompt() expects a string or user message mapping")


def _resolve_options_map(options: AgentOptions | Mapping[str, Any]) -> Mapping[str, Any]:
    if isinstance(options, Mapping):
        return options
    return dict(options)
