from __future__ import annotations

import gzip
import json
import random
import time
from collections.abc import Sequence
from email.utils import parsedate_to_datetime
from typing import Dict, Optional, cast

import httpx

from .types import (
    CompressionOptions,
    GzipCompressionOptions,
    RetryOptions,
    SendOptions,
    WyndDataIngestError,
    WyndDataIngestPayload,
    WyndDataIngestPayloadMode,
    WyndDataIngestRequestResult,
    WyndDataIngestTableOperation,
)

DEFAULT_GZIP_COMPRESSION = GzipCompressionOptions(min_bytes=100 * 1_024)


class WyndDataIngestClient:
    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.Client] = None,
    ) -> None:
        self._mode = mode
        self._timeout_ms = timeout_ms
        self._retry = retry or RetryOptions()
        self._compression = compression
        self._default_headers = {
            "content-type": _resolve_content_type(mode),
            "user-agent": user_agent,
            **_resolve_table_headers(table_name, table_operation),
            **_lowercase_headers(headers),
        }

        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,
            )
            timeout = httpx.Timeout(timeout_ms / 1000.0)
            self._client = httpx.Client(limits=limits, timeout=timeout, http2=False)
            self._owns_client = True

        self._url = url

    def send(
        self,
        payload: WyndDataIngestPayload,
        options: Optional[SendOptions] = None,
    ) -> WyndDataIngestRequestResult:
        body = self._serialize_payload(payload)
        options = options or SendOptions()
        headers = {
            **self._default_headers,
            **_lowercase_headers(options.headers),
        }
        if options.content_type:
            headers["content-type"] = options.content_type

        request_body = self._prepare_request_body(body, headers)
        timeout = self._client.timeout if options.timeout_ms is None else httpx.Timeout(options.timeout_ms / 1000.0)

        attempt = 1
        while True:
            try:
                response = self._client.post(
                    self._url,
                    content=request_body,
                    headers=headers,
                    timeout=timeout,
                )
                response_headers = _headers_to_record(response.headers)
                if 200 <= response.status_code < 300:
                    return WyndDataIngestRequestResult(
                        status_code=response.status_code,
                        ok=True,
                        headers=response_headers,
                        attempts=attempt,
                    )

                response_body = response.text
                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
                ):
                    raise WyndDataIngestError(
                        f"WyndDataIngest request failed with status {response.status_code}.",
                        attempts=attempt,
                        status_code=response.status_code,
                        response_body=response_body,
                    )

                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:
                    raise WyndDataIngestError(
                        "WyndDataIngest request failed.",
                        attempts=attempt,
                        cause=exc,
                    ) from exc

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

    def close(self) -> None:
        if self._owns_client:
            self._client.close()

    def _compute_delay(self, attempt: int, retry_after_ms: Optional[int] = None) -> int:
        if retry_after_ms is not None:
            return min(retry_after_ms, self._retry.max_delay_ms)

        exponential = min(
            int(self._retry.initial_delay_ms * (self._retry.factor ** max(0, attempt - 1))),
            self._retry.max_delay_ms,
        )
        jitter_offset = int(exponential * self._retry.jitter_ratio * random.random())
        return exponential + jitter_offset

    def _serialize_payload(self, payload: WyndDataIngestPayload) -> bytes:
        if isinstance(payload, (bytes, bytearray)):
            return bytes(payload)

        if self._mode == "text":
            if isinstance(payload, str):
                return payload.encode("utf-8")

            if isinstance(payload, Sequence) and all(isinstance(value, str) for value in payload):
                return "\n".join(cast(Sequence[str], payload)).encode("utf-8")

            raise TypeError("Text mode requires a string or a sequence of strings.")

        if self._mode == "ndjson":
            if not isinstance(payload, Sequence) or isinstance(payload, (str, bytes, bytearray)):
                raise TypeError("NDJSON mode requires a sequence payload.")

            serialized = "\n".join(json.dumps(item) for item in payload) + "\n"
            return serialized.encode("utf-8")

        return json.dumps(payload).encode("utf-8")

    def _prepare_request_body(self, body: bytes, headers: Dict[str, str]) -> bytes:
        gzip_options = _normalize_gzip_compression(self._compression.gzip if self._compression else None)
        if not gzip_options or "content-encoding" in headers or len(body) < gzip_options.min_bytes:
            return body

        headers["content-encoding"] = "gzip"
        compresslevel = gzip_options.level if gzip_options.level is not None else 9
        return gzip.compress(body, compresslevel=compresslevel)


def _resolve_content_type(mode: WyndDataIngestPayloadMode) -> str:
    if mode == "json":
        return "application/json"
    if mode == "ndjson":
        return "application/x-ndjson"
    if mode == "text":
        return "text/plain; charset=utf-8"
    return "application/octet-stream"


def _normalize_gzip_compression(
    value: bool | GzipCompressionOptions | None,
) -> Optional[GzipCompressionOptions]:
    if value is True:
        return DEFAULT_GZIP_COMPRESSION
    if not value:
        return None
    return GzipCompressionOptions(min_bytes=value.min_bytes, level=value.level)


def _resolve_table_headers(
    table_name: Optional[str],
    table_operation: Optional[WyndDataIngestTableOperation],
) -> Dict[str, str]:
    if not table_name:
        return {}

    if table_operation in (None, "insert", "new"):
        return {"table-name": table_name}
    if table_operation == "update":
        return {"table-name": f"{table_name}-updates"}
    return {"table-name": f"{table_name}-deletes"}


def _lowercase_headers(headers: Optional[Dict[str, str]]) -> Dict[str, str]:
    if not headers:
        return {}
    return {key.lower(): value for key, value in headers.items()}


def _headers_to_record(headers: httpx.Headers) -> Dict[str, str]:
    return dict(headers.items())


def _parse_retry_after_ms(value: Optional[str]) -> Optional[int]:
    if not value:
        return None

    try:
        return max(0, int(float(value) * 1000))
    except ValueError:
        pass

    try:
        target = parsedate_to_datetime(value).timestamp()
        return max(0, int((target - time.time()) * 1000))
    except Exception:
        return None


def _is_retryable_error(exc: httpx.HTTPError) -> bool:
    return isinstance(
        exc,
        (
            httpx.ConnectError,
            httpx.ReadError,
            httpx.ReadTimeout,
            httpx.WriteError,
            httpx.WriteTimeout,
            httpx.PoolTimeout,
            httpx.RemoteProtocolError,
        ),
    )
