from pathlib import Path
from typing import Optional

from phidata.asset.file import File
from phidata.asset import DataAsset, DataAssetArgs
from phidata.infra.aws.resource.s3.bucket import S3Bucket
from phidata.utils.cli_console import print_info, print_error
from phidata.utils.log import logger


class S3ObjectArgs(DataAssetArgs):
    key: str
    bucket: S3Bucket
    update_ok: bool = True


class S3Object(DataAsset):
    def __init__(
        self,
        key: str,
        bucket: S3Bucket,
        update_ok: bool = False,
        version: Optional[str] = None,
        enabled: bool = True,
    ) -> None:

        super().__init__()
        try:
            self.args: S3ObjectArgs = S3ObjectArgs(
                key=key,
                bucket=bucket,
                name=key,
                update_ok=update_ok,
                version=version,
                enabled=enabled,
            )
        except Exception as e:
            print_error(f"Args for {self.__class__.__name__} are not valid")
            raise

    @property
    def key(self) -> str:
        return self.args.key

    @key.setter
    def key(self, key: str) -> None:
        if key is not None:
            self.args.key = key

    @property
    def bucket(self) -> S3Bucket:
        return self.args.bucket

    @bucket.setter
    def bucket(self, bucket: S3Bucket) -> None:
        if bucket is not None:
            self.args.bucket = bucket

    ######################################################
    ## Write to s3
    ######################################################

    def upload_file(self, file: File) -> bool:

        ######################################################
        ## Validate
        ######################################################

        # S3Object not yet initialized
        if self.args is None:
            return False

        # Check file_path is available
        file_path: Optional[Path] = file.file_path
        if file_path is None or not isinstance(file_path, Path):
            print_error("FilePath invalid")
            return False

        bucket_name = self.bucket.name
        object_key = self.args.key
        print_info("Uploading")
        print_info(f"  File: {file_path}")
        print_info(f"  To  : {bucket_name}/{object_key}")

        ######################################################
        ## Upload
        ######################################################
        try:
            from botocore.exceptions import ClientError

            s3_bucket = self.bucket.read()
            # logger.info("Bucket: {}".format(s3_bucket))
            # logger.info("Bucket type: {}".format(type(s3_bucket)))

            if s3_bucket is None:
                print_error(
                    f"Could not reach bucket: {bucket_name}. Please confirm it exists."
                )
                return False

            # upload file to bucket
            s3_object_size = 0
            s3_object = s3_bucket.Object(key=object_key)
            try:
                s3_object.load()
                s3_object_size = s3_object.content_length
                # logger.info("s3_object_size: {}".format(s3_object_size))
            except ClientError as e:
                pass
            if s3_object_size >= 1 and not self.args.update_ok:
                logger.info(f"Object {bucket_name}/{object_key} already exists")
                return False
            # logger.info("s3_object type: {}".format(type(s3_object)))
            # logger.info("s3_object: {}".format(s3_object))

            logger.info("Uploading S3Object")
            s3_object.upload_file(Filename=str(file_path.resolve()))
            logger.info("S3Object uploaded")
            return True
        except Exception as e:
            print_error("Could not upload S3Object, please try again")
            print_info("--- stacktrace ---")
            logger.exception(e)
            print_info("--- stacktrace ---")
        return False

    ######################################################
    ## Read from s3
    ######################################################

    def download_file(self, file: File) -> bool:

        ######################################################
        ## Validate
        ######################################################

        # S3Object not yet initialized
        if self.args is None:
            return False

        # Check file_path is available
        file_path: Optional[Path] = file.file_path
        if file_path is None or not isinstance(file_path, Path):
            print_error("FilePath invalid")
            return False

        bucket_name = self.bucket.name
        object_key = self.args.key
        print_info("Download")
        print_info(f"  Object : {bucket_name}/{object_key}")
        print_info(f"  To     : {file_path}")

        ######################################################
        ## Download
        ######################################################
        try:
            from botocore.exceptions import ClientError

            s3_bucket = self.bucket.read()
            # logger.info("Bucket: {}".format(s3_bucket))
            # logger.info("Bucket type: {}".format(type(s3_bucket)))

            if s3_bucket is None:
                print_error(
                    f"Could not reach bucket: {bucket_name}. Please confirm it exists."
                )
                return False

            # download object to file
            s3_object_size = 0
            s3_object = s3_bucket.Object(key=object_key)
            try:
                s3_object.load()
                s3_object_size = s3_object.content_length
                logger.info("s3_object_size: {}".format(s3_object_size))
            except ClientError as e:
                pass
            if s3_object_size == 0:
                logger.info(f"Object {bucket_name}/{object_key} does not exist")
                return False
            logger.info("s3_object type: {}".format(type(s3_object)))
            logger.info("s3_object: {}".format(s3_object))

            logger.info("Download S3Object to File")
            s3_object.download_file(Filename=str(file_path.resolve()))
            logger.info("S3Object downloaded")
            return True
        except Exception as e:
            print_error("Could not download S3Object, please try again")
            print_info("--- stacktrace ---")
            logger.exception(e)
            print_info("--- stacktrace ---")
        return False
