from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from types import TracebackType
from typing import Dict, Optional, Type

import httpx

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


class AsyncWyndDataIngestClient(_BaseWyndDataIngestClient):
    def __init__(
        self,
        *,
        url: str,
        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,
        batching: Optional[BatchingOptions] = None,
        client: Optional[httpx.AsyncClient] = None,
    ) -> None:
        super().__init__(
            url=url,
            table_name=table_name,
            table_operation=table_operation,
            headers=headers,
            timeout_ms=timeout_ms,
            user_agent=user_agent,
            retry=retry,
            compression=compression,
            batching=batching,
        )
        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
        self._current_batch = _AsyncPendingBatch()
        self._ready_batches: list[_AsyncPendingBatch] = []
        self._batch_lock = asyncio.Lock()
        self._flush_handle: Optional[asyncio.TimerHandle] = None
        self._in_flight_flushes: set[asyncio.Task[None]] = set()

    async def send(
        self,
        payload: WyndDataIngestPayload,
        options: Optional[SendOptions] = None,
    ) -> WyndDataIngestRequestResult:
        if not self._can_batch(payload, options):
            await self._drain_pending_batches()
            request_body, headers, timeout = self._prepare_send(payload, options)
            return await self._send_serialized(request_body, headers, timeout)

        lines = self._serialize_payload_lines(payload)
        return await self._enqueue_batch(lines)

    async def aclose(self) -> None:
        self._clear_flush_handle()
        await self._drain_pending_batches()

        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()

    async def _send_serialized(
        self,
        request_body: bytes,
        headers: Dict[str, str],
        timeout: httpx.Timeout,
    ) -> WyndDataIngestRequestResult:
        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

    def _can_batch(
        self,
        payload: WyndDataIngestPayload,
        options: Optional[SendOptions],
    ) -> bool:
        if not self._batching_enabled or isinstance(payload, (bytes, bytearray)):
            return False

        if options is None:
            return True

        return (
            options.headers is None
            and options.content_type is None
            and options.timeout_ms is None
        )

    async def _enqueue_batch(self, lines: list[str]) -> WyndDataIngestRequestResult:
        loop = asyncio.get_running_loop()
        future: asyncio.Future[WyndDataIngestRequestResult] = loop.create_future()

        async with self._batch_lock:
            self._current_batch.lines.extend(lines)
            self._current_batch.bytes += _get_lines_size(lines)
            self._current_batch.futures.append(future)

            if self._should_flush_current_batch_immediately_locked():
                self._cancel_flush_handle_locked()
                self._enqueue_current_batch_for_flush_locked()
                self._dispatch_ready_batches_locked(loop)
            else:
                self._schedule_flush_locked(loop)

        return await future

    def _should_flush_current_batch_immediately_locked(self) -> bool:
        return (
            len(self._current_batch.lines) >= self._batching.max_records
            or self._current_batch.bytes >= self._batching.max_bytes
            or self._batching.linger_ms <= 0
        )

    def _schedule_flush_locked(self, loop: asyncio.AbstractEventLoop) -> None:
        if self._flush_handle is not None or not self._current_batch.lines:
            return

        self._flush_handle = loop.call_later(
            self._batching.linger_ms / 1000.0,
            self._schedule_flush_task_from_callback,
        )

    def _schedule_flush_task_from_callback(self) -> None:
        self._flush_handle = None
        loop = asyncio.get_running_loop()
        self._enqueue_current_batch_for_flush_locked()
        self._dispatch_ready_batches_locked(loop)

    def _clear_flush_handle(self) -> None:
        if self._flush_handle is not None:
            self._flush_handle.cancel()
            self._flush_handle = None

    def _cancel_flush_handle_locked(self) -> None:
        if self._flush_handle is not None:
            self._flush_handle.cancel()
            self._flush_handle = None

    def _enqueue_current_batch_for_flush_locked(self) -> None:
        if not self._current_batch.lines:
            return

        self._ready_batches.append(self._current_batch)
        self._current_batch = _AsyncPendingBatch()

    def _dispatch_ready_batches_locked(self, loop: asyncio.AbstractEventLoop) -> None:
        while (
            len(self._in_flight_flushes) < self._batching.max_inflight_batches
            and self._ready_batches
        ):
            batch = self._ready_batches.pop(0)
            task = loop.create_task(self._flush_batch_task(batch))
            self._in_flight_flushes.add(task)
            task.add_done_callback(self._handle_flush_task_done)

    async def _flush_batch_task(self, batch: _AsyncPendingBatch) -> None:
        try:
            request_body, headers, timeout = self._prepare_serialized_batch(batch.lines)
            result = await self._send_serialized(request_body, headers, timeout)
        except BaseException as exc:
            for future in batch.futures:
                if not future.done():
                    future.set_exception(exc)
        else:
            for future in batch.futures:
                if not future.done():
                    future.set_result(result)

    def _handle_flush_task_done(self, task: asyncio.Task[None]) -> None:
        self._in_flight_flushes.discard(task)
        asyncio.get_running_loop().create_task(self._after_flush_completion())

    async def _after_flush_completion(self) -> None:
        async with self._batch_lock:
            loop = asyncio.get_running_loop()
            self._dispatch_ready_batches_locked(loop)

            if self._ready_batches:
                return

            if not self._current_batch.lines:
                return

            if self._should_flush_current_batch_immediately_locked():
                self._cancel_flush_handle_locked()
                self._enqueue_current_batch_for_flush_locked()
                self._dispatch_ready_batches_locked(loop)
                return

            self._schedule_flush_locked(loop)

    async def _drain_pending_batches(self) -> None:
        async with self._batch_lock:
            loop = asyncio.get_running_loop()
            self._cancel_flush_handle_locked()
            self._enqueue_current_batch_for_flush_locked()
            self._dispatch_ready_batches_locked(loop)

        while True:
            tasks = list(self._in_flight_flushes)
            if not tasks:
                async with self._batch_lock:
                    if not self._current_batch.lines and not self._ready_batches:
                        return

                    loop = asyncio.get_running_loop()
                    self._enqueue_current_batch_for_flush_locked()
                    self._dispatch_ready_batches_locked(loop)
                    tasks = list(self._in_flight_flushes)
                    if not tasks and not self._current_batch.lines and not self._ready_batches:
                        return

            await asyncio.gather(*tasks, return_exceptions=True)

    def _prepare_serialized_batch(
        self,
        lines: list[str],
    ) -> tuple[bytes, Dict[str, str], httpx.Timeout]:
        headers = dict(self._default_headers)
        body = _format_ndjson_lines(lines)
        request_body = self._prepare_request_body(body, headers)
        return request_body, headers, self._default_timeout


@dataclass(slots=True)
class _AsyncPendingBatch:
    lines: list[str] = field(default_factory=list)
    bytes: int = 0
    futures: list[asyncio.Future[WyndDataIngestRequestResult]] = field(default_factory=list)
