import asyncio
import contextlib
import functools
import inspect
import random
import time
from collections.abc import Awaitable, Callable
from logging import getLogger
from typing import ParamSpec, TypeVar

from httpx import AsyncClient, AsyncHTTPTransport
from httpx_sse import ServerSentEvent, aconnect_sse
from pydantic import JsonValue, TypeAdapter

from progtc.types import (
    ExecuteCodeError,
    ExecuteCodeMessage,
    ExecuteCodeStepMessage,
    ExecuteCodeSuccess,
    ExecutionEvent,
    ExecutionEventHandler,
    ToolCall,
    ToolCallCompletedEvent,
    ToolCallReceivedEvent,
    ToolCallSentEvent,
    ToolHandler,
)

ExecuteCodeMessageAdapter = TypeAdapter[ExecuteCodeMessage](ExecuteCodeMessage)
ExecuteCodeStepMessageAdapter = TypeAdapter[ExecuteCodeStepMessage](
    ExecuteCodeStepMessage
)

logger = getLogger(__name__)

T = TypeVar("T")
P = ParamSpec("P")


def _calculate_backoff(
    attempt: int,
    exp_multiplier: float = 1.0,
    exp_min: float = 0.0,
    exp_max: float = 10.0,
    exp_base: float = 2.0,
    random_min: float = 0.0,
    random_max: float = 1.0,
) -> float:
    exp_delay = exp_multiplier * (exp_base**attempt)
    delay = max(min(exp_delay, exp_max), exp_min)
    delay += random.uniform(random_min, random_max)
    return delay


def _retry(
    stop_after_attempt: int = 3,
    exp_multiplier: float = 1.0,
    exp_min: float = 0.0,
    exp_max: float = 10.0,
    exp_base: float = 2.0,
    random_min: float = 0.0,
    random_max: float = 1.0,
) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
    def decorator(fn: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
        @functools.wraps(fn)
        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            for attempt in range(stop_after_attempt):
                try:
                    return await fn(*args, **kwargs)
                except Exception:
                    if attempt == stop_after_attempt - 1:
                        raise
                    delay = _calculate_backoff(
                        attempt,
                        exp_multiplier=exp_multiplier,
                        exp_min=exp_min,
                        exp_max=exp_max,
                        exp_base=exp_base,
                        random_min=random_min,
                        random_max=random_max,
                    )
                    logger.warning(
                        "Error in function %s, retrying in %s seconds",
                        fn.__name__,
                        delay,
                        exc_info=True,
                    )
                    await asyncio.sleep(delay)
            raise RuntimeError("Unreachable")

        return wrapper

    return decorator


DEFAULT_PING_TIMEOUT = 5.0
DEFAULT_WAIT_FOR_SERVER_TIMEOUT = 10.0
DEFAULT_CODE_EXECUTION_TIMEOUT = 30.0
DEFAULT_SEND_TOOL_RESULT_TIMEOUT = 10.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 _handle_event(
        self, event: ExecutionEvent, event_handler: ExecutionEventHandler | None
    ) -> None:
        if event_handler is None:
            return
        elif inspect.iscoroutinefunction(event_handler):
            await event_handler(event)
        else:
            event_handler(event)

    async def execute_code(
        self,
        code: str,
        tools: dict[str, ToolHandler],
        timeout: float | None = DEFAULT_CODE_EXECUTION_TIMEOUT,
        headers: dict[str, str] | None = None,
        event_handler: ExecutionEventHandler | None = None,
        send_tool_result_timeout: float | None = DEFAULT_SEND_TOOL_RESULT_TIMEOUT,
    ) -> ExecuteCodeSuccess | ExecuteCodeError:
        merged_headers = {
            **self._headers,
            **(headers or {}),
            "X-Progtc-API-Key": self._api_key,
        }
        async with (
            AsyncHTTPTransport(retries=2) as transport,
            AsyncClient(
                base_url=self._base_url,
                headers=merged_headers,
                timeout=timeout,
                transport=transport,
            ) 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:
                await self._handle_event(
                    ToolCallReceivedEvent(tool_call=tool_call), event_handler
                )
                handler = tools[tool_call.tool_name]
                result = await handler(*tool_call.args, **tool_call.kwargs)
                await self._handle_event(
                    ToolCallCompletedEvent(tool_call=tool_call, result=result),
                    event_handler,
                )

                @_retry()
                async def post_tool_result() -> None:
                    response = await client.post(
                        "/tool-results",
                        json={
                            "execution_id": tool_call.execution_id,
                            "tool_call_id": tool_call.id,
                            "result": result,
                        },
                        timeout=send_tool_result_timeout,
                    )
                    response.raise_for_status()

                await post_tool_result()
                await self._handle_event(
                    ToolCallSentEvent(tool_call=tool_call, result=result), event_handler
                )

            pending_tasks: set[asyncio.Task[None]] = set()
            event_iter = event_source.aiter_sse()
            event_task: asyncio.Task[ServerSentEvent] | None = None

            while True:
                if event_task is None:
                    event_task = asyncio.create_task(anext(event_iter))

                all_tasks: set[asyncio.Task[ServerSentEvent] | asyncio.Task[None]] = {
                    event_task,
                    *pending_tasks,
                }

                done, _ = await asyncio.wait(
                    all_tasks, return_when=asyncio.FIRST_COMPLETED
                )

                # Check for failed tool call tasks
                done_tool_tasks = done & pending_tasks
                pending_tasks -= done_tool_tasks
                for task in done_tool_tasks:
                    if (exc := task.exception()) is not None:
                        event_task.cancel()
                        with contextlib.suppress(asyncio.CancelledError):
                            await event_task
                        raise exc

                # Process event if received
                if event_task in done:
                    try:
                        event = event_task.result()
                    except StopAsyncIteration as err:
                        raise RuntimeError("Stream ended without result") from err
                    finally:
                        event_task = None

                    event_obj = ExecuteCodeMessageAdapter.validate_python(event.json())
                    if isinstance(event_obj, ToolCall):
                        task = asyncio.create_task(handle_tool_call(event_obj))
                        pending_tasks.add(task)
                    else:
                        if pending_tasks:
                            await asyncio.gather(*pending_tasks)
                        return event_obj

    async def execute_code_step(
        self,
        code: str,
        tool_names: list[str],
        state: str | None = None,
        tool_results: list[JsonValue] | None = None,
        timeout: float | None = DEFAULT_CODE_EXECUTION_TIMEOUT,
        headers: dict[str, str] | None = None,
    ) -> ExecuteCodeStepMessage:
        merged_headers = {
            **self._headers,
            **(headers or {}),
            "X-Progtc-API-Key": self._api_key,
        }
        body = {
            "tool_names": tool_names,
            "code": code,
            "state": state,
            "tool_results": tool_results,
        }
        async with AsyncClient(
            base_url=self._base_url,
            headers=merged_headers,
            timeout=timeout,
        ) as client:
            response = await client.post(
                "/execute-code-step",
                json=body,
            )
            response.raise_for_status()
            return ExecuteCodeStepMessageAdapter.validate_python(response.json())
