"""
Lightweight client for dispatching Claude SDK turns via gRPC.

PM2 agent code creates an instance, calls `start()` to dispatch a turn,
and provides one-shot callbacks that fire when the sidecar finishes.
All communication goes through the existing gRPC broker -- no HTTP.
"""
import asyncio
import uuid
from typing import Any, Callable, Dict


class ClaudeAgentCallbacks:
    """Typed holder for one-shot result callbacks."""

    def __init__(
        self,
        agent_completed: Callable[[Any], None],
        agent_failed: Callable[[Any], None],
    ):
        self.agent_completed = agent_completed
        self.agent_failed = agent_failed


class ClaudeAgentClient:
    """Dispatches Claude turns over gRPC and routes results to callbacks.

    Created once per PM2 process (bound to the StreamManager). Each task
    calls ``start()`` with its own context, params, and one-shot callbacks.
    """

    def __init__(self, stream_manager: Any):
        self._stream_manager = stream_manager

        # requestId -> { future, callbacks }
        self._dispatch_waiters: Dict[str, Dict[str, Any]] = {}
        # jobId -> ClaudeAgentCallbacks
        self._result_callbacks: Dict[str, ClaudeAgentCallbacks] = {}

        # Register message hook on StreamManager (one-time per instance)
        self._stream_manager.on_message(self._handle_message)

    def _handle_message(self, payload: Any) -> bool:
        msg_type = payload.get("_type") if isinstance(payload, dict) else None

        if msg_type == "claude_agent.dispatch_ack":
            request_id = payload.get("requestId")
            waiter = self._dispatch_waiters.pop(request_id, None)
            if waiter:
                self._result_callbacks[payload["jobId"]] = waiter["callbacks"]
                waiter["future"].set_result(
                    {"session_id": payload["sessionId"], "job_id": payload["jobId"]}
                )
            return True

        if msg_type == "claude_agent.dispatch_error":
            request_id = payload.get("requestId")
            waiter = self._dispatch_waiters.pop(request_id, None)
            if waiter:
                waiter["future"].set_exception(
                    RuntimeError(payload.get("error", "Dispatch failed"))
                )
            return True

        if msg_type == "claude_agent.result":
            job_id = payload.get("jobId")
            cbs = self._result_callbacks.pop(job_id, None)  # auto-unregister
            if cbs:
                if payload.get("success"):
                    cbs.agent_completed(payload.get("data"))
                else:
                    cbs.agent_failed(payload.get("error"))
            return True

        return False  # not handled -- pass to regular task processing

    async def start(
        self,
        params: Dict[str, Any],
        callbacks: ClaudeAgentCallbacks,
        context: Dict[str, Any],
    ) -> Dict[str, str]:
        """Dispatch a Claude SDK turn and register one-shot callbacks.

        Args:
            params: Claude SDK parameters (prompt, model, etc.)
            callbacks: One-shot agentCompleted / agentFailed callbacks.
            context: Per-task context with query_id, agent_id, webhook_groups.

        Returns ``{ "session_id": ..., "job_id": ... }`` once dispatch is acked.
        The actual turn result arrives later via the registered callbacks.
        """
        request_id = str(uuid.uuid4())
        loop = asyncio.get_running_loop()
        future = loop.create_future()

        self._dispatch_waiters[request_id] = {
            "future": future,
            "callbacks": callbacks,
        }

        await self._stream_manager.send_message_to_server(
            {
                "_type": "claude_agent.start",
                "requestId": request_id,
                "machineAgentId": self._stream_manager.id,
                "queryId": context.get("query_id"),
                "agentId": context.get("agent_id"),
                "webhookGroups": context.get("webhook_groups"),
                "params": params,
            }
        )

        return await future
