"""STSUploadSession — core upload functionality using STS credentials."""

from __future__ import annotations

import logging
import os

import boto3
from boto3.s3.transfer import S3Transfer
from botocore.config import Config as BotoConfig

from vcm_file_uploader._types import TransferConfig, UploadResult, UploadSummary
from vcm_file_uploader.plan import UploadPlan
from vcm_file_uploader.transfer import create_transfer_config

logger = logging.getLogger(__name__)

# Default retry configuration for transient S3 errors
DEFAULT_MAX_RETRIES = 5
DEFAULT_RETRY_MODE = "adaptive"


class PrefixError(ValueError):
    """Raised when an S3 key falls outside the allowed prefix."""


class STSUploadSession:
    """Job-bound upload session backed by temporary STS credentials.

    Accepts caller-provided STS credentials and uploads files to a
    scoped S3 prefix using boto3 TransferManager. Configures exponential
    backoff retries for transient S3 failures.
    """

    def __init__(
        self,
        aws_access_key_id: str,
        aws_secret_access_key: str,
        aws_session_token: str,
        region: str,
        bucket: str,
        prefix: str,
        transfer_config: TransferConfig | None = None,
        max_retries: int = DEFAULT_MAX_RETRIES,
        retry_mode: str = DEFAULT_RETRY_MODE,
    ) -> None:
        self.bucket = bucket
        self.prefix = prefix.rstrip("/") + "/"
        self._boto_session = boto3.Session(
            aws_access_key_id=aws_access_key_id,
            aws_secret_access_key=aws_secret_access_key,
            aws_session_token=aws_session_token,
            region_name=region,
        )
        boto_config = BotoConfig(
            retries={"max_attempts": max_retries, "mode": retry_mode},  # type: ignore[arg-type]
        )
        self._s3_client = self._boto_session.client("s3", config=boto_config)
        self._boto_transfer_config = create_transfer_config(transfer_config)
        self._transfer = S3Transfer(self._s3_client, self._boto_transfer_config)  # type: ignore[arg-type]
        self._closed = False

    def _enforce_prefix(self, key: str) -> None:
        """Raise PrefixError if key does not start with the configured prefix."""
        if not key.startswith(self.prefix):
            raise PrefixError(
                f"Key {key!r} is outside allowed prefix {self.prefix!r}"
            )

    def upload_file(
        self,
        local_path: str,
        key: str,
        *,
        extra_args: dict | None = None,
    ) -> UploadResult:
        """Upload a single file to S3.

        The key must start with the configured prefix.
        """
        self._enforce_prefix(key)
        try:
            size = os.path.getsize(local_path)
            self._transfer.upload_file(
                local_path,
                self.bucket,
                key,
                extra_args=extra_args,
            )
            logger.info("Uploaded %s -> s3://%s/%s (%d bytes)", local_path, self.bucket, key, size)
            return UploadResult(local_path=local_path, s3_key=key, size=size, success=True)
        except Exception as exc:
            logger.error("Failed to upload %s: %s", local_path, exc)
            return UploadResult(
                local_path=local_path, s3_key=key, size=0, success=False, error=str(exc)
            )

    def upload_files(self, plan: UploadPlan) -> UploadSummary:
        """Upload all files from an UploadPlan.

        Files are uploaded in descending size order (plan iteration order).
        Continues on individual file failures; the summary reports all outcomes.
        """
        summary = UploadSummary()
        total = len(plan)
        for i, entry in enumerate(plan, 1):
            key = self.prefix + entry.path
            logger.info("Uploading file %d/%d: %s", i, total, entry.path)
            result = self.upload_file(entry.full_path, key)
            summary.results.append(result)

        if summary.failed:
            logger.warning(
                "Upload completed with failures: %d/%d succeeded, %d/%d failed",
                len(summary.succeeded), total, len(summary.failed), total,
            )
            for r in summary.failed:
                logger.warning("  FAILED: %s -> %s: %s", r.local_path, r.s3_key, r.error)
        elif summary.results:
            logger.info("All %d files uploaded successfully (%d bytes)", total, summary.total_bytes)

        return summary

    def close(self) -> None:
        """Clean up resources."""
        if not self._closed:
            self._closed = True

    def __enter__(self) -> STSUploadSession:
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()
