"""
Async S3-compatible object storage client using aioboto3.

Features:
- Works with AWS S3, Ceph RGW, MinIO, and other S3-compatible services
- Persistent client connection (safe for high concurrency)
- Same API as sync S3Client, all methods async

Usage:
    from investify_utils.s3 import AsyncS3Client

    client = AsyncS3Client(
        endpoint_url="https://s3.example.com",
        access_key="access_key",
        secret_key="secret_key",
    )

    # Get object as bytes
    data = await client.get_object("my-bucket", "remote.pdf")

    # Put object from bytes/string
    await client.put_object("my-bucket", "file.txt", b"content", content_type="text/plain")

    # Close when done (or use as async context manager)
    await client.close()
"""

import asyncio
import os
from typing import IO

import aioboto3
from botocore.config import Config
from botocore.exceptions import ClientError


class AsyncS3Client:
    """
    Async S3-compatible object storage client.

    Uses a persistent client connection — safe for concurrent use from
    multiple asyncio tasks. The client is lazily created on first use
    and reused for all subsequent operations.

    Args:
        endpoint_url: S3 endpoint URL (e.g., https://s3.amazonaws.com)
        access_key: AWS access key ID
        secret_key: AWS secret access key
        region: AWS region (default: None)
    """

    def __init__(
        self,
        endpoint_url: str,
        access_key: str,
        secret_key: str,
        region: str | None = None,
    ):
        self._endpoint_url = endpoint_url
        self._access_key = access_key
        self._secret_key = secret_key
        self._region = region
        self._client = None
        self._client_ctx = None
        self._lock = asyncio.Lock()

    async def _get_client(self):
        """Get or create the persistent S3 client."""
        if self._client is None:
            async with self._lock:
                if self._client is None:
                    session = aioboto3.Session(
                        aws_access_key_id=self._access_key,
                        aws_secret_access_key=self._secret_key,
                        region_name=self._region,
                    )
                    self._client_ctx = session.client(
                        "s3",
                        endpoint_url=self._endpoint_url,
                        config=Config(signature_version="s3v4"),
                    )
                    self._client = await self._client_ctx.__aenter__()
        return self._client

    async def close(self) -> None:
        """Close the persistent client connection."""
        if self._client_ctx is not None:
            await self._client_ctx.__aexit__(None, None, None)
            self._client = None
            self._client_ctx = None

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        await self.close()

    async def list_buckets(self) -> list[str]:
        """List all buckets."""
        client = await self._get_client()
        response = await client.list_buckets()
        return [bucket["Name"] for bucket in response["Buckets"]]

    async def list_objects(
        self,
        bucket: str,
        prefix: str = "",
        max_keys: int | None = None,
    ) -> list[dict]:
        """
        List objects in a bucket with optional prefix filter.

        Args:
            bucket: Bucket name
            prefix: Filter objects by prefix (e.g., "folder/")
            max_keys: Maximum number of objects to return (None = all)

        Returns:
            List of object metadata dicts with keys: Key, Size, LastModified
        """
        client = await self._get_client()
        objects = []
        paginator = client.get_paginator("list_objects_v2")
        async for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            for obj in page.get("Contents", []):
                objects.append(
                    {
                        "Key": obj["Key"],
                        "Size": obj["Size"],
                        "LastModified": obj["LastModified"],
                    }
                )
                if max_keys and len(objects) >= max_keys:
                    return objects
        return objects

    async def upload_file(self, file_path: str, bucket: str, key: str | None = None) -> None:
        """
        Upload a local file to S3.

        Args:
            file_path: Local file path
            bucket: Bucket name
            key: Object key (default: basename of file_path)
        """
        if key is None:
            key = os.path.basename(file_path)
        client = await self._get_client()
        await client.upload_file(file_path, bucket, key)

    async def download_file(self, bucket: str, key: str, file_path: str) -> None:
        """
        Download an object to a local file.

        Args:
            bucket: Bucket name
            key: Object key
            file_path: Local file path to save to
        """
        client = await self._get_client()
        await client.download_file(bucket, key, file_path)

    async def get_object(self, bucket: str, key: str) -> bytes:
        """
        Get object content as bytes.

        Args:
            bucket: Bucket name
            key: Object key

        Returns:
            Object content as bytes
        """
        client = await self._get_client()
        response = await client.get_object(Bucket=bucket, Key=key)
        return await response["Body"].read()

    async def put_object(
        self,
        bucket: str,
        key: str,
        data: str | bytes | IO[bytes],
        content_type: str | None = None,
        content_disposition: str | None = None,
    ) -> None:
        """
        Upload data directly to S3.

        Args:
            bucket: Bucket name
            key: Object key
            data: Content as string, bytes, or file-like object
            content_type: MIME type (e.g., "application/pdf")
            content_disposition: Content-Disposition header value
        """
        params: dict = {"Bucket": bucket, "Key": key, "Body": data}
        if content_type:
            params["ContentType"] = content_type
        if content_disposition:
            params["ContentDisposition"] = content_disposition
        client = await self._get_client()
        await client.put_object(**params)

    async def delete_object(self, bucket: str, key: str) -> None:
        """Delete a single object."""
        client = await self._get_client()
        await client.delete_object(Bucket=bucket, Key=key)

    async def delete_prefix(self, bucket: str, prefix: str) -> int:
        """
        Delete all objects with a given prefix.

        Args:
            bucket: Bucket name
            prefix: Prefix to match (e.g., "folder/" deletes all in folder)

        Returns:
            Number of objects deleted
        """
        client = await self._get_client()
        deleted_count = 0
        paginator = client.get_paginator("list_objects_v2")
        async for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            contents = page.get("Contents", [])
            if not contents:
                continue
            delete_keys = [{"Key": obj["Key"]} for obj in contents]
            await client.delete_objects(Bucket=bucket, Delete={"Objects": delete_keys})
            deleted_count += len(delete_keys)
        return deleted_count

    async def exists(self, bucket: str, key: str) -> bool:
        """
        Check if an object exists.

        Args:
            bucket: Bucket name
            key: Object key

        Returns:
            True if object exists, False otherwise
        """
        try:
            client = await self._get_client()
            await client.head_object(Bucket=bucket, Key=key)
            return True
        except ClientError as e:
            if e.response["Error"]["Code"] == "404":
                return False
            raise
