import enum
import io
import json

import boto3
import oss2
from qcloud_cos import CosConfig, CosS3Client
from google.cloud import storage
from google.oauth2 import service_account
from google.oauth2.credentials import Credentials


class ObjectStorageType(enum.Enum):
    OSS = 0
    COS = 1
    S3 = 2
    GCS = 3


class ObjectStorageClient(object):
    def __init__(self, objectStorageType, ak_id, ak_secret, token, endpoint, bucket=None):
        self.type = objectStorageType
        self.ak_id = ak_id
        self.ak_secret = ak_secret
        self.token = token
        self.endpoint = endpoint
        self.bucket = bucket
        self.client = None
        if self.type == ObjectStorageType.OSS:
            auth = oss2.StsAuth(self.ak_id, self.ak_secret, self.token)
            self.client = oss2.Bucket(auth, self.endpoint, self.bucket)
        elif self.type == ObjectStorageType.COS:
            cos_config = CosConfig(Region=self.endpoint, SecretId=self.ak_id, SecretKey=self.ak_secret, Token=self.token)
            self.client = CosS3Client(cos_config)
        elif self.type == ObjectStorageType.S3:
            session = boto3.Session(aws_session_token=self.token,
                                    aws_access_key_id=self.ak_id,
                                    aws_secret_access_key=self.ak_secret,
                                    region_name=self.endpoint)
            self.client = session.resource('s3')
        elif self.type == ObjectStorageType.GCS:
            credentials = None
            if self.ak_id is not None and len(self.ak_id) > 0:
                credentials = service_account.Credentials.from_service_account_info(json.loads(self.ak_id))
            elif self.token is not None and len(self.token) > 0:
                credentials = Credentials(token=self.token)
            self.client = storage.Client(project='clickzetta-connector', credentials=credentials)

    def get_stream(self, bucket_name, object_name):
        stream = None
        if self.type == ObjectStorageType.OSS:
            stream = self.client.get_object(object_name).read()
        elif self.type == ObjectStorageType.COS:
            stream = self.client.get_object(Bucket=bucket_name, Key=object_name)['Body'].get_raw_stream()
        elif self.type == ObjectStorageType.S3:
            stream = self.client.Object(bucket_name, object_name).get()['Body'].read()
        elif self.type == ObjectStorageType.GCS:
            bucket = self.client.bucket(bucket_name)
            blob = bucket.blob(object_name)
            file_obj = io.BytesIO()
            blob.download_to_file(file_obj)
            file_obj.seek(0)
            stream = file_obj
        return stream




