import os
import gzip
import boto3
from botocore.config import Config
from os import environ
from concurrent.futures import ThreadPoolExecutor
from io import IOBase, BytesIO
from typing import List
from botocore.exceptions import ClientError
from mypy_boto3_s3.service_resource import S3ServiceResource
from google.cloud.storage import Client, Blob
from google.cloud.exceptions import NotFound
from azure.storage.blob import BlobServiceClient
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import DefaultAzureCredential

S3_MAX_FETCH_THREADS = 100


class ObjectStore:
    def list(self, bucket: str, prefix: str) -> List[str]:
        raise NotImplementedError

    def lists(self, bucket: str, prefixes: List[str]) -> List[str]:
        thread_count = len(prefixes)
        if thread_count > S3_MAX_FETCH_THREADS:
            thread_count = S3_MAX_FETCH_THREADS
        keys = []
        with ThreadPoolExecutor(thread_count) as executor:
            rs = executor.map(self.list, [bucket] * len(prefixes), prefixes)
            for r in rs:
                for k in r:
                    keys.append(k)
            return keys

    def object(self, bucket: str, key: str) -> IOBase:
        raise NotImplementedError

    def objects(self, bucket: str, keys: List[str]) -> List[IOBase]:
        thread_count = len(keys)
        if thread_count > S3_MAX_FETCH_THREADS:
            thread_count = S3_MAX_FETCH_THREADS
        with ThreadPoolExecutor(thread_count) as executor:
            return list(executor.map(self.object, [bucket] * len(keys), keys))

    def upload(self, bucket: str, key: str, filepath: str):
        raise NotImplementedError

    def download(self, bucket: str, key: str, filepath: str) -> bool:
        raise NotImplementedError

    def existed(self, bucket: str, key: str) -> bool:
        raise NotImplementedError


class Local(ObjectStore):
    path: str

    def __init__(self, path: str):
        self.path = path

    def list(self, bucket: str, prefix: str) -> List[str]:
        dirname = os.path.dirname(prefix)
        basename = os.path.basename(prefix)
        try:
            filenames = os.listdir(os.path.join(self.path, dirname))
        except FileNotFoundError:
            return []
        if basename:
            filenames = [f for f in filenames if f.startswith(basename)]
        keys = []
        for filename in filenames:
            key = os.path.join(dirname, filename)
            if os.path.isfile(os.path.join(self.path, key)):
                keys.append(key)
        return keys

    def lists(self, bucket: str, prefixes: List[str]) -> List[str]:
        filenames = []
        for prefix in prefixes:
            filenames.extend(self.list(bucket, prefix))
        return filenames

    def object(self, bucket: str, key: str) -> IOBase:
        try:
            filepath = os.path.join(self.path, key)
            if key.endswith(".gz"):
                return gzip.open(filepath)
            return open(filepath, "rb")
        except FileNotFoundError:
            return None

    def objects(self, bucket: str, keys: List[str]) -> List[IOBase]:
        files = []
        for key in keys:
            files.append(self.object(bucket, key))
        return files

    def upload(self, bucket: str, key: str, filepath: str):
        dst = os.path.join(self.path, key)
        os.makedirs(os.path.dirname(dst), exist_ok=True)
        with open(filepath, "rb") as r:
            if filepath.endswith(".gz"):
                with open(dst, "wb") as w:
                    w.write(r.read())
            else:
                with gzip.open(dst, "wb") as w:
                    w.write(r.read())

    def download(self, bucket: str, key: str, filepath: str) -> bool:
        if key.endswith(".gz") and not filepath.endswith(".gz"):
            filepath += ".gz"
        try:
            dst = os.path.join(self.path, key)
            with open(dst, "rb") as r, open(filepath, "wb") as w:
                w.write(r.read())
        except FileNotFoundError:
            return False
        if filepath.endswith(".gz"):
            with gzip.open(filepath) as r, open(filepath.removesuffix(".gz"), "wb") as w:
                w.write(r.read())
            os.unlink(filepath)
        return True

    def existed(self, bucket: str, key: str) -> bool:
        return os.path.exists(os.path.join(self.path, key))


