"""AskHumanTool — Human-in-the-Loop tool for Chat mode.

When registered in a Chat-mode Agent, this tool allows the ReAct loop
to pause and ask the user a question, then wait for a reply via the
WebSocket connection.
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from typing import Any

from agentkit.tools.base import Tool

logger = logging.getLogger(__name__)


class AskHumanTool(Tool):
    """Tool that asks the human user a question and waits for a reply.

    Only functional in Chat mode where a WebSocket connection exists.
    In Task mode, this tool should not be registered.

    Usage in ReAct loop:
        The Agent calls this tool when it needs clarification or
        a decision from the user. The question is pushed to the
        client via WebSocket, and the tool blocks until the user
        replies or a timeout expires.
    """

    def __init__(self, timeout: float = 60.0):
        super().__init__(
            name="ask_human",
            description="Ask the human user a question and wait for their reply. "
                        "Use this when you need clarification, a decision, or "
                        "confirmation from the user before proceeding.",
        )
        self._timeout = timeout
        # Shared dict injected by the Chat WebSocket handler:
        # request_id -> asyncio.Future
        self._pending_replies: dict[str, asyncio.Future] | None = None
        # Callback to push question to client
        self._ask_callback: Any = None

    def configure(
        self,
        pending_replies: dict[str, asyncio.Future] | None = None,
        ask_callback: Any = None,
    ) -> None:
        """Configure the tool with WebSocket communication channels.

        Args:
            pending_replies: Dict mapping request_id to Future that will
                be resolved when the user replies.
            ask_callback: Async callable(request_id, question, options)
                that pushes the question to the client.
        """
        self._pending_replies = pending_replies
        self._ask_callback = ask_callback

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "question": {
                    "type": "string",
                    "description": "The question to ask the user",
                },
                "options": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Optional list of choices for the user",
                },
            },
            "required": ["question"],
        }

    async def execute(self, **kwargs: Any) -> dict:
        """Ask the user a question and wait for their reply.

        Args:
            question: The question to ask.
            options: Optional list of choices.

        Returns:
            Dict with "reply" key containing the user's response.
        """
        question = kwargs.get("question", "")
        options = kwargs.get("options")

        if self._pending_replies is None or self._ask_callback is None:
            # Not in Chat mode — return a default response
            logger.warning("AskHumanTool called outside Chat mode, returning default response")
            default = options[0] if options else "confirmed"
            return {"reply": default}

        request_id = str(uuid.uuid4())[:8]

        # Create and register future BEFORE calling callback so the
        # callback (or any concurrent task) can resolve it immediately.
        loop = asyncio.get_event_loop()
        future = loop.create_future()
        self._pending_replies[request_id] = future

        # Push question to client
        await self._ask_callback(request_id, question, options)

        try:
            reply = await asyncio.wait_for(future, timeout=self._timeout)
            return {"reply": str(reply)}
        except asyncio.TimeoutError:
            logger.warning(f"AskHumanTool timeout for request {request_id}")
            default = options[0] if options else "timeout — no response received"
            return {"reply": default}
        finally:
            self._pending_replies.pop(request_id, None)