from __future__ import annotations

import asyncio
import queue
from multiprocessing.queues import Queue as MPQueueType
from types import TracebackType

from pydantic import JsonValue

from progtc.server.config import server_config
from progtc.types import (
    ExecuteCodeStepState,
    ToolCall,
    ToolCallResult,
    _tool_call_hash,
)


class DeterminismError(Exception):
    pass


class ToolCallManager:
    def __init__(
        self,
        execution_id: str,
        tool_call_queue: MPQueueType[ToolCall],
        tool_call_results_queue: MPQueueType[ToolCallResult],
        state: ExecuteCodeStepState | None = None,
    ) -> None:
        self._execution_id = execution_id
        self._tool_call_queue = tool_call_queue
        self._tool_call_results_queue = tool_call_results_queue
        self._tool_call_results_split_queue: dict[
            str,
            asyncio.Queue[ToolCallResult],
        ] = {}
        self._handle_results_task: asyncio.Task[None] | None = None
        self._shutdown_event = asyncio.Event()
        self._state = state
        self._tool_call_count = 0
        self._lock = asyncio.Lock()

    async def _handle_results_queue(self) -> None:
        while not self._shutdown_event.is_set():
            try:
                tool_call_result = self._tool_call_results_queue.get_nowait()
                await self._tool_call_results_split_queue[
                    tool_call_result.tool_call_id
                ].put(tool_call_result)
            except queue.Empty:
                await asyncio.sleep(0.1)

    async def __aenter__(self) -> ToolCallManager:
        self._handle_results_task = asyncio.create_task(self._handle_results_queue())
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        self._shutdown_event.set()
        if self._handle_results_task is not None:
            await self._handle_results_task

    async def call_tool(
        self,
        tool_name: str,
        *args: JsonValue,
        **kwargs: JsonValue,
    ) -> JsonValue:
        async with self._lock:
            tool_call = ToolCall(
                execution_id=self._execution_id,
                tool_name=tool_name,
                args=args,
                kwargs=kwargs,
            )

            # If we have already received the result for this tool call, return it.
            if self._state is not None and self._tool_call_count < len(
                self._state.tool_results
            ):
                expected_hash = self._state.tool_call_hashes[self._tool_call_count]
                actual_hash = _tool_call_hash(tool_name, args, kwargs)
                if actual_hash != expected_hash:
                    raise DeterminismError(
                        f"Tool call {self._tool_call_count} changed on "
                        f"replay: expected {tool_call.tool_name}() to "
                        f"match the original call, but the arguments "
                        f"differ. The executed code must be deterministic."
                    )
                result_value = self._state.tool_results[self._tool_call_count]
                self._tool_call_count += 1
                return result_value

            self._tool_call_results_split_queue[tool_call.id] = asyncio.Queue[
                ToolCallResult
            ]()
            self._tool_call_queue.put(tool_call)

        try:
            result = await asyncio.wait_for(
                self._tool_call_results_split_queue[tool_call.id].get(),
                timeout=server_config.tool_call_timeout,
            )
            return result.result
        finally:
            del self._tool_call_results_split_queue[tool_call.id]