class S3(ObjectStore):
    s3: S3ServiceResource

    def __init__(self):
        endpoint = None
        if "AWS_ENDPOINT_S3" in os.environ:
            s3_config = Config(
                s3={'addressing_style': 'virtual',
                'payload_signing_enabled': False},
                # Disable chunked encoding which causes issues with OSS
                signature_version='s3v4'
            )
            endpoint = os.environ["AWS_ENDPOINT_S3"]
            self.s3 = boto3.resource('s3', endpoint_url=endpoint, config=s3_config)
        else:
            self.s3 = boto3.resource("s3")

    def list(self, bucket: str, prefix: str) -> List[str]:
        return [obj.key for obj in self.s3.Bucket(bucket).objects.filter(Prefix=prefix)]

    def object(self, bucket: str, key: str) -> IOBase:
        try:
            content = self.s3.Bucket(bucket).Object(key).get()["Body"]
            if key.endswith(".gz"):
                return gzip.open(content)
            return content
        except ClientError as e:
            if e.response["Error"]["Code"] == "404" or e.response["Error"]["Code"] == "NoSuchKey":
                return None
            else:
                raise e

    def upload(self, bucket: str, key: str, filepath: str):
        with open(filepath, "rb") as f:
            content = f.read()
        if key.endswith(".gz") and not filepath.endswith(".gz"):
            content = gzip.compress(content)
        extra_args = {}
        if key.endswith(".gz"):
            extra_args = {"ContentEncoding": "gzip"}
        self.s3.Bucket(bucket).upload_fileobj(BytesIO(content), key, ExtraArgs=extra_args)

    def download(self, bucket: str, key: str, filepath: str) -> bool:
        if key.endswith(".gz") and not filepath.endswith(".gz"):
            filepath += ".gz"
        try:
            self.s3.Bucket(bucket).download_file(key, filepath)
        except ClientError as e:
            if e.response["Error"]["Code"] == "404" or e.response["Error"]["Code"] == "NoSuchKey":
                return False
            else:
                raise e
        if filepath.endswith(".gz"):
            with gzip.open(filepath) as r, open(filepath.removesuffix(".gz"), "wb") as w:
                w.write(r.read())
            os.unlink(filepath)
        return True

    def existed(self, bucket: str, key: str) -> bool:
        try:
            self.s3.Bucket(bucket).Object(key).load()
            return True
        except ClientError as e:
            if e.response["Error"]["Code"] == "404" or e.response["Error"]["Code"] == "NoSuchKey":
                return False
            else:
                raise


class GCS(ObjectStore):
    client: Client

    def __init__(self):
        self.client = Client()

    def list(self, bucket: str, prefix: str) -> List[str]:
        blobs = self.client.list_blobs(bucket, prefix=prefix)
        return [blob.name for blob in blobs]

    def object(self, bucket: str, key: str) -> IOBase:
        try:
            bucket = self.client.get_bucket(bucket)
            blob = Blob(key, bucket)
            body = blob.download_as_bytes()
            return BytesIO(body)
        except NotFound:
            return None

    def upload(self, bucket: str, key: str, filepath: str):
        bucket = self.client.get_bucket(bucket)
        blob = Blob(key, bucket)
        blob.upload_from_filename(filepath)

    def download(self, bucket: str, key: str, filepath: str) -> bool:
        bucket = self.client.get_bucket(bucket)
        blob = Blob(key, bucket)
        try:
            blob.download_to_filename(filepath)
            return True
        except NotFound:
            return False

    def existed(self, bucket: str, key: str) -> bool:
        try:
            bucket = self.client.get_bucket(bucket)
            blob = Blob(key, bucket)
            blob.reload()
            return True
        except NotFound:
            return False


# Implement Azure Blob Storage
class AzureBlob(ObjectStore):
    blob_service_client: BlobServiceClient

    def __init__(self):
        credential = DefaultAzureCredential()
        self.blob_service_client = BlobServiceClient("https://{}.blob.core.windows.net".format(environ["AZURE_STORAGE_ACCOUNT"]), credential)

    def list(self, bucket: str, prefix: str) -> List[str]:
        container_client = self.blob_service_client.get_container_client(bucket)
        return [blob.name for blob in container_client.list_blobs(name_starts_with=prefix)]

    def object(self, bucket: str, key: str) -> IOBase:
        try:
            blob_client = self.blob_service_client.get_blob_client(container=bucket, blob=key)
            stream = BytesIO()
            blob_client.download_blob().readinto(stream)
            stream.seek(0)
            # if key.endswith(".gz"):
            #     return gzip.open(stream)
            return stream
        except ResourceNotFoundError:
            return None

    def upload(self, bucket: str, key: str, filepath: str):
        blob_client = self.blob_service_client.get_blob_client(container=bucket, blob=key)
        with open(filepath, "rb") as data:
            blob_client.upload_blob(data, overwrite=True)

    def download(self, bucket: str, key: str, filepath: str) -> bool:
        try:
            blob_client = self.blob_service_client.get_blob_client(container=bucket, blob=key)
            with open(filepath, "wb") as file:
                blob_client.download_blob().readinto(file)
            return True
        except ResourceNotFoundError:
            return False

    def existed(self, bucket: str, key: str) -> bool:
        try:
            blob_client = self.blob_service_client.get_blob_client(container=bucket, blob=key)
            blob_client.get_blob_properties()
            return True
        except ResourceNotFoundError:
            return False