import base64
import fnmatch
import json
import logging
import re
from datetime import timedelta
from enum import Enum
from functools import lru_cache
from json import JSONDecodeError
from typing import Optional, Union
from urllib.parse import urlparse

import google.auth
import google.cloud.storage as gcs
from core.utils.common import get_ttl_hash
from django.conf import settings
from google.auth.exceptions import DefaultCredentialsError
from google.oauth2 import service_account

logger = logging.getLogger(__name__)

Base64 = bytes


class GCS(object):
    _client_cache = {}
    _credentials_cache = None
    DEFAULT_GOOGLE_PROJECT_ID = gcs.client._marker

    class ConvertBlobTo(Enum):
        NOTHING = 1
        JSON = 2
        JSON_DICT = 3
        BASE64 = 4

    @classmethod
    @lru_cache(maxsize=1)
    def get_bucket(
        cls,
        ttl_hash: int,
        google_project_id: Optional[str] = None,
        google_application_credentials: Optional[Union[str, dict]] = None,
        bucket_name: Optional[str] = None,
    ) -> gcs.Bucket:

        client = cls.get_client(
            google_project_id=google_project_id, google_application_credentials=google_application_credentials
        )

        return client.get_bucket(bucket_name)

    @classmethod
    def get_client(
        cls, google_project_id: str = None, google_application_credentials: Union[str, dict] = None
    ) -> gcs.Client:
        """
        :param google_project_id:
        :param google_application_credentials:
        :return:
        """
        google_project_id = google_project_id or GCS.DEFAULT_GOOGLE_PROJECT_ID
        cache_key = google_application_credentials

        if cache_key not in GCS._client_cache:

            # use credentials from LS Cloud Storage settings
            if google_application_credentials:
                if isinstance(google_application_credentials, str):
                    try:
                        google_application_credentials = json.loads(google_application_credentials)
                    except JSONDecodeError as e:
                        # change JSON error to human-readable format
                        raise ValueError(f'Google Application Credentials must be valid JSON string. {e}')
                credentials = service_account.Credentials.from_service_account_info(google_application_credentials)
                GCS._client_cache[cache_key] = gcs.Client(project=google_project_id, credentials=credentials)

            # use Google Application Default Credentials (ADC)
            else:
                GCS._client_cache[cache_key] = gcs.Client(project=google_project_id)

        return GCS._client_cache[cache_key]

    @classmethod
    def validate_connection(
        cls,
        bucket_name: str,
        google_project_id: str = None,
        google_application_credentials: Union[str, dict] = None,
        prefix: str = None,
        use_glob_syntax: bool = False,
    ):
        logger.debug('Validating GCS connection')
        client = cls.get_client(
            google_application_credentials=google_application_credentials, google_project_id=google_project_id
        )
        logger.debug('Validating GCS bucket')
        bucket = client.get_bucket(bucket_name)

        # Dataset storages uses glob syntax and we want to add explicit checks
        # In the future when GCS lib supports it
        if use_glob_syntax:
            pass
        else:
            if prefix:
                blobs = list(bucket.list_blobs(prefix=prefix, max_results=1))
                if not blobs:
                    raise ValueError(f"No blobs found in {bucket_name}/{prefix} or prefix doesn't exist")

    @classmethod
    def iter_blobs(
        cls,
        client: gcs.Client,
        bucket_name: str,
        prefix: str = None,
        regex_filter: str = None,
        limit: int = None,
        return_key: bool = False,
        recursive_scan: bool = True,
    ):
        """
        Iterate files on the bucket. Optionally return limited number of files that match provided extensions
        :param client: GCS Client obj
        :param bucket_name: bucket name
        :param prefix: bucket prefix
        :param regex_filter: RegEx filter
        :param limit: specify limit for max files
        :param return_key: return object key string instead of gcs.Blob object
        :return: Iterator object
        """
        total_read = 0
        # Normalize prefix to end with '/'
        normalized_prefix = (str(prefix).rstrip('/') + '/') if prefix else ''
        # Use delimiter for non-recursive listing
        if recursive_scan:
            blob_iter = client.list_blobs(bucket_name, prefix=normalized_prefix or None)
        else:
            blob_iter = client.list_blobs(bucket_name, prefix=normalized_prefix or None, delimiter='/')
        prefix = normalized_prefix
        regex = re.compile(str(regex_filter)) if regex_filter else None
        for blob in blob_iter:
            # skip directory entries at any level (directories end with '/')
            if blob.name.endswith('/'):
                continue
            # check regex pattern filter
            if regex and not regex.match(blob.name):
                logger.debug(blob.name + ' is skipped by regex filter')
                continue
            if return_key:
                yield blob.name
            else:
                yield blob
            total_read += 1
            if limit and total_read == limit:
                break

    @classmethod
    def _get_default_credentials(cls):
        """Get default GCS credentials for LS Cloud Storages"""
        # TODO: remove this func with fflag_fix_back_lsdv_4902_force_google_adc_16052023_short
        try:
            # check if GCS._credentials_cache is None, we don't want to try getting default credentials again
            credentials = GCS._credentials_cache.get('credentials') if GCS._credentials_cache else None
            if GCS._credentials_cache is None or (credentials and credentials.expired):
                # try to get credentials from the current environment
                credentials, _ = google.auth.default(['https://www.googleapis.com/auth/cloud-platform'])
                # apply & refresh credentials
                auth_req = google.auth.transport.requests.Request()
                credentials.refresh(auth_req)
                # set cache
                GCS._credentials_cache = {
                    'service_account_email': credentials.service_account_email,
                    'access_token': credentials.token,
                    'credentials': credentials,
                }

        except DefaultCredentialsError as exc:
            logger.warning(f'Label studio could not load default GCS credentials from env. {exc}', exc_info=True)
            GCS._credentials_cache = {}

        return GCS._credentials_cache

    @classmethod
    def generate_http_url(
        cls,
        url: str,
        presign: bool,
        google_application_credentials: Union[str, dict] = None,
        google_project_id: str = None,
        presign_ttl: int = 1,
    ) -> str:
        """
        Gets gs:// like URI string and returns presigned https:// URL
        :param url: input URI
        :param presign: Whether to generate presigned URL. If false, will generate base64 encoded data URL
        :param google_application_credentials:
        :param google_project_id:
        :param presign_ttl: Presign TTL in minutes
        :return: Presigned URL string
        """
        r = urlparse(url, allow_fragments=False)
        bucket_name = r.netloc
        blob_name = r.path.lstrip('/')

        """Generates a v4 signed URL for downloading a blob.

        Note that this method requires a service account key file. You can not use
        this if you are using Application Default Credentials from Google Compute
        Engine or from the Google Cloud SDK.
        """
        bucket = cls.get_bucket(
            ttl_hash=get_ttl_hash(),
            google_application_credentials=google_application_credentials,
            google_project_id=google_project_id,
            bucket_name=bucket_name,
        )

        blob = bucket.blob(blob_name)

        # this flag should be OFF, maybe we need to enable it for 1-2 customers, we have to check it
        if settings.GCS_CLOUD_STORAGE_FORCE_DEFAULT_CREDENTIALS:
            # google_application_credentials has higher priority,
            # use Application Default Credentials (ADC) when google_application_credentials is empty only
            maybe_credentials = {} if google_application_credentials else cls._get_default_credentials()
            maybe_client = None if google_application_credentials else cls.get_client()
        else:
            maybe_credentials = {}
            maybe_client = None

        if not presign:
            blob.reload(client=maybe_client)  # needed to know the content type
            blob_bytes = blob.download_as_bytes(client=maybe_client)
            return f'data:{blob.content_type};base64,{base64.b64encode(blob_bytes).decode("utf-8")}'

        url = blob.generate_signed_url(
            version='v4',
            # This URL is valid for 15 minutes
            expiration=timedelta(minutes=presign_ttl),
            # Allow GET requests using this URL.
            method='GET',
            **maybe_credentials,
        )

        logger.debug('Generated GCS signed url: ' + url)
        return url

    @classmethod
    def iter_images_base64(cls, client, bucket_name, max_files):
        for image in cls.iter_blobs(client, bucket_name, max_files):
            yield GCS.read_base64(image)

    @classmethod
    def iter_images_filename(cls, client, bucket_name, max_files):
        for image in cls.iter_blobs(client, bucket_name, max_files):
            yield image.name

    @classmethod
    def get_uri(cls, bucket_name, key):
        return f'gs://{bucket_name}/{key}'

    @classmethod
    def read_file(
        cls, client: gcs.Client, bucket_name: str, key: str, convert_to: ConvertBlobTo = ConvertBlobTo.NOTHING
    ):
        bucket = client.get_bucket(bucket_name)
        blob = bucket.blob(key)
        blob = blob.download_as_bytes()

        if convert_to == cls.ConvertBlobTo.BASE64:
            return base64.b64encode(blob)

        return blob

    @classmethod
    def read_base64(cls, f: gcs.Blob) -> Base64:
        return base64.b64encode(f.download_as_bytes())

    @classmethod
    def get_blob_metadata(
        cls,
        url: str,
        google_application_credentials: Union[str, dict] = None,
        google_project_id: str = None,
        properties_name: list = [],
    ) -> dict:
        """
        Gets object metadata like size and updated date from GCS in dict format
        :param url: input URI
        :param google_application_credentials:
        :param google_project_id:
        :return: Object metadata dict("name": "value")
        """
        r = urlparse(url, allow_fragments=False)
        bucket_name = r.netloc
        blob_name = r.path.lstrip('/')

        client = cls.get_client(
            google_application_credentials=google_application_credentials, google_project_id=google_project_id
        )
        bucket = client.get_bucket(bucket_name)
        # Get blob instead of Blob() is used to make an http request and get metadata
        blob = bucket.get_blob(blob_name)
        if not properties_name:
            return blob._properties
        return {key: value for key, value in blob._properties.items() if key in properties_name}

    @classmethod
    def validate_pattern(cls, storage, pattern, glob_pattern=True):
        """
        Validate pattern against Google Cloud Storage
        :param storage: Google Cloud Storage instance
        :param pattern: Pattern to validate
        :param glob_pattern: If True, pattern is a glob pattern, otherwise it is a regex pattern
        :return: Message if pattern is not valid, empty string otherwise
        """
        client = storage.get_client()
        blob_iter = client.list_blobs(
            storage.bucket, prefix=storage.prefix, page_size=settings.CLOUD_STORAGE_CHECK_FOR_RECORDS_PAGE_SIZE
        )
        prefix = str(storage.prefix) if storage.prefix else ''
        # compile pattern to regex
        if glob_pattern:
            pattern = fnmatch.translate(pattern)
        regex = re.compile(str(pattern))
        for index, blob in enumerate(blob_iter):
            # skip directories
            if blob.name == (prefix.rstrip('/') + '/'):
                continue
            # check regex pattern filter
            if pattern and regex.match(blob.name):
                logger.debug(blob.name + ' matches file pattern')
                return ''
        return 'No objects found matching the provided glob pattern'