"""S3 export handler — writes records and annotations as JSON arrays to S3."""

from __future__ import annotations

import json
import random
import time
from datetime import UTC, datetime

import structlog
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential_jitter

from caliper._annotation import Annotation
from caliper._config import S3Partition
from caliper._record import Record

logger = structlog.get_logger()

try:
    import boto3
except ImportError:
    boto3 = None  # type: ignore[assignment]

_PERMANENT_S3_ERROR_CODES: frozenset[str] = frozenset(
    {
        "AccessDenied",
        "InvalidAccessKeyId",
        "NoSuchBucket",
        "SignatureDoesNotMatch",
        "AccountProblem",
        "AllAccessDisabled",
        "InvalidBucketName",
        "InvalidSecurity",
    }
)


def _is_retryable_s3(exc: BaseException) -> bool:
    """Return True for transient S3 failures worth retrying."""
    try:
        from botocore.exceptions import BotoCoreError, ClientError
    except ImportError:
        return False

    if isinstance(exc, ClientError):
        error_code = exc.response.get("Error", {}).get("Code", "")
        return error_code not in _PERMANENT_S3_ERROR_CODES
    if isinstance(exc, BotoCoreError):
        return True
    return False


def _classify_s3_error(exc: Exception) -> tuple[str, dict]:
    """Classify an S3 error as transient, permanent, or unknown.

    Returns (classification, context_dict) for structured logging.
    """
    try:
        from botocore.exceptions import BotoCoreError, ClientError
    except ImportError:
        return "unknown", {}

    if isinstance(exc, ClientError):
        error_code = exc.response.get("Error", {}).get("Code", "")
        status_code = exc.response.get("ResponseMetadata", {}).get("HTTPStatusCode")
        context: dict = {"s3_error_code": error_code}
        if status_code is not None:
            context["status_code"] = status_code
        if error_code in _PERMANENT_S3_ERROR_CODES:
            return "permanent", context
        return "transient", context
    if isinstance(exc, BotoCoreError):
        return "transient", {"transport_error": type(exc).__name__}
    return "unknown", {}


def _generate_ulid() -> str:
    """Generate a ULID-like string: 12-char timestamp (ms hex) + 20-char random hex."""
    timestamp_ms = int(time.time() * 1000)
    random_part = random.randbytes(10).hex()
    return f"{timestamp_ms:012x}{random_part}"


class S3Handler:
    """Uploads records and annotations as JSON arrays to S3 with configurable time partitioning."""

    def __init__(
        self,
        *,
        bucket: str,
        prefix: str = "",
        region: str = "us-east-1",
        access_key: str | None = None,
        secret_key: str | None = None,
        endpoint: str | None = None,
        max_retries: int = 3,
        partition: S3Partition = "hour",
    ) -> None:
        if boto3 is None:
            raise ImportError("boto3 is required for target='s3'. Install with: pip install caliper-sdk[s3]")

        self._bucket = bucket
        self._prefix = prefix.strip("/")
        self._partition = partition

        kwargs: dict = {"region_name": region}
        if access_key and secret_key:
            kwargs["aws_access_key_id"] = access_key
            kwargs["aws_secret_access_key"] = secret_key
        if endpoint:
            kwargs["endpoint_url"] = endpoint

        self._max_retries = max_retries
        self._client = boto3.client("s3", **kwargs)

    def _build_key(self, category: str) -> str:
        now = datetime.now(UTC)
        ulid = _generate_ulid()
        filename = f"batch_{ulid}.json"

        format_by_partition = {
            "none": "",
            "year": now.strftime("%Y"),
            "month": now.strftime("%Y/%m"),
            "day": now.strftime("%Y/%m/%d"),
            "hour": now.strftime("%Y/%m/%d/%H"),
        }
        date_path = format_by_partition[self._partition]

        parts = [p for p in [self._prefix, category, date_path, filename] if p]
        return "/".join(parts)

    def _upload_json(self, key: str, payload: list[dict]) -> None:
        body = json.dumps(payload, separators=(",", ":")).encode()
        self._client.put_object(Bucket=self._bucket, Key=key, Body=body, ContentType="application/json")

    def _upload_with_retry(self, key: str, payload: list[dict]) -> None:
        @retry(
            stop=stop_after_attempt(self._max_retries),
            wait=wait_exponential_jitter(initial=0.5, max=10.0, jitter=1.0),
            retry=retry_if_exception(_is_retryable_s3),
            reraise=True,
        )
        def _do_upload() -> None:
            self._upload_json(key, payload)

        _do_upload()

    def export(self, records: list[Record]) -> None:
        key = self._build_key("records")
        payload = [r.to_dict() for r in records]
        try:
            self._upload_with_retry(key, payload)
        except Exception as exc:
            error_class, context = _classify_s3_error(exc)
            log = logger.error if error_class == "permanent" else logger.warning
            log(
                "caliper.export_failed",
                target="s3",
                record_count=len(records),
                bucket=self._bucket,
                error_class=error_class,
                **context,
                exc_info=True,
            )

    def export_annotations(self, annotations: list[Annotation]) -> None:
        key = self._build_key("annotations")
        payload = [a.to_dict() for a in annotations]
        try:
            self._upload_with_retry(key, payload)
        except Exception as exc:
            error_class, context = _classify_s3_error(exc)
            log = logger.error if error_class == "permanent" else logger.warning
            log(
                "caliper.export_annotations_failed",
                target="s3",
                annotation_count=len(annotations),
                bucket=self._bucket,
                error_class=error_class,
                **context,
                exc_info=True,
            )

    def close(self) -> None:
        pass
