import functools
from collections.abc import Callable
from typing import Concatenate, ParamSpec, TypeVar
from uuid import UUID

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError
from src.logger import logger
from src.settings import Settings

settings = Settings()

# 引数と戻り値の型をキャプチャする変数
_P = ParamSpec("_P")
_RT = TypeVar("_RT")


class S3Client:
    ACCESS_TOKEN_DURATION_SECONDS = 1800  # 30 minutes

    def __init__(self, tenant_id: UUID) -> None:
        self.role_arn = settings.s3_role_arn
        self.tenant_id = tenant_id
        self.__set_s3_client()

    @staticmethod
    def _retry_on_expired_token(
        max_retries: int = 5,
    ) -> Callable[
        [Callable[Concatenate["S3Client", _P], _RT]],
        Callable[Concatenate["S3Client", _P], _RT],
    ]:
        """アクセストークンが失効した場合に最大max_retries回まで再試行するデコレータ

        "ExpiredToken" エラー発生時は S3 クライアントを再生成して最大 `max_retries` 回まで再試行します。

        Args:
            max_retries (int, optional): 最大リトライ回数(デフォルトは5)。
        Returns:
            Callable: リトライ処理を付与したデコレータ。
        Raises:
            RuntimeError: 最大リトライ回数を超えた場合。
            ClientError: "ExpiredToken" 以外のエラー、またはリトライ失敗時。

        """

        def decorator(
            func: Callable[Concatenate["S3Client", _P], _RT],
        ) -> Callable[Concatenate["S3Client", _P], _RT]:
            @functools.wraps(func)
            def _wrapper(
                self: "S3Client",
                *args: _P.args,
                **kwargs: _P.kwargs,
            ) -> _RT:
                last_exception = None
                for i in range(max_retries):
                    try:
                        return func(self, *args, **kwargs)
                    except ClientError as e:  # noqa: PERF203 リトライ処理として妥当なため
                        if e.response["Error"]["Code"] == "ExpiredToken":
                            last_exception = e
                            if i < max_retries - 1:
                                logger.warning(f"Token expired. Retrying... (attempt {i + 1}/{max_retries})")
                                self.__set_s3_client()
                        else:
                            raise
                raise RuntimeError("Max retries exceeded due to ExpiredToken.") from last_exception

            return _wrapper

        return decorator

    def __set_s3_client(self) -> None:
        """特定のテナントデータにのみアクセスできる一時的な認証情報で S3 クライアントを初期化します。

        セッションにはテナントIDのタグを付与し、認証情報の有効期間は ACCESS_TOKEN_DURATION_SECONDS で指定します。
        """
        sts = boto3.client("sts")
        assumed = sts.assume_role(
            RoleArn=self.role_arn,
            RoleSessionName=f"s3-session-{self.tenant_id}",
            Tags=[{"Key": "tenant_id", "Value": str(self.tenant_id)}],
            DurationSeconds=self.ACCESS_TOKEN_DURATION_SECONDS,
        )
        self.client = boto3.client(
            "s3",
            config=Config(retries={"max_attempts": 3, "mode": "standard"}),
            aws_access_key_id=assumed["Credentials"]["AccessKeyId"],
            aws_secret_access_key=assumed["Credentials"]["SecretAccessKey"],
            aws_session_token=assumed["Credentials"]["SessionToken"],
        )

    @_retry_on_expired_token()
    def get_object(self, bucket_name: str, key: str) -> bytes:
        response = self.client.get_object(Bucket=bucket_name, Key=key)
        body = response["Body"].read()
        logger.debug(f"File '{key}' retrieved from S3 bucket '{bucket_name}' successfully.")
        return body

    @_retry_on_expired_token()
    def upload_object(self, bucket_name: str, key: str, body: bytes, *, if_none_match: bool) -> None:
        put_object_args = {"Bucket": bucket_name, "Key": key, "Body": body}
        if if_none_match:
            put_object_args["IfNoneMatch"] = "*"
        self.client.put_object(**put_object_args)
        logger.debug(f"File '{key}' uploaded to S3 bucket '{bucket_name}' successfully.")

    @_retry_on_expired_token()
    def delete_object(self, bucket_name: str, key: str) -> None:
        self.client.delete_object(Bucket=bucket_name, Key=key)
        logger.debug(f"File '{key}' deleted from S3 bucket '{bucket_name}' successfully.")
