"""
s3的公共操作
版本号见setup.py
"""
# coding=utf-8
import boto3
import time
import random
import pathlib
import os
import gzip
import pprint
from tfduck.common.defines import BMOBJ, Et
from botocore.exceptions import ClientError
from concurrent.futures import ThreadPoolExecutor, as_completed, wait, ALL_COMPLETED, FIRST_COMPLETED
from io import StringIO, BytesIO
from botocore.client import Config


class S3Oper(object):
    """
    @des: S3的公共操作, 上传文件
    """

    def __init__(self, ctx={}, ak="", sk="", bucket="", region_name=""):
        """
        @des: 初始化
        """
        self.ctx = ctx
        """
        aws key
        """
        self.access_key = ak
        self.secret_key = sk
        """
        aws s3
        """
        self.bucket = bucket
        self.region_name = region_name
        """
        初始化client
        """
        self.set_s3_client()

    def get_s3_config(self):
        """
        @des: 设置上传文件下载文件，超时的一些设置
        """
        # connect_timeout 和 read_timeout默认值都是60秒
        # max_pool_connections 最大的连接池，默认10
        config = Config(connect_timeout=60,
                        read_timeout=600,
                        retries={'max_attempts': 0}, # 将默认重试次数设为0，不重试
                        max_pool_connections=60)   # upload_foler的max_workers必须比这个值小
        # s3 = boto3.client('s3', config=config)
        return config

    def get_s3_session(self):
        """
        @des: 初始化session
        """
        s3_session = boto3.Session(
            aws_access_key_id=self.access_key,
            aws_secret_access_key=self.secret_key,
            region_name=self.region_name
        )
        return s3_session

    def get_s3_resource(self):
        """
        @des: 初始化resource
        """
        s3_session = self.get_s3_session()
        config = self.get_s3_config()
        s3_resource = s3_session.resource('s3')
        # s3_resource = s3_session.resource('s3', config=config)
        return s3_resource

    def get_s3_bucket(self):
        """
        @des: 初始化bucket
        #####
        bucket = self.get_bucket()
        res = bucket.upload_file(Filename=local_path, Key=remote_path)
        #
        bucket = self.get_bucket()
        bucket.download_file(remote_path, local_path)
        ####
        """
        s3_resource = self.get_s3_resource()
        bucket = s3_resource.Bucket(self.bucket)
        return bucket

    def get_s3_client(self):
        """
        @des: 初始化client
        """
        s3_session = self.get_s3_session()
        config = self.get_s3_config()
        s3_client = s3_session.client("s3", config=config)
        return s3_client
    
    def set_s3_client(self):
        """
        @des: 初始化client
        """
        client = self.get_s3_client()
        self.s3_client = client

    def find_prefix(self, remote_path, isrm=True):
        """
        @des: 遍历s3文件夹，找到所有文件路径列表
        @params: isrm ------是否递归查找子目录下的文件
        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#S3.Client.list_objects_v2
        """
        ctx = self.ctx
        bucket = self.bucket
        file_list = []
        client = self.get_s3_client()
        """
        这种方式不能超过1000
        """
        # # print(help(client.list_objects_v2))
        # response = client.list_objects_v2(
        #         Bucket=self.bucket,
        #         Delimiter="",
        #         Prefix =remote_path,
        #         MaxKeys=1000
        #         ) # MaxKeys的最大值为1000， 不能超过1000,
        # print(response)
        # #
        # for obj in response['Contents']:
        #     key = obj['Key']
        #     file_list.append(key)
        #     print(key)
        # return file_list
        """
        可以超过1000
        """
        paginator = client.get_paginator(
            'list_objects_v2')  # MaxKeys的最大值为1000， 不能超过1000,
        # 所以用分页的方式，这样就可以通过分很多页超过1000
        pages = paginator.paginate(
            Bucket=self.bucket, Prefix=remote_path, MaxKeys=1000)
        remote_path_parents_count = len(pathlib.Path(remote_path).parents)
        for page in pages:
            for obj in page['Contents']:
                key = obj['Key']
                if not isrm:
                    if len(pathlib.Path(key).parents)-1 == remote_path_parents_count:
                        file_list.append(key)
                else:
                    file_list.append(key)
        # pprint.pprint(file_list)
        return file_list

    def upload_file(self, local_path, remote_path):
        """
        @des: 将本地文件上传到s3
        """
        s3_client = self.get_s3_client()
        try:
            _ = s3_client.upload_file(local_path, self.bucket, remote_path)
        except ClientError as e:
            BMOBJ.clog(self.ctx, "upload s3 file error:", e)
            return False
        return True

    def upload_fileobj(self, io_obj, remote_path):
        """
        @des: 将bytesIO上传到s3

        上传普通文件
        ########
        s3 = S3Oper(...)
        with BytesIO() as r:
            r.write(b'')
            r.seek(0)
            s3.upload_fileobj(r,
                        "sagemaker/yxtest/SUCCESS")
        ########

        上传压缩文件
        ########
        s3 = S3Oper(...)
        with BytesIO() as f:
            with gzip.open(f, 'wb') as r:
                r.write(b'')
            f.seek(0)
            s3.upload_fileobj(f,
                        "sagemaker/yxtest/SUCCESS")
        ########

        """
        s3_client = self.get_s3_client()
        try:
            _ = s3_client.upload_fileobj(io_obj, self.bucket, remote_path)
        except ClientError as e:
            BMOBJ.clog(self.ctx, "upload s3 file error:", e)
            return False
        return True

    def download_file(self, local_path, remote_path):
        """
        @des: 将文件下载到本地
        """
        s3_client = self.get_s3_client()
        try:
            _ = s3_client.download_file(self.bucket, remote_path, local_path)
            BMOBJ.clog(self.ctx, 111, _)
        except ClientError as e:
            BMOBJ.clog(self.ctx, "download s3 file error:", e)
            return False
        return True

    def is_exists(self, remote_path):
        """
        @des: 判断资源是否存在
        s3 = boto3.resource('s3')
        try:
            s3.head_object(Bucket=bucket, Key=key)
        except ClientError as e:
            return int(e.response['Error']['Code']) != 404
        return True
        或者
        try:
            s3.Object(bucket, key).load()
        except ClientError as e:
            return int(e.response['Error']['Code']) != 404
        return True
        """
        s3_client = self.get_s3_client()
        try:
            s3_client.head_object(Bucket=self.bucket, Key=remote_path)
        except ClientError as e:
            return int(e.response['Error']['Code']) != 404
        return True

    def delete_file(self, remote_path):
        """
        @des: 删除文件
        """
        s3_client = self.get_s3_client()
        try:
            s3_client.delete_object(Bucket=self.bucket, Key=remote_path)
        except ClientError as e:
            BMOBJ.clog(self.ctx, "delete s3 file error:", e)
            return False
        return True

    def delete_folder(self, remote_path):
        """
        @des: 删除目录
        """
        bucket = self.get_s3_bucket()
        bucket.objects.filter(Prefix=remote_path).delete()
        return True

    def _upload_i(self,  remote_file, local_file):
        """
        @des: 多线程批量上传
        """
        ctx = self.ctx
        # s3_client = self.get_s3_client()
        s3_client = self.s3_client
        #
        retry_count = 3
        for i in range(retry_count):  # 最多重试三次，由于网络不稳定等问题
            try:
                _s = time.time()
                _ = s3_client.upload_file(
                    local_file, self.bucket, remote_file)  # 返回值为None
                _e = time.time()
                BMOBJ.clog(
                    ctx, f"{local_file} upload success, sub time {_e-_s}", )
                break
            except Exception as e:
                BMOBJ.clog(
                    ctx, f"{local_file} upload  fail, repeat {i}, error: {e}")
                #
                if i < retry_count-1:
                    sleep_time = random.randint(60, 120)
                    time.sleep(sleep_time)
                    continue
                else:
                    BMOBJ.clog(
                        ctx, f"{local_file} upload  finally fail: {e}")
                    raise Et(2, f"upload fail {remote_file} {local_file}")

    def upload_folder(self, local_path, remote_path, add_success=False, add_empty=False, max_workers=50, isrm=True, isdel=True):
        """
        @des: 上传到s3---多线程上传---上传文件夹
        """
        ctx = self.ctx
        s = time.time()
        # 删除已经存在的remote_folder
        self.delete_folder(remote_path)
        # 上传
        if isrm:  # 遍历文件夹和子文件夹
            subobjs = list(pathlib.Path(local_path).rglob("*"))
        else:  # 只遍历当前文件夹
            subobjs = list(pathlib.Path(local_path).glob("*"))
        subfiles = [subobj for subobj in subobjs if subobj.is_dir() == False]
        # 打印文件信息
        total_files = []
        total_size = 0
        for subfile in subfiles:
            size = round(subfile.stat().st_size/1024/1024, 4)
            total_size += size
            name = subfile.name
            total_files.append(f"{size}M {name}")
        _infos = '\n'.join(total_files)
        BMOBJ.clog(
            ctx, f"""upload file info *  file total count {len(subfiles)}  file total size {total_size}M""", _infos)

        # 参考 https://www.jianshu.com/p/b9b3d66aa0be
        # 控制最大队列数200，记得修改settings.py的redis队列数必须大于这个
        if subfiles:
            executor = ThreadPoolExecutor(max_workers=max_workers)
            all_tasks = []
            for subfile in subfiles:
                if not isrm:
                    remote_file = os.path.join(remote_path, subfile.name)
                else:
                    l_name = str(subfile).replace(local_path, "")
                    l_name = l_name.lstrip("/")  # 去掉开始的全部斜杠
                    remote_file = os.path.join(remote_path, l_name)
                local_file = str(subfile)
                # 通过submit函数提交执行的函数到线程池中，submit函数立即返回，不阻塞
                task_i = executor.submit(
                    self._upload_i, *(remote_file, local_file))
                all_tasks.append(task_i)
            # 等待所有任务完成后
            # wait(all_tasks, timeout=timeout, return_when=ALL_COMPLETED)
            for future in as_completed(all_tasks):  # 这个子线程出错会抛出来
                _ = future.result()
            # 判断是否上传一个成功的文件
            if add_success:
                # 上传成功后，上传一个空文件代表成功
                with BytesIO() as f:
                    with gzip.open(f, 'wb') as r:
                        r.write(b'')
                    f.seek(0)
                    self.upload_fileobj(
                        f, os.path.join(remote_path, "_SUCCESS"))
        else:
            if add_empty:
                # 上传一个empty文件，代表没有数据
                with BytesIO() as f:
                    with gzip.open(f, 'wb') as r:
                        r.write(b'')
                    f.seek(0)
                    self.upload_fileobj(f, os.path.join(remote_path, "_EMPTY"))
        e = time.time()
        # 删除所有本地文件
        if isdel:
            if len(pathlib.Path(local_path).parents) == 0:
                raise Et(2, "cannt del root folder")
            BMOBJ.remove_folder(local_path)
        #
        BMOBJ.clog(
            ctx, f"{remote_path} upload all time", e-s)

    def _download_i(self, remote_file, local_file):
        """
        @des: 多线程批量下载
        """
        ctx = self.ctx
        # s3_client = self.get_s3_client()
        s3_client = self.s3_client
        #
        retry_count = 3
        for i in range(retry_count):  # 最多重试三次，由于网络不稳定等问题
            try:
                _s = time.time()
                _ = s3_client.download_file(
                    self.bucket, remote_file, local_file)  # 返回值为None
                _e = time.time()
                BMOBJ.clog(
                    ctx, f"{local_file} download success, sub time {_e-_s}", )
                break
            except Exception as e:
                BMOBJ.clog(
                    ctx, f"{local_file} download  fail, repeat {i}, error: {e}")
                if i < retry_count-1:
                    sleep_time = random.randint(60, 120)
                    time.sleep(sleep_time)
                    continue
                else:
                    BMOBJ.clog(
                        ctx, f"{local_file} download  finally fail: {e}")
                    raise Et(2, f"download fail {remote_file} {local_file}")

    def download_folder(self, local_path, remote_path, max_workers=50, isrm=True, isdel=True):
        """
        @des: 下载到本地---多线程下载---下载文件夹--下载后删除s3的文件
        """
        ctx = self.ctx
        # bucket = self.bucket
        s = time.time()
        # 删除本地已经存在的文件,重新创建本地路径
        BMOBJ.remove_folder(local_path)
        os.makedirs(local_path, exist_ok=True)
        # 下载
        subfiles = self.find_prefix(remote_path, isrm=isrm)
        if subfiles:
            executor = ThreadPoolExecutor(max_workers=max_workers)
            all_tasks = []
            for subfile in subfiles:
                remote_file = subfile
                subfile_name = pathlib.PurePath(remote_file).name
                if not isrm:
                    local_file = os.path.join(local_path, subfile_name)
                else:
                    l_name = str(pathlib.PurePath(remote_file)
                                 ).replace(remote_path, "")
                    l_name = l_name.lstrip("/")
                    local_file = os.path.join(local_path, l_name)
                    os.makedirs(os.path.dirname(local_file),
                                exist_ok=True)  # 创建不存在的子文件夹
                # 通过submit函数提交执行的函数到线程池中，submit函数立即返回，不阻塞
                task_i = executor.submit(
                    self._download_i, *(remote_file, local_file))
                all_tasks.append(task_i)
            # 等待所有任务完成后
            # wait(all_tasks, timeout=timeout, return_when=ALL_COMPLETED)
            for future in as_completed(all_tasks):  # 这个子线程出错会抛出来
                _ = future.result()
        e = time.time()
        # 删除s3已经存在的part_date的文件---内网端
        print(remote_path)
        if isdel:
            if len(pathlib.Path(remote_path).parents) == 0:
                raise Et(2, "cannt del root folder")
            self.delete_folder(remote_path)
        #
        BMOBJ.clog(
            ctx, f"{remote_path} download all time", e-s)


if __name__ == '__main__':  # 打版本的时候一定记得记得脱敏
    pass
    # s3 = S3Oper(ctx = {}, ak="xx", sk="yy", bucket="xx", region_name="us-east-2")
    # s3.upload_folder(local_path="/Users/yuanxiao/Downloads/train/samples",
    #                  remote_path="sagemaker/yxtest/samples", isdel=False, add_empty=True, add_success=True, max_workers=50)
    # s3.find_prefix("sagemaker/yxtest", True)
    # s3.download_folder(local_path="/Users/yuanxiao/Downloads/train/samples",
    #                    remote_path="sagemaker/yxtest/samples", isdel=False, max_workers=50)
