import datetime
import logging
import pathlib
from typing import Callable, List, Optional

import boto3
import botocore.client
import pydantic

import shovl.schemas as schemas
import shovl.services as services

from .service import ServiceProvider

logger = logging.getLogger("shovl.providers")


class TemporaryCredentials(pydantic.BaseModel):
    access_key_id: str = pydantic.Field(..., alias="AccessKeyId")
    secret_access_key: str = pydantic.Field(..., alias="SecretAccessKey")
    session_token: str = pydantic.Field(..., alias="SessionToken")
    expiration: datetime.datetime = pydantic.Field(..., alias="Expiration")


class BucketProvider(ServiceProvider):
    config: schemas.BucketConfig

    temp_credentials: TemporaryCredentials | None

    session: boto3.Session | None
    client: botocore.client.BaseClient | None

    def __init__(self, config: schemas.BucketConfig) -> None:
        """
        Intialize the bucket provider with the given configuration.
        """
        super().__init__()
        self.config = config

    def _cleanup(self) -> None:
        """
        Cleanup the provider's state.
        """
        self.temp_credentials = None
        self.session = None
        self.client = None

    @property
    def connected(self) -> bool:
        """
        Check if the provider is connected to a bucket.
        """
        return self.client is not None

    def connect(self) -> None:
        """
        Connect to a S3 bucket.
        """
        if self.connected:
            self.disconnect()

        self._create_session()

        if self.config.use_sts:
            self._get_temporary_credentials()

        self._create_s3_client()

        if not self.client:
            raise Exception("Failed to create S3 client.")

        if self.config.test_connection:
            self.client.head_bucket(
                Bucket=services.resolve_credentials(self.config.bucket_name)
            )

    def disconnect(self) -> None:
        """
        Disconnect from a S3 bucket.
        """
        self._cleanup()

    def inspect(
        self, prefix: Optional[str] = None
    ) -> List[schemas.BucketEntity]:
        """
        Inspect the contents of the S3 bucket at the given prefix.
        """
        if not self.client:
            raise Exception("Not connected to a bucket.")

        if prefix and not prefix.endswith("/"):
            raise Exception("Prefix must end with a '/' character.")

        if self.config.head_prefix and (not prefix or len(prefix) == 0):
            prefix = f"{self.config.head_prefix.rstrip('/')}/"

        paginator = self.client.get_paginator("list_objects_v2")

        if prefix:
            logger.debug(
                f"Calling list_objects_v2 with prefix='{prefix}' and "
                f"Delimiter='/'"
            )
            page_iterator = paginator.paginate(
                Bucket=services.resolve_credentials(self.config.bucket_name),
                Prefix=prefix,
                Delimiter="/",
            )
        else:
            logger.debug("Calling list_objects_v2 with Delimiter='/'")
            page_iterator = paginator.paginate(
                Bucket=services.resolve_credentials(self.config.bucket_name),
                Delimiter="/",
            )

        entities = []
        for page in page_iterator:
            # "Folders" are found in CommonPrefixes
            if "CommonPrefixes" in page:
                for obj in page["CommonPrefixes"]:
                    entities.append(
                        schemas.BucketEntity(
                            name=obj["Prefix"],
                            type=schemas.BucketEntityType.folder,
                        )
                    )

            # Top-level files (if any) are found in 'Contents'
            if "Contents" in page:
                for obj in page["Contents"]:
                    if obj["Key"].endswith("/"):
                        continue
                    entities.append(
                        schemas.BucketEntity(
                            name=obj["Key"],
                            size=obj["Size"],
                            last_modified=obj["LastModified"],
                            type=schemas.BucketEntityType.file,
                        )
                    )

        return entities

    def download(
        self,
        key: str,
        target: pathlib.Path,
        recursive: bool = True,
        progress_callback: Callable[[int, int], None] = lambda x, y: None,
    ) -> None:
        """
        Download a file from the S3 bucket with the given key.
        """
        if recursive:
            entities = self.client.list_objects_v2(  # type: ignore
                Bucket=services.resolve_credentials(self.config.bucket_name),
                Prefix=key,
            )
        else:
            entities = self.client.list_objects_v2(  # type: ignore
                Bucket=services.resolve_credentials(self.config.bucket_name),
                Prefix=key,
                Delimiter="/",
            )

        if "Contents" not in entities:
            return

        entities = entities["Contents"]

        files = [
            entity["Key"]
            for entity in entities
            if "Key" in entity and not entity["Key"].endswith("/")
        ]

        parent = str(pathlib.Path(key).parent)
        if parent == ".":
            parent = ""

        for index, file in enumerate(files):
            progress_callback(index + 1, len(files))

            relative_path = file[len(parent) :].lstrip("/")
            local_filename = target / relative_path

            if not local_filename.parent.exists():
                local_filename.parent.mkdir(parents=True, exist_ok=True)

            self.client.download_file(  # type: ignore
                Bucket=services.resolve_credentials(self.config.bucket_name),
                Key=file,
                Filename=str(local_filename),
            )

    def upload(
        self,
        key: str,
        type: schemas.BucketEntityType,
        source: pathlib.Path,
        recursive: bool = True,
        progress_callback: Callable[[int, int], None] = lambda x, y: None,
    ) -> None:
        """
        Upload files to the S3 bucket at the given key.
        """
        if not source.exists():
            raise Exception(f"Source path '{source}' does not exist.")

        if source.is_file():
            entities = [source]
        elif recursive:
            entities = pathlib.Path(source).rglob("*")
        else:
            entities = pathlib.Path(source).glob("*")

        entities = list(entities)

        entities = [entity for entity in entities if entity.is_file()]

        for index, entity in enumerate(list(entities)):
            progress_callback(index + 1, len(entities))

            relative_path = entity.relative_to(source.parent)

            if type == schemas.BucketEntityType.folder:
                s3_key = f"{key.rstrip('/')}/{relative_path.as_posix()}"
            else:
                s3_key = relative_path.as_posix()

            self.client.upload_file(  # type: ignore
                Filename=str(entity),
                Bucket=services.resolve_credentials(self.config.bucket_name),
                Key=s3_key,
            )

    def delete(
        self,
        key: str,
        recursive: bool = True,
        progress_callback: Callable[[int, int], None] = lambda x, y: None,
    ) -> None:
        """
        Delete a file from the S3 bucket with the given key.
        """
        if recursive:
            entities = self.client.list_objects_v2(  # type: ignore
                Bucket=services.resolve_credentials(self.config.bucket_name),
                Prefix=key,
            )
        else:
            entities = self.client.list_objects_v2(  # type: ignore
                Bucket=services.resolve_credentials(self.config.bucket_name),
                Prefix=key,
                Delimiter="/",
            )

        if "Contents" not in entities:
            return

        entities = entities["Contents"]

        files = [
            entity["Key"]
            for entity in entities
            if "Key" in entity and not entity["Key"].endswith("/")
        ]

        for index, file in enumerate(files):
            progress_callback(index + 1, len(files))

            self.client.delete_object(  # type: ignore
                Bucket=self.config.bucket_name,
                Key=file,
            )

    def _create_session(self) -> None:
        """
        Create a boto3 session using the provided configuration.
        """
        logger.info("Creating AWS session...")
        self.session = boto3.Session(
            aws_access_key_id=services.resolve_credentials(
                self.config.aws_access_key_id
            )
            if self.config.aws_access_key_id
            else None,
            aws_secret_access_key=services.resolve_credentials(
                self.config.aws_secret_access_key
            )
            if self.config.aws_secret_access_key
            else None,
            region_name=services.resolve_credentials(self.config.region_name)
            if self.config.region_name
            else None,
            profile_name=services.resolve_credentials(self.config.profile_name)
            if self.config.profile_name
            else None,
            aws_account_id=services.resolve_credentials(
                self.config.aws_account_id
            )
            if self.config.aws_account_id
            else None,
        )
        logger.info("AWS session created successfully.")

    def _create_s3_client(self) -> None:
        """
        Create a boto3 bucket resource using the session.
        """
        logger.info("Creating S3 client...")
        if not self.session:
            raise Exception("AWS session not created.")

        if not self.config.use_sts:
            self.client = self.session.client("s3")
        else:
            if not self.temp_credentials:
                raise Exception("Temporary credentials not obtained.")

            self.client = boto3.client(
                "s3",
                aws_access_key_id=self.temp_credentials.access_key_id,
                aws_secret_access_key=self.temp_credentials.secret_access_key,
                aws_session_token=self.temp_credentials.session_token,
                region_name=services.resolve_credentials(
                    self.config.region_name
                )
                if self.config.region_name
                else None,
            )

        logger.info("S3 client created successfully.")

    def _get_temporary_credentials(self) -> None:
        """
        Get temporary credentials for accessing the S3 bucket.
        """
        logger.debug("Fetching temporary credentials using AWS STS...")
        if not self.session:
            raise Exception("AWS session not created.")

        account_id = (
            services.resolve_credentials(self.config.aws_account_id)
            if self.config.aws_account_id
            else None
        )

        role_name = (
            services.resolve_credentials(self.config.role_name)
            if self.config.role_name
            else None
        )

        sts = self.session.client("sts")
        response = sts.assume_role(
            RoleArn=f"arn:aws:iam::{account_id}:role/{role_name}",
            RoleSessionName="ShovlBucketProviderSession",
        )

        if "Credentials" not in response:
            raise Exception("Failed to assume role for temporary credentials.")

        self.temp_credentials = TemporaryCredentials(**response["Credentials"])
        logger.debug("Temporary credentials obtained successfully.")
