from __future__ import annotations

import atexit
import os
import signal
import sys
import threading
import time
from dataclasses import dataclass, field
from types import TracebackType
from typing import Any, 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,
)

_SIGNALS = (signal.SIGINT, signal.SIGTERM)
_REGISTERED_CLIENTS: set["WyndDataIngestClient"] = set()
_PROCESS_HOOKS_ATTACHED = False
_ORIGINAL_SIGNAL_HANDLERS: dict[signal.Signals, Any] = {}
_ORIGINAL_EXCEPTHOOK = sys.excepthook
_ORIGINAL_THREADING_EXCEPTHOOK = threading.excepthook


class WyndDataIngestClient(_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.Client] = 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.Client(
                limits=limits,
                timeout=self._default_timeout,
                http2=False,
            )
            self._owns_client = True
        self._batch_condition = threading.Condition()
        self._current_batch = _SyncPendingBatch()
        self._ready_batches: list[_SyncPendingBatch] = []
        self._flush_timer: Optional[threading.Timer] = None
        self._in_flight_flushes = 0

        if self._batching_enabled:
            _register_client_for_process_flush(self)

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

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

    def close(self) -> None:
        _unregister_client_for_process_flush(self)
        self._clear_flush_timer()
        self._drain_pending_batches()

        if self._owns_client:
            self._client.close()

    def flush_on_process_lifecycle(self) -> None:
        self._clear_flush_timer()
        self._drain_pending_batches()

    def __enter__(self) -> "WyndDataIngestClient":
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> None:
        self.close()

    def _send_serialized(
        self,
        request_body: bytes,
        headers: Dict[str, str],
        timeout: httpx.Timeout,
    ) -> WyndDataIngestRequestResult:
        attempt = 1
        while True:
            try:
                response = 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,
                    )

                time.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)

                time.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
        )

    def _enqueue_batch(self, lines: list[str]) -> WyndDataIngestRequestResult:
        deferred = _SyncDeferred()

        with self._batch_condition:
            self._current_batch.lines.extend(lines)
            self._current_batch.bytes += _get_lines_size(lines)
            self._current_batch.deferreds.append(deferred)

            if self._should_flush_current_batch_immediately_locked():
                self._clear_flush_timer_locked()
                self._enqueue_current_batch_for_flush_locked()
                self._dispatch_ready_batches_locked()
            else:
                self._schedule_flush_locked()

        return deferred.wait()

    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) -> None:
        if self._flush_timer is not None or not self._current_batch.lines:
            return

        self._flush_timer = threading.Timer(
            self._batching.linger_ms / 1000.0,
            self._flush_timer_callback,
        )
        self._flush_timer.daemon = True
        self._flush_timer.start()

    def _flush_timer_callback(self) -> None:
        with self._batch_condition:
            self._flush_timer = None
            self._enqueue_current_batch_for_flush_locked()
            self._dispatch_ready_batches_locked()

    def _clear_flush_timer(self) -> None:
        with self._batch_condition:
            self._clear_flush_timer_locked()

    def _clear_flush_timer_locked(self) -> None:
        if self._flush_timer is not None:
            self._flush_timer.cancel()
            self._flush_timer = 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 = _SyncPendingBatch()

    def _dispatch_ready_batches_locked(self) -> None:
        while (
            self._in_flight_flushes < self._batching.max_inflight_batches
            and self._ready_batches
        ):
            batch = self._ready_batches.pop(0)
            self._in_flight_flushes += 1
            worker = threading.Thread(
                target=self._flush_batch_worker,
                args=(batch,),
                daemon=True,
            )
            worker.start()

    def _flush_batch_worker(self, batch: _SyncPendingBatch) -> None:
        try:
            request_body, headers, timeout = self._prepare_serialized_batch(batch.lines)
            result = self._send_serialized(request_body, headers, timeout)
        except BaseException as exc:
            for deferred in batch.deferreds:
                deferred.reject(exc)
        else:
            for deferred in batch.deferreds:
                deferred.resolve(result)
        finally:
            with self._batch_condition:
                self._in_flight_flushes -= 1

                if self._ready_batches:
                    self._dispatch_ready_batches_locked()
                elif self._current_batch.lines:
                    if self._should_flush_current_batch_immediately_locked():
                        self._clear_flush_timer_locked()
                        self._enqueue_current_batch_for_flush_locked()
                        self._dispatch_ready_batches_locked()
                    else:
                        self._schedule_flush_locked()

                self._batch_condition.notify_all()

    def _drain_pending_batches(self) -> None:
        with self._batch_condition:
            self._clear_flush_timer_locked()
            self._enqueue_current_batch_for_flush_locked()
            self._dispatch_ready_batches_locked()

            while self._current_batch.lines or self._ready_batches or self._in_flight_flushes > 0:
                self._batch_condition.wait()

    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


def _register_client_for_process_flush(client: WyndDataIngestClient) -> None:
    global _PROCESS_HOOKS_ATTACHED

    _REGISTERED_CLIENTS.add(client)
    if _PROCESS_HOOKS_ATTACHED:
        return

    atexit.register(_handle_atexit)
    sys.excepthook = _handle_uncaught_exception
    threading.excepthook = _handle_thread_exception

    for registered_signal in _SIGNALS:
        _ORIGINAL_SIGNAL_HANDLERS[registered_signal] = signal.getsignal(registered_signal)
        signal.signal(registered_signal, _handle_signal)

    _PROCESS_HOOKS_ATTACHED = True


def _unregister_client_for_process_flush(client: WyndDataIngestClient) -> None:
    _REGISTERED_CLIENTS.discard(client)


def _flush_registered_clients() -> None:
    for client in list(_REGISTERED_CLIENTS):
        try:
            client.flush_on_process_lifecycle()
        except Exception:
            continue


def _handle_atexit() -> None:
    _flush_registered_clients()


def _handle_signal(signum: int, frame: object) -> None:
    _flush_registered_clients()

    previous_handler = _ORIGINAL_SIGNAL_HANDLERS.get(signal.Signals(signum), signal.SIG_DFL)
    if callable(previous_handler) and previous_handler is not _handle_signal:
        previous_handler(signum, frame)
        return

    if previous_handler == signal.SIG_IGN:
        return

    signal.signal(signum, signal.SIG_DFL)
    os.kill(os.getpid(), signum)


def _handle_uncaught_exception(
    exc_type: type[BaseException],
    exc_value: BaseException,
    traceback: TracebackType | None,
) -> None:
    _flush_registered_clients()
    _ORIGINAL_EXCEPTHOOK(exc_type, exc_value, traceback)


def _handle_thread_exception(args: threading.ExceptHookArgs) -> None:
    _flush_registered_clients()
    _ORIGINAL_THREADING_EXCEPTHOOK(args)


@dataclass(slots=True)
class _SyncDeferred:
    _event: threading.Event = field(default_factory=threading.Event)
    _result: Optional[WyndDataIngestRequestResult] = None
    _error: Optional[BaseException] = None

    def resolve(self, value: WyndDataIngestRequestResult) -> None:
        self._result = value
        self._event.set()

    def reject(self, error: BaseException) -> None:
        self._error = error
        self._event.set()

    def wait(self) -> WyndDataIngestRequestResult:
        self._event.wait()
        if self._error is not None:
            raise self._error
        if self._result is None:
            raise RuntimeError("Batch result was not resolved.")
        return self._result


@dataclass(slots=True)
class _SyncPendingBatch:
    lines: list[str] = field(default_factory=list)
    bytes: int = 0
    deferreds: list[_SyncDeferred] = field(default_factory=list)
