import logging
from io import BytesIO
from urllib.parse import urlparse

import boto3
from botocore.exceptions import ClientError

from ..utils import IOPathWrapper
from .base_backend import BaseBackend
from .progress import TqdmCallback

logger = logging.getLogger(__name__)


class S3Backend(BaseBackend):
    def __init__(self, bucket: str, prefix: str, session: boto3.Session | None = None) -> None:
        # Use default boto3 session if none is provided
        if not session:
            logger.debug("No boto3 session provided, creating default session.")
            session = boto3.Session()

        # Create a new s3 client from the given session
        self.__client = session.client("s3")

        # Write the bucket and prefix to the private attributes
        self.__bucket = bucket
        self.__prefix = prefix
        logger.debug("S3Backend(bucket=%s, prefix=%s) initialized.", bucket, prefix)

    @property
    def remote_source(self) -> bool:
        # This backend is always remote
        return True

    @property
    def scheme(self) -> str:
        return "s3"

    @property
    def bucket(self) -> str:
        return self.__bucket

    @property
    def prefix(self) -> str:
        return self.__prefix

    def read_file(self, file_path: str) -> IOPathWrapper:
        """Read a file from S3.

        Args:
            file_path (str): Path to file, either full S3 URI or key within bucket

        Returns:
            IOPathWrapper: File contents

        Raises:
            FileNotFoundError: If file doesn't exist
        """
        bucket, key = self._parse_input(file_path)
        data = BytesIO()
        with TqdmCallback(total=self.get_file_size(file_path), desc=f"Downloading {key}") as pbar:
            self.__client.download_fileobj(bucket, key, data, Callback=pbar.callback)
        return IOPathWrapper(data)

    def write_file(self, data: IOPathWrapper, file_path: str) -> None:
        """Write file to S3.

        Args:
            data (IOPathWrapper): File data to write
            file_path (str): Destination path
        """
        bucket, key = self._parse_input(file_path)

        # Reset file object position
        data.file_obj.seek(0)

        # Upload to S3 (Always show progress)
        try:
            with TqdmCallback(total=data.file_obj.getbuffer().nbytes, desc=f"Uploading {key}") as pbar:
                self.__client.upload_fileobj(data.file_obj, bucket, key, Callback=pbar.callback)
        except ClientError as e:
            raise RuntimeError(f"Failed to upload file to S3: {e}") from e

    def exists(self, file_path: str) -> bool:
        """Check if a file exists in S3.

        Args:
            file_path (str): Path to check

        Returns:
            bool: True if file exists
        """
        bucket, key = self._parse_input(file_path)

        try:
            self.__client.head_object(Bucket=bucket, Key=key)
            return True
        except ClientError as e:
            if self._is_not_found_error(e):
                return False
            raise

    def close(self) -> None:
        """Close connections.

        For S3, we need to close the underlying boto3 client.
        """
        self.__client.close()

    def get_file_list(self, extensions: tuple[str, ...] | None = None, num_files: int | None = None) -> list[str]:
        """Get list of files from S3 with optional filtering.

        Args:
            extensions (tuple[str, ...], optional): Filter by file extensions
            num_files (int, optional): Limit number of files returned

        Returns:
            list[str]: List of file paths
        """
        # List all objects with the given prefix
        paginator = self.__client.get_paginator("list_objects_v2")

        files = []
        for page in paginator.paginate(Bucket=self.bucket, Prefix=self.prefix):
            if "Contents" in page:
                for obj in page["Contents"]:
                    key = obj["Key"]
                    # Skip directories (keys ending with /)
                    if key.endswith("/"):
                        continue

                    # Add full S3 path
                    files.append(self._to_url(self.bucket, key))

        # Filter by extension if provided
        if extensions:
            files = [f for f in files if any(f.lower().endswith(ext.lower()) for ext in extensions)]

        # Sort for deterministic behavior
        files.sort()

        # Limit results if specified
        if num_files is not None:
            files = files[:num_files]

        return files

    def get_file_size(self, file_path: str) -> int:
        """Return the size of the file on s3 bucket in bytes."""
        bucket, key = self._parse_input(file_path)

        try:
            response = self.__client.head_object(Bucket=bucket, Key=key)
            return response["ContentLength"]
        except ClientError as e:
            if self._is_not_found_error(e):
                raise FileNotFoundError(f"File not found: {file_path}") from e
            raise

    def get_file_identity(self, file_path: str) -> str:
        """Return object identity (ETag) for a file on S3.

        Args:
            file_path (str): Path to file on S3.

        Returns:
            str: ETag value.

        Raises:
            FileNotFoundError: If file doesn't exist.
        """
        bucket, key = self._parse_input(file_path)

        try:
            response = self.__client.head_object(Bucket=bucket, Key=key)
            etag = response.get("ETag")
            if etag is None:
                raise RuntimeError(f"ETag unavailable for file: {file_path}")
            return etag
        except ClientError as e:
            if self._is_not_found_error(e):
                raise FileNotFoundError(f"File not found: {file_path}") from e
            raise

    def download_file(self, source_path: str, destination_path: str) -> None:
        """Download a file from S3 to local filesystem.

        Args:
            source_path (str): Source S3 path
            destination_path (str): Destination local path
        """
        bucket, key = self._parse_input(source_path)

        # Download the file
        with TqdmCallback(total=self.get_file_size(source_path), desc=f"Downloading {key}") as progress:
            self.__client.download_file(bucket, key, destination_path, Callback=progress.callback)

    def _parse_input(self, file_path: str) -> tuple[str, str]:
        """Convert S3 URI to (bucket, key) tuple.

        Args:
            file_path: Either "s3://bucket/key" or just "key"

        Returns:
            tuple[str, str]: (bucket, key)
        """
        parsed = urlparse(file_path)
        if parsed.scheme == "s3":
            # s3://bucket/key/path -> bucket="bucket", key="key/path"
            bucket = parsed.netloc
            key = parsed.path.lstrip("/")  # Remove leading slash
            return bucket, key

        # Not a URI - treat as absolute key within default bucket
        # No prefix prepending - user must provide full key
        return self.bucket, file_path

    def _to_url(self, bucket: str, key: str) -> str:
        """Convert (bucket, key) to S3 URI.

        Args:
            bucket: S3 bucket name
            key: Object key

        Returns:
            str: S3 URI like "s3://bucket/key"
        """
        return f"s3://{bucket}/{key}"

    def _is_not_found_error(self, e: ClientError) -> bool:
        """Check if ClientError indicates file not found.

        AWS S3 returns different error codes for "not found":
        - '404' - from head_object() operation
        - 'NoSuchKey' - from get_object() and other operations
        - 'NoSuchBucket' - bucket doesn't exist
        - '403' - can indicate missing file if user lacks s3:ListBucket permission

        Args:
            e: ClientError from boto3

        Returns:
            bool: True if this is a not-found error
        """
        error_code = e.response.get("Error", {}).get("Code", "")
        return error_code in ("404", "403", "NoSuchKey", "NoSuchBucket")
