from __future__ import annotations

import asyncio
import inspect
from typing import Any, Awaitable, Callable, Literal, Sequence

from autogen_agentchat.agents import UserProxyAgent, BaseChatAgent
from autogen_agentchat.base import Response
from autogen_agentchat.messages import ChatMessage, TextMessage
from autogen_core import CancellationToken

UserInputCallback = Callable[[Sequence[ChatMessage]], str | Awaitable[str]]


def _safe_to_text(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    return str(value)


def _build_native_user_proxy_agent(
    *,
    name: str,
    description: str,
    input_func: Callable[..., Awaitable[str]],
    timeout: float | None,
) -> UserProxyAgent | None:
    """
    Best-effort build of AutoGen UserProxyAgent across constructor variants.

    Falls back to custom agents if the installed UserProxyAgent signature does
    not support callback-based input.
    """
    try:
        sig = inspect.signature(UserProxyAgent)
        params = sig.parameters
        kwargs: dict[str, Any] = {}

        if "name" in params:
            kwargs["name"] = name
        if "description" in params:
            kwargs["description"] = description

        input_param = None
        for candidate in ("input_func", "input_function", "input_callback"):
            if candidate in params:
                input_param = candidate
                break
        if input_param is None:
            return None
        kwargs[input_param] = input_func

        if timeout is not None and "timeout" in params:
            kwargs["timeout"] = timeout

        return UserProxyAgent(**kwargs)
    except Exception:
        return None


class _BaseUserProxyAgent(BaseChatAgent):
    """Shared behaviors for custom UserProxy implementations."""

    @property
    def produced_message_types(self) -> Sequence[type[ChatMessage]]:
        return [TextMessage]

    def _response(self, content: Any) -> Response:
        return Response(
            chat_message=TextMessage(
                content=_safe_to_text(content),
                source=self.name,
            )
        )


class WebUserProxyAgent(_BaseUserProxyAgent):
    """A UserProxy that gets input from an async queue instead of console."""

    def __init__(
        self,
        name: str = "UserProxy",
        description: str = "Human user proxy that accepts input via WebSocket",
        input_queue: asyncio.Queue | None = None,
        timeout: float = 3600.0,  # 1 hour timeout
    ) -> None:
        super().__init__(name=name, description=description)
        self._input_queue = input_queue
        self._timeout = timeout

    async def on_messages(
        self,
        messages: Sequence[ChatMessage],
        cancellation_token: CancellationToken,
    ) -> Response:
        """Wait for user input from the queue."""
        if self._input_queue is None:
            return self._response("No input queue configured. Please provide input via WebSocket.")

        try:
            # Wait for input with timeout
            user_input = await asyncio.wait_for(
                self._input_queue.get(),
                timeout=self._timeout,
            )
            return self._response(user_input)
        except asyncio.TimeoutError:
            return self._response("User input timed out. Please provide input and retry.")

    async def on_reset(self, cancellation_token: CancellationToken) -> None:
        """Reset the agent state."""
        # Clear any pending items in the queue
        if self._input_queue is not None:
            while not self._input_queue.empty():
                try:
                    self._input_queue.get_nowait()
                except asyncio.QueueEmpty:
                    break


class ConsoleUserProxyAgent(_BaseUserProxyAgent):
    """A UserProxy that reads input from local console (CLI mode)."""

    def __init__(
        self,
        name: str = "UserProxy",
        description: str = "Human user proxy that accepts input from console",
        prompt: str = "User input> ",
        timeout: float | None = None,
    ) -> None:
        super().__init__(name=name, description=description)
        self._prompt = prompt
        self._timeout = timeout

    async def on_messages(
        self,
        messages: Sequence[ChatMessage],
        cancellation_token: CancellationToken,
    ) -> Response:
        try:
            if self._timeout is None:
                user_input = await asyncio.to_thread(input, self._prompt)
            else:
                user_input = await asyncio.wait_for(
                    asyncio.to_thread(input, self._prompt),
                    timeout=self._timeout,
                )
            return self._response(user_input)
        except asyncio.TimeoutError:
            return self._response("User input timed out. Please provide input and retry.")
        except EOFError:
            return self._response("No console input available.")

    async def on_reset(self, cancellation_token: CancellationToken) -> None:
        return


class CallbackUserProxyAgent(_BaseUserProxyAgent):
    """A UserProxy that obtains input from a callback (SDK mode)."""

    def __init__(
        self,
        name: str = "UserProxy",
        description: str = "Human user proxy that accepts input from callback",
        input_callback: UserInputCallback | None = None,
        timeout: float | None = None,
    ) -> None:
        super().__init__(name=name, description=description)
        self._input_callback = input_callback
        self._timeout = timeout

    async def on_messages(
        self,
        messages: Sequence[ChatMessage],
        cancellation_token: CancellationToken,
    ) -> Response:
        if self._input_callback is None:
            return self._response("No input callback configured.")

        async def _invoke_callback() -> Any:
            result = self._input_callback(messages)
            if asyncio.iscoroutine(result):
                return await result
            return result

        try:
            if self._timeout is None:
                user_input = await _invoke_callback()
            else:
                user_input = await asyncio.wait_for(
                    _invoke_callback(),
                    timeout=self._timeout,
                )
            return self._response(user_input)
        except asyncio.TimeoutError:
            return self._response("User input timed out. Please provide input and retry.")

    async def on_reset(self, cancellation_token: CancellationToken) -> None:
        return


def build_user_proxy_agent(
    *,
    name: str = "UserProxy",
    mode: Literal["web", "console", "callback"] | None = None,
    input_queue: asyncio.Queue | None = None,
    input_callback: UserInputCallback | None = None,
    timeout: float = 3600.0,
    prefer_native: bool = False,
) -> BaseChatAgent:
    """
    Build a UserProxy agent for API/SDK/CLI.

    Compatibility:
    - Existing callers passing only input_queue keep WebUserProxy behavior.
    - If mode is omitted: queue -> web, callback -> callback, otherwise console.
    """
    if mode is None:
        if input_queue is not None:
            mode = "web"
        elif input_callback is not None:
            mode = "callback"
        else:
            mode = "console"

    if mode == "web":
        async def _web_input_func(*args: Any, **kwargs: Any) -> str:
            if input_queue is None:
                return "No input queue configured. Please provide input via WebSocket."
            try:
                return _safe_to_text(await asyncio.wait_for(input_queue.get(), timeout=timeout))
            except asyncio.TimeoutError:
                return "User input timed out. Please provide input and retry."

        if prefer_native:
            native = _build_native_user_proxy_agent(
                name=name,
                description="Human user proxy that accepts input via WebSocket",
                input_func=_web_input_func,
                timeout=timeout,
            )
            if native is not None:
                return native

        return WebUserProxyAgent(name=name, input_queue=input_queue, timeout=timeout)

    if mode == "callback":
        async def _callback_input_func(*args: Any, **kwargs: Any) -> str:
            if input_callback is None:
                return "No input callback configured."

            messages: Sequence[ChatMessage] = []
            if args and isinstance(args[0], Sequence):
                candidate = args[0]
                if not isinstance(candidate, (str, bytes)):
                    messages = candidate  # type: ignore[assignment]

            result = input_callback(messages)
            if asyncio.iscoroutine(result):
                result = await result
            return _safe_to_text(result)

        if prefer_native:
            native = _build_native_user_proxy_agent(
                name=name,
                description="Human user proxy that accepts input from callback",
                input_func=_callback_input_func,
                timeout=timeout,
            )
            if native is not None:
                return native

        return CallbackUserProxyAgent(
            name=name,
            input_callback=input_callback,
            timeout=timeout,
        )

    if mode == "console":
        async def _console_input_func(*args: Any, **kwargs: Any) -> str:
            prompt = "User input> "
            if args and isinstance(args[0], str):
                prompt = args[0]
            try:
                return _safe_to_text(await asyncio.to_thread(input, prompt))
            except EOFError:
                return "No console input available."

        if prefer_native:
            native = _build_native_user_proxy_agent(
                name=name,
                description="Human user proxy that accepts input from console",
                input_func=_console_input_func,
                timeout=None,
            )
            if native is not None:
                return native

        # Console mode defaults to no timeout unless caller specifies one explicitly.
        return ConsoleUserProxyAgent(name=name, timeout=None)
    raise ValueError(f"Unsupported user proxy mode: {mode}")

