"""
An ADK ArtifactService implementation using Amazon S3 compatible storage.
"""

import asyncio
import logging
import unicodedata

import boto3
from botocore.client import BaseClient
from botocore.exceptions import BotoCoreError, ClientError, NoCredentialsError
from google.adk.artifacts import BaseArtifactService
from google.genai import types as adk_types
from typing_extensions import override

logger = logging.getLogger(__name__)


class S3ArtifactService(BaseArtifactService):
    """
    An artifact service implementation using Amazon S3 compatible storage.

    Stores artifacts in an S3-compatible bucket with a structured key format:
    {app_name}/{user_id}/{session_id_or_user}/{filename}/{version}

    Supports AWS S3 and S3-compatible APIs like MinIO.

    Required S3 Permissions:
    The IAM user or role must have the following minimum permissions for the specific bucket:
    - s3:GetObject: Read artifacts from the bucket
    - s3:PutObject: Store new artifacts to the bucket
    - s3:DeleteObject: Delete artifacts from the bucket

    Example IAM Policy (replace 'your-bucket-name' with actual bucket):
    {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": [
                    "s3:GetObject",
                    "s3:PutObject",
                    "s3:DeleteObject"
                ],
                "Resource": "arn:aws:s3:::your-bucket-name/*"
            },
            {
                "Effect": "Allow",
                "Action": "s3:ListBucket",
                "Resource": "arn:aws:s3:::your-bucket-name"
            }
        ]
    }
    """

    def __init__(
        self,
        bucket_name: str,
        s3_client: BaseClient | None = None,
        **kwargs,
    ):
        """
        Args:
            bucket_name: The name of the S3 bucket to use.
            s3_client: Optional pre-configured S3 client. If None, creates a new client.
            **kwargs: Optional parameters for boto3 client configuration.

        Raises:
            ValueError: If bucket_name is not provided.
            NoCredentialsError: If AWS credentials are not available.
        """
        if not bucket_name:
            raise ValueError("bucket_name cannot be empty for S3ArtifactService")

        self.bucket_name = bucket_name

        if s3_client is None:
            try:
                self.s3 = boto3.client("s3", **kwargs)
            except NoCredentialsError as e:
                logger.error("AWS credentials not found. Please configure credentials.")
                raise ValueError(
                    "AWS credentials not found. Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables or configure AWS credentials."
                ) from e
        else:
            self.s3 = s3_client

        try:
            self.s3.head_bucket(Bucket=self.bucket_name)
            logger.info(
                "S3ArtifactService initialized successfully. Bucket: %s",
                self.bucket_name,
            )
        except ClientError as e:
            error_code = e.response.get("Error", {}).get("Code", "Unknown")
            if error_code == "404":
                logger.error("S3 bucket '%s' does not exist", self.bucket_name)
                raise ValueError(
                    f"S3 bucket '{self.bucket_name}' does not exist"
                ) from e
            elif error_code == "403":
                logger.error("Access denied to S3 bucket '%s'", self.bucket_name)
                raise ValueError(
                    f"Access denied to S3 bucket '{self.bucket_name}'"
                ) from e
            else:
                logger.error("Failed to access S3 bucket '%s': %s", self.bucket_name, e)
                raise ValueError(
                    f"Failed to access S3 bucket '{self.bucket_name}': {e}"
                ) from e

    def _file_has_user_namespace(self, filename: str) -> bool:
        return filename.startswith("user:")

    def _get_object_key(
        self,
        app_name: str,
        user_id: str,
        session_id: str,
        filename: str,
        version: int | str,
    ) -> str:
        """Constructs the S3 object key for an artifact."""
        filename = self._normalize_filename_unicode(filename)

        if self._file_has_user_namespace(filename):
            filename_clean = filename.split(":", 1)[1]
            return f"{app_name}/{user_id}/user/{filename_clean}/{version}"
        return f"{app_name}/{user_id}/{session_id}/{filename}/{version}"

    def _normalize_filename_unicode(self, filename: str) -> str:
        """Normalizes Unicode characters in a filename to their standard form."""
        return unicodedata.normalize("NFKC", filename)

    @override
    async def save_artifact(
        self,
        *,
        app_name: str,
        user_id: str,
        session_id: str,
        filename: str,
        artifact: adk_types.Part,
    ) -> int:
        log_prefix = f"[S3Artifact:Save:{filename}] "

        if not artifact.inline_data or artifact.inline_data.data is None:
            raise ValueError("Artifact Part has no inline_data to save.")

        filename = self._normalize_filename_unicode(filename)

        # Get existing versions to determine next version number
        versions = await self.list_versions(
            app_name=app_name,
            user_id=user_id,
            session_id=session_id,
            filename=filename,
        )
        version = 0 if not versions else max(versions) + 1

        object_key = self._get_object_key(
            app_name, user_id, session_id, filename, version
        )

        try:

            def _put_object():
                return self.s3.put_object(
                    Bucket=self.bucket_name,
                    Key=object_key,
                    Body=artifact.inline_data.data,
                    ContentType=artifact.inline_data.mime_type,
                    Metadata={
                        "original_filename": filename,
                        "user_id": user_id,
                        "session_id": session_id,
                        "version": str(version),
                    },
                )

            await asyncio.to_thread(_put_object)

            logger.info(
                "%sSaved artifact '%s' version %d successfully to S3 key: %s",
                log_prefix,
                filename,
                version,
                object_key,
            )
            return version

        except ClientError as e:
            logger.error(
                "%sFailed to save artifact '%s' version %d to S3: %s",
                log_prefix,
                filename,
                version,
                e,
            )
            raise OSError(
                f"Failed to save artifact version {version} to S3: {e}"
            ) from e
        except BotoCoreError as e:
            logger.error(
                "%sBotoCore error saving artifact '%s' version %d: %s",
                log_prefix,
                filename,
                version,
                e,
            )
            raise OSError(
                f"BotoCore error saving artifact version {version}: {e}"
            ) from e

    @override
    async def load_artifact(
        self,
        *,
        app_name: str,
        user_id: str,
        session_id: str,
        filename: str,
        version: int | None = None,
    ) -> adk_types.Part | None:
        log_prefix = f"[S3Artifact:Load:{filename}] "
        filename = self._normalize_filename_unicode(filename)

        load_version = version
        if load_version is None:
            versions = await self.list_versions(
                app_name=app_name,
                user_id=user_id,
                session_id=session_id,
                filename=filename,
            )
            if not versions:
                logger.debug("%sNo versions found for artifact.", log_prefix)
                return None
            load_version = max(versions)
            logger.debug("%sLoading latest version: %d", log_prefix, load_version)
        else:
            logger.debug("%sLoading specified version: %d", log_prefix, load_version)

        object_key = self._get_object_key(
            app_name, user_id, session_id, filename, load_version
        )

        try:

            def _get_object():
                return self.s3.get_object(Bucket=self.bucket_name, Key=object_key)

            response = await asyncio.to_thread(_get_object)
            data = response["Body"].read()
            mime_type = response.get("ContentType", "application/octet-stream")

            artifact_part = adk_types.Part.from_bytes(data=data, mime_type=mime_type)

            logger.info(
                "%sLoaded artifact '%s' version %d successfully (%d bytes, %s)",
                log_prefix,
                filename,
                load_version,
                len(data),
                mime_type,
            )
            return artifact_part

        except ClientError as e:
            error_code = e.response.get("Error", {}).get("Code", "Unknown")
            if error_code == "NoSuchKey":
                logger.debug("%sArtifact not found: %s", log_prefix, object_key)
                return None
            else:
                logger.error(
                    "%sFailed to load artifact '%s' version %d from S3: %s",
                    log_prefix,
                    filename,
                    load_version,
                    e,
                )
                return None
        except BotoCoreError as e:
            logger.error(
                "%sBotoCore error loading artifact '%s' version %d: %s",
                log_prefix,
                filename,
                load_version,
                e,
            )
            return None

    @override
    async def list_artifact_keys(
        self, *, app_name: str, user_id: str, session_id: str
    ) -> list[str]:
        log_prefix = "[S3Artifact:ListKeys] "
        filenames = set()

        session_prefix = f"{app_name}/{user_id}/{session_id}/"
        try:

            def _list_session_objects():
                paginator = self.s3.get_paginator("list_objects_v2")
                return paginator.paginate(
                    Bucket=self.bucket_name, Prefix=session_prefix
                )

            session_pages = await asyncio.to_thread(_list_session_objects)
            for page in session_pages:
                for obj in page.get("Contents", []):
                    parts = obj["Key"].split("/")
                    if len(parts) >= 5:  # scope/user/session/filename/version
                        filename = parts[3]
                        filenames.add(filename)
        except ClientError as e:
            logger.warning(
                "%sError listing session objects with prefix '%s': %s",
                log_prefix,
                session_prefix,
                e,
            )

        user_prefix = f"{app_name}/{user_id}/user/"
        try:

            def _list_user_objects():
                paginator = self.s3.get_paginator("list_objects_v2")
                return paginator.paginate(Bucket=self.bucket_name, Prefix=user_prefix)

            user_pages = await asyncio.to_thread(_list_user_objects)
            for page in user_pages:
                for obj in page.get("Contents", []):
                    parts = obj["Key"].split("/")
                    if len(parts) >= 5:  # scope/user/user/filename/version
                        filename = parts[3]
                        filenames.add(f"user:{filename}")
        except ClientError as e:
            logger.warning(
                "%sError listing user objects with prefix '%s': %s",
                log_prefix,
                user_prefix,
                e,
            )

        sorted_filenames = sorted(list(filenames))
        logger.debug("%sFound %d artifact keys.", log_prefix, len(sorted_filenames))
        return sorted_filenames

    @override
    async def delete_artifact(
        self, *, app_name: str, user_id: str, session_id: str, filename: str
    ) -> None:
        log_prefix = f"[S3Artifact:Delete:{filename}] "
        filename = self._normalize_filename_unicode(filename)

        # Get all versions to delete
        versions = await self.list_versions(
            app_name=app_name,
            user_id=user_id,
            session_id=session_id,
            filename=filename,
        )

        if not versions:
            logger.debug("%sNo versions found to delete for artifact.", log_prefix)
            return

        # Delete all versions
        for version in versions:
            object_key = self._get_object_key(
                app_name, user_id, session_id, filename, version
            )
            try:

                def _delete_object():
                    return self.s3.delete_object(
                        Bucket=self.bucket_name, Key=object_key
                    )

                await asyncio.to_thread(_delete_object)
                logger.debug(
                    "%sDeleted version %d: %s", log_prefix, version, object_key
                )
            except ClientError as e:
                logger.warning(
                    "%sFailed to delete version %d (%s): %s",
                    log_prefix,
                    version,
                    object_key,
                    e,
                )

        logger.info(
            "%sDeleted artifact '%s' (%d versions)",
            log_prefix,
            filename,
            len(versions),
        )

    @override
    async def list_versions(
        self, *, app_name: str, user_id: str, session_id: str, filename: str
    ) -> list[int]:
        log_prefix = f"[S3Artifact:ListVersions:{filename}] "
        filename = self._normalize_filename_unicode(filename)

        # Get the prefix for this specific artifact (without version)
        prefix = self._get_object_key(app_name, user_id, session_id, filename, "")
        versions = []

        try:

            def _list_objects():
                paginator = self.s3.get_paginator("list_objects_v2")
                return paginator.paginate(Bucket=self.bucket_name, Prefix=prefix)

            pages = await asyncio.to_thread(_list_objects)
            for page in pages:
                for obj in page.get("Contents", []):
                    parts = obj["Key"].split("/")
                    if len(parts) >= 5:  # scope/user/session_or_user/filename/version
                        try:
                            version = int(parts[4])
                            versions.append(version)
                        except ValueError:
                            continue  # Skip non-integer versions

        except ClientError as e:
            logger.error(
                "%sError listing versions with prefix '%s': %s",
                log_prefix,
                prefix,
                e,
            )
            return []

        sorted_versions = sorted(versions)
        logger.debug("%sFound versions: %s", log_prefix, sorted_versions)
        return sorted_versions
