from __future__ import annotations

import asyncio
from types import TracebackType
from typing import Dict, Optional, Type

import httpx

from ._shared import (
    _BaseWyndDataIngestClient,
    _is_retryable_error,
    _parse_retry_after_ms,
)
from .types import (
    CompressionOptions,
    RetryOptions,
    SendOptions,
    WyndDataIngestPayload,
    WyndDataIngestPayloadMode,
    WyndDataIngestRequestResult,
    WyndDataIngestTableOperation,
)


class AsyncWyndDataIngestClient(_BaseWyndDataIngestClient):
    def __init__(
        self,
        *,
        url: str,
        mode: WyndDataIngestPayloadMode = "ndjson",
        connections: int = 16,
        table_name: Optional[str] = None,
        table_operation: Optional[WyndDataIngestTableOperation] = None,
        headers: Optional[Dict[str, str]] = None,
        timeout_ms: int = 30_000,
        user_agent: str = "wynd-dataingest/1.0.0",
        retry: Optional[RetryOptions] = None,
        compression: Optional[CompressionOptions] = None,
        client: Optional[httpx.AsyncClient] = None,
    ) -> None:
        super().__init__(
            url=url,
            mode=mode,
            table_name=table_name,
            table_operation=table_operation,
            headers=headers,
            timeout_ms=timeout_ms,
            user_agent=user_agent,
            retry=retry,
            compression=compression,
        )
        if client is not None:
            self._client = client
            self._owns_client = False
        else:
            limits = httpx.Limits(
                max_connections=connections,
                max_keepalive_connections=connections,
                keepalive_expiry=10.0,
            )
            self._client = httpx.AsyncClient(
                limits=limits,
                timeout=self._default_timeout,
                http2=False,
            )
            self._owns_client = True

    async def send(
        self,
        payload: WyndDataIngestPayload,
        options: Optional[SendOptions] = None,
    ) -> WyndDataIngestRequestResult:
        request_body, headers, timeout = self._prepare_send(payload, options)

        attempt = 1
        while True:
            try:
                response = await self._client.post(
                    self._url,
                    content=request_body,
                    headers=headers,
                    timeout=timeout,
                )
                if 200 <= response.status_code < 300:
                    return self._build_result(response, attempt)

                retry_after_ms = _parse_retry_after_ms(response.headers.get("retry-after"))
                if (
                    response.status_code not in self._retry.retryable_status_codes
                    or attempt >= self._retry.max_attempts
                ):
                    self._raise_request_error(
                        attempt=attempt,
                        status_code=response.status_code,
                        response_body=response.text,
                    )

                await asyncio.sleep(self._compute_delay(attempt, retry_after_ms) / 1000.0)
                attempt += 1
            except httpx.HTTPError as exc:
                if not _is_retryable_error(exc) or attempt >= self._retry.max_attempts:
                    self._raise_request_error(attempt=attempt, cause=exc)

                await asyncio.sleep(self._compute_delay(attempt) / 1000.0)
                attempt += 1

    async def aclose(self) -> None:
        if self._owns_client:
            await self._client.aclose()

    async def __aenter__(self) -> "AsyncWyndDataIngestClient":
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> None:
        await self.aclose()
