import asyncio
import time

from httpx import AsyncClient
from httpx_sse import aconnect_sse
from pydantic import TypeAdapter

from progtc.types import (
    ExecuteCodeError,
    ExecuteCodeMessage,
    ExecuteCodeSuccess,
    ToolCall,
    ToolHandler,
)

ExecuteCodeMessageAdapter = TypeAdapter[ExecuteCodeMessage](ExecuteCodeMessage)


DEFAULT_PING_TIMEOUT = 5.0
DEFAULT_WAIT_FOR_SERVER_TIMEOUT = 10.0
DEFAULT_CODE_EXECUTION_TIMEOUT = 30.0
DEFAULT_TOOL_CALL_RESULT_TIMEOUT = 5.0


class AsyncProgtcClient:
    def __init__(
        self,
        base_url: str,
        api_key: str,
        headers: dict[str, str] | None = None,
    ):
        self._base_url = base_url
        self._api_key = api_key
        self._headers = headers or {}

    async def ping(
        self,
        timeout: float | None = DEFAULT_PING_TIMEOUT,
        headers: dict[str, str] | None = None,
    ) -> bool:
        merged_headers = {**self._headers, **(headers or {})}
        async with AsyncClient(
            base_url=self._base_url, headers=merged_headers, timeout=timeout
        ) as client:
            try:
                response = await client.get("/ping")
                return response.status_code == 200
            except Exception:
                return False

    async def wait_for_server(
        self,
        timeout: float | None = DEFAULT_WAIT_FOR_SERVER_TIMEOUT,
        headers: dict[str, str] | None = None,
    ) -> None:
        t0 = time.monotonic()
        ping_timeout = (
            min(timeout, DEFAULT_PING_TIMEOUT)
            if timeout is not None
            else DEFAULT_PING_TIMEOUT
        )
        while not await self.ping(timeout=ping_timeout, headers=headers):
            await asyncio.sleep(0.5)
            if timeout is not None and time.monotonic() - t0 > timeout:
                raise TimeoutError(f"Progtc server not ready after {timeout}s")

    async def execute_code(
        self,
        code: str,
        tools: dict[str, ToolHandler],
        timeout: float | None = DEFAULT_CODE_EXECUTION_TIMEOUT,
        headers: dict[str, str] | None = None,
    ) -> ExecuteCodeSuccess | ExecuteCodeError:
        merged_headers = {
            **self._headers,
            **(headers or {}),
            "X-Progtc-API-Key": self._api_key,
        }
        async with (
            AsyncClient(
                base_url=self._base_url,
                headers=merged_headers,
                timeout=timeout,
            ) as client,
            aconnect_sse(
                client,
                "POST",
                "/execute-code",
                json={
                    "tool_names": list(tools.keys()),
                    "code": code,
                },
            ) as event_source,
        ):

            async def handle_tool_call(tool_call: ToolCall) -> None:
                handler = tools[tool_call.tool_name]
                result = await handler(*tool_call.args, **tool_call.kwargs)
                await client.post(
                    "/tool-results",
                    json={
                        "execution_id": tool_call.execution_id,
                        "tool_call_id": tool_call.id,
                        "result": result,
                    },
                    timeout=DEFAULT_TOOL_CALL_RESULT_TIMEOUT,
                )

            pending_tasks: list[asyncio.Task[None]] = []
            async for event in event_source.aiter_sse():
                event_obj = ExecuteCodeMessageAdapter.validate_python(event.json())
                if isinstance(event_obj, ToolCall):
                    task = asyncio.create_task(handle_tool_call(event_obj))
                    pending_tasks.append(task)
                else:
                    if pending_tasks:
                        await asyncio.gather(*pending_tasks)
                    return event_obj

        raise RuntimeError("Stream ended without result")
