# -*-coding:utf-8 -*-

import os
import random
import string
import zipfile
import tarfile
from urllib.parse import quote_plus
import traceback
from obs import ObsClient


# The /tmp directory where packages will be downloaded and decompressed is max. 512 MB.
# To increase the space, mount a disk to your function and specify LOCAL_MOUNT_PATH.
LOCAL_MOUNT_PATH = '/tmp/'
LETTERS = string.ascii_letters


class FileProcessing:

    def __init__(self, context):
        self.logger = context.getLogger()
        obs_endpoint = context.getUserData("obs_server")
        self.obs_client = new_obs_client(context, obs_endpoint)
        self.target_bucket = context.getUserData("resource_bucket")
        self.target_prefix = context.getUserData("target_prefix", "")
        self.download_dir = gen_local_download_path()
        self.extract_dir = ""

    def run(self, record):
        (bucket, object_key) = get_obs_obj_info(record)
        self.logger.info("src bucket:" + bucket)
        self.logger.info("src object_key:" + object_key)
        (path, file) = os.path.split(object_key)
        download_path = self.download_dir + "/" + file
        download_result = self.download_file_from_obs(bucket, quote_plus(object_key), download_path)
        if not download_result:
            return "FAILED"
        (filename, ext) = os.path.splitext(file)
        self.logger.info('start to extract file %s', download_path)
        self.extract_dir = self.download_dir + "/" + filename
        if not self.extract_download_file(download_path):
            return "FAILED"
        self.logger.info('succeeded to extract file %s', download_path)
        self.logger.info('start to upload files in %s to obs', self.extract_dir)

        self.upload_local_file(self.extract_dir)
        self.logger.info('succeeded to upload files in %s to obs', self.extract_dir)
        return 'SUCCESS'

    def download_file_from_obs(self, bucket, obj_name, download_path):
        self.logger.info(f'start to download object %s from obs %s to local %s',
                         obj_name, bucket, download_path)
        try:
            resp = self.obs_client.getObject(bucket, obj_name,
                                             downloadPath=download_path)
            if resp.status < 300:
                self.logger.info(
                    f'succeeded to download object %s from obs %s to local %s',
                    obj_name, bucket, download_path)
                return True
            else:
                self.logger.error(
                    f"failed to download object {obj_name} from obs {bucket}, "
                    f"errorCode:{resp.errorCode} errorMessage:{resp.errorMessage}")
        except:
            self.logger.error(
                f"failed to download file {obj_name} from obs bucket{bucket}, "
                f"exp:{traceback.format_exc()}")

    def upload_local_file(self, directory):
        for filename in os.listdir(directory):
            file_path = os.path.join(directory, filename)
            if os.path.isfile(file_path):
                object_key = self.get_object_key(file_path)
                self.upload_file_to_obs(object_key, file_path)
            elif os.path.isdir(file_path):
                self.upload_local_file(file_path)

    def clean_local_files(self, file_path):
        if os.path.isfile(file_path):
            self.delete_local_file(file_path)
        elif os.path.isdir(file_path):
            for filename in os.listdir(file_path):
                sub_file_path = os.path.join(file_path, filename)
                if os.path.isfile(sub_file_path):
                    self.delete_local_file(sub_file_path)
                elif os.path.isdir(sub_file_path):
                    self.clean_local_files(sub_file_path)
                    os.rmdir(sub_file_path)

    def upload_file_to_obs(self, object_key, file_path):
        try:
            resp = self.obs_client.putFile(self.target_bucket, object_key,
                                           file_path)
            if resp.status > 300:
                self.logger.error(f"failed to upload file {file_path} to "
                                  f"obs bucket {self.target_bucket}, "
                                  f"errorCode:{resp.errorCode} "
                                  f"errorMessage:{resp.errorMessage}")
                return None
        except:
            self.logger.error(f"failed to upload file {file_path} to "
                              f"obs bucket {self.target_bucket}, "
                              f"exp:{traceback.format_exc()}")

    def delete_local_file(self, file_path):
        try:
            os.remove(file_path)
        except:
            self.logger.error(
                f"failed to delete local file {file_path}, "
                f"exp:{traceback.format_exc()}")

    def get_object_key(self, file_path):
        object_key = file_path
        if self.target_prefix != "":
            object_key = self.target_prefix + file_path
        object_key = object_key.replace(self.extract_dir, "")
        if object_key.startswith("/"):
            object_key = object_key[1:]
        return object_key

    def extract_download_file(self, download_path):
        if zipfile.is_zipfile(download_path):
            zipfile.ZipFile(download_path).extractall(self.extract_dir)
            return True

        if tarfile.is_tarfile(download_path):
            tarfile.open(download_path).extractall(self.extract_dir)
            return True

        self.logger.error("failed to extract file %s', only support "
                          "compression formats: zip and tar "
                          "(gzip, bz2 and lzma)", download_path)


def new_obs_client(context, obs_server):
    ak = context.getAccessKey()
    sk = context.getSecretKey()
    return ObsClient(access_key_id=ak, secret_access_key=sk, server=obs_server)


def get_obs_obj_info(record):
    if 's3' in record:
        s3 = record['s3']
        return s3['bucket']['name'], s3['object']['key']


def gen_local_download_path():
    download_dir = LOCAL_MOUNT_PATH + ''.join(
        random.choice(LETTERS) for i in range(16))
    os.makedirs(download_dir)
    return download_dir