import datetime
import os
import shutil
import sys
import time

from tqdm import tqdm
# from alive_progress import alive_bar
import tpys
import zips


scflag = 'yssc'
had_ysdir = "已压缩"  # 放置已压缩的压缩包
restart_number = 2  # 重试次数

tpgs = ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG', '.bmp', '.BMP', '.gif', '.GIF']
ysgs = ['.rar', '.zip', '.RAR', '.ZIP', '.7z', '.7Z']


def strofsize(integer, remainder, level):
    if integer >= 1024:
        remainder = integer % 1024
        integer //= 1024
        level += 1
        return strofsize(integer, remainder, level)
    else:
        return integer, round(remainder / 1024, 2), level


def StrOfSize(size):
    units = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
    integer, remainder, level = strofsize(size, 0, 0)
    if level + 1 > len(units):
        level = -1
    return '{} {}'.format(integer + remainder, units[level])


def find_zip(mydir, yysdir=had_ysdir):  # 寻找符合条件的压缩包
    file_list = []
    file_list2 = []
    for folderName, subfolders, filenames in os.walk(mydir):
        for filename in filenames:
            if os.path.splitext(filename)[1] is not None and os.path.splitext(filename)[1] in ysgs and os.path.splitext(filename)[1] != '':
                if yysdir not in os.path.join(folderName, filename):
                    # 寻找所有压缩包，去除在已压缩目录下的压缩包
                    file_list.append(os.path.join(folderName, filename))
    for flist in file_list:  # 寻找符合条件的压缩包
        if zips.pd_ziprar(flist) == 1:
            file_list2.append(flist)
    return list(set(file_list) ^ set(file_list2))


def unzips(file_list, code=0, flag='123'):  # 解压压缩包
    zip_number = 0
    with tqdm(total=len(file_list)) as pbar:
        for flist in file_list:
            if zips.un_ziprar(flist, code=code, flag=flag) == 0:
                zip_number += 1
                pbar.update(1)
                print("\n")
    return zip_number


# 检查文件是否为图像文件
def is_image_file(filename):
    _, file_extension = os.path.splitext(filename)
    return file_extension.lower() in tpgs

# 获取目录下的图片文件列表


def get_image_files(directory):
    image_files = []
    for filename in os.listdir(directory):
        if is_image_file(filename):
            image_files.append(os.path.join(directory, filename))
    return image_files

# 检查路径中是否包含特定文件名


def contains_specific_name(path):
    specific_name = had_ysdir
    return specific_name in path

# 移动并重命名父目录下的图片文件到含有图片的子目录中


def move_and_rename_parent_directory_images(root_dir):
    count = 0
    for dirpath, dirnames, _ in os.walk(root_dir):
        for subdir in dirnames:
            subdir_path = os.path.join(dirpath, subdir)
            if any(is_image_file(filename) for filename in os.listdir(subdir_path)):
                # 子目录含有图片文件
                images_in_parent = get_image_files(dirpath)
                images_count = len(images_in_parent)
                if images_count < 3:
                    # 如果图片数量小于3，则将父目录下的图片移动到子目录中并对父目录图片进行重命名
                    idx = 1
                    count += 1
                    for image in images_in_parent:
                        new_name = f"image_{idx}{os.path.splitext(image)[1]}"
                        new_path = os.path.join(subdir_path, new_name)
                        shutil.move(image, new_path)
                        idx += 1

                else:
                    # 如果图片数量大于等于3，则新建文件夹并移动图片到新建的文件夹中
                    count += 1
                    new_folder_name = f"Other_Images_{images_count}"
                    new_folder_path = os.path.join(dirpath, new_folder_name)
                    os.makedirs(new_folder_path, exist_ok=True)
                    for image in images_in_parent:
                        shutil.move(image, new_folder_path)
                    break
    return count

# 寻找深度为1的文件夹


def find_depth_one_folder(dirpath, root_dir):
    while dirpath.count(os.path.sep) - root_dir.count(os.path.sep) > 1:
        dirpath = os.path.dirname(dirpath)
    return dirpath

# 遍历文件夹
# 定义一个函数来遍历文件夹并返回包含图片的子文件夹的路径集合


def find_folders_with_images(root_folder, uexcluded):
    folders_with_images = set()
    parent_folders = {}
    excluded_folders = set()

    for root, dirs, files in os.walk(root_folder):
        for dir in dirs:
            folder_path = os.path.join(root, dir)
            parent_folder = os.path.dirname(folder_path)

            # 如果文件夹名为“已压缩”或“已压缩”文件夹的子文件夹，则添加到排除集合中
            if dir == uexcluded or folder_path in excluded_folders:
                excluded_folders.add(folder_path)
                continue

            # 跟踪每个上级目录的子文件夹数量
            if parent_folder in parent_folders:
                parent_folders[parent_folder] += 1
            else:
                parent_folders[parent_folder] = 1

            contains_image = False

            # 检查文件夹中的文件是否为图片
            for file in os.listdir(folder_path):
                file_path = os.path.join(folder_path, file)
                if os.path.isfile(file_path) and is_image_file(file):
                    contains_image = True
                    break

            # 如果包含图片，则将文件夹路径添加到集合中
            if contains_image:
                folders_with_images.add(folder_path)

    # 排除具有相同上级目录且数量超过一个的子文件夹，但保留根目录下的子文件夹，并排除“已压缩”及其所有子文件夹
    for folder_path in folders_with_images.copy():
        parent_folder = os.path.dirname(folder_path)
        if parent_folder != root_folder and parent_folders[parent_folder] > 1 or any(
                folder_path.startswith(excluded) for excluded in excluded_folders):
            folders_with_images.discard(folder_path)

    return folders_with_images


# 检查文件夹及其父目录是否在排除列表中
def is_folder_or_parent_excluded(root_dir, folder, excluded_folders):
    while folder != root_dir:
        if folder in excluded_folders:
            return True
        folder = os.path.dirname(folder)
    return False

# 遍历文件夹(排除已压缩)


def search_image_folders2(root_dir, excluded_folders, excluded_folder):
    image_folders = set()
    for dirpath, dirnames, filenames in os.walk(root_dir):
        if excluded_folder in dirnames:
            dirnames.remove(excluded_folder)
        if any(is_image_file(file) for file in filenames):
            if not is_folder_or_parent_excluded(root_dir, dirpath, excluded_folders):
                image_folders.add(dirpath)
    return list(image_folders)


def getdirsize(mydir):  # 获取文件夹大小
    size = 0
    for root, dirs, files in os.walk(winapi_path(mydir)):
        size += sum([os.path.getsize(os.path.join(root, name))
                    for name in files])
    return size


def winapi_path(dos_path, encoding=None):
    if (not isinstance(dos_path, str) and encoding is not None):
        dos_path = dos_path.decode(encoding)
    path = os.path.abspath(dos_path)
    if path.startswith(u"\\\\"):
        return dos_path
    return u"\\\\?\\" + path


def ys_rootdir(dir_list1, rootdir, model, works, code, flag=scflag, code2=1, reys=restart_number):  # 压缩图片并打包
    pd1 = 0
    pd2 = 0
    compresssize = 0
    if len(dir_list1) != 0:
        # with tqdm(total=len(dir_list1), ncols=70) as pbar:
        with tqdm(total=len(dir_list1)) as pbar:
            for list1 in dir_list1:
                print("开始压缩" + list1 + "内部图片")
                st_size = getdirsize(list1)
                pd2 = tpys.dirys(mydir=list1, model=model,
                                 works=works, code=code, flag=flag)
                if pd2 != 0:
                    print(list1 + "压缩完毕,共出现" + str(pd2) + "处错误")  # 重试
                    while (pd2 != 0 and pd1 < reys):
                        print("开始重试，第" + str(pd1) + "次")
                        pd1 += 1
                        pd2 = tpys.dirys(mydir=list1, model=model,
                                         works=works, code=code, flag=flag)
                        if pd2 == 0:
                            break
                    if pd2 != 0:
                        print(list1 + "压缩失败")
                    else:
                        if code == 1:
                            end_size = getdirsize(list1)
                            if end_size < st_size:
                                print(list1 + "压缩完毕")
                                print("共节约" + StrOfSize(st_size - end_size) +
                                      "空间，压缩率" + '{:.0%}'.format(1 - end_size / st_size))
                                compresssize += (st_size -
                                                 end_size) / (1024 * 1024)
                            else:
                                print("反而没有节约空间")
                        else:
                            print(list1 + "压缩完毕(不删除原始文件)")
                elif pd2 == 0:
                    if code == 1:
                        end_size = getdirsize(list1)
                        if end_size < st_size:
                            print(list1 + "压缩完毕")
                            print("共节约" + StrOfSize(st_size - end_size) +"空间，压缩率" + '{:.0%}'.format(1 - end_size / st_size))
                            compresssize += (st_size -end_size) / (1024 * 1024)
                        else:
                            print("反而没有节约空间")
                    else:
                        print(list1 + "压缩完毕(不删除原始文件)")
                if code2 == 1:  # code3=1 不需要打包
                    pbar.update(1)
                    continue
                ys_dir = find_depth_one_folder(list1, rootdir)  # 文件夹在根目录下
                if zips.zip_file(src_dir=ys_dir, code=code, flag=scflag) != 1:
                    print(list1 + "打包完成")  # 不是二级目录的话移动到指定位置
                    new_dir = mydir + "\\" + had_ysdir + "\\"
                    if os.path.isdir(new_dir):
                        pass
                    else:
                        os.mkdir(new_dir)
                    try:
                        if os.path.exists(ys_dir + ".zip"):
                            shutil.move(ys_dir + ".zip", new_dir)
                            print(ys_dir + ".zip" + "移动文件成功")
                        else:
                            print("移动文件失败")
                    except BaseException:
                        print("移动文件失败")
                else:
                    print(list1 + "打包失败")
                pbar.update(1)  # 更新进度条
                print("\n")
    return compresssize


def ys_nrootdir(dir_list1, model, works, code, flag='no'):  # 压缩图片
    pd2 = 0
    compresssize = 0
    if len(dir_list1) != 0:
        with tqdm(total=len(dir_list1)) as pbar:
            for list1 in dir_list1:
                print("开始压缩" + list1 + "内部图片")
                st_size = getdirsize(list1)
                pd2 = tpys.dirys(mydir=list1, model=model,
                                 works=works, code=code, flag=flag)
                if pd2 != 0:
                    print(list1 + "压缩完毕,共出现" + str(pd2) + "处错误")
                elif pd2 == 0:
                    if code == 1:
                        end_size = getdirsize(list1)
                        if end_size < st_size:
                            print(list1 + "压缩完毕")
                            print("共节约" + StrOfSize(st_size - end_size) +
                                  "空间，压缩率" + '{:.0%}'.format(1 - end_size / st_size))
                            compresssize += (st_size -
                                             end_size) / (1024 * 1024)
                        else:
                            print("反而没有节约空间")
                    else:
                        print(list1 + "压缩完毕")
                pbar.update(1)
                print("\n")
    return compresssize
# 输入开始


while True:  # 判断输入地址是否正确，输入cs可退出
    mydir = input("请输入待处理文件地址：")
    if not os.path.isdir(mydir):
        if mydir != 'cs':
            print("请输入正确的地址")
        else:
            print("退出")
            input("按任意键结束")
            sys.exit()
    else:
        break

print("0 不删除原始压缩包，新压缩包重命名  1 删除原始压缩包，新压缩包不重命名")
while True:
    code1 = input("是否删除原压缩包（默认代码1）：")
    if code1 != '0' and code1 != '1' and code1 != '' and code1 != 'cs':
        print("请输入正确的代码")
    elif code1 == '1' or code1 == '':
        code1 = '1'
        break
    elif code1 == '0':
        break
    elif code1 == 'cs':
        print("退出")
        input("按任意键结束")
        sys.exit()

print("0 智能判断压缩方式  1 采用png压缩（若图片非png则调用jpg压缩）")
print("2 转换为WEBP       3 采用jpg模式1压缩（jpegtran）")
print("4 采用jpg模式2压缩（mozcjpeg）（默认调用）")
while True:
    model1 = input("选择压缩模式(默认代码4)：")
    if model1 != '0' and model1 != '1' and model1 != '2' and model1 != '3' and model1 != '4' and model1 != '' and model1 != 'cs':
        print("请输入正确的代码")
    elif model1 == '4' or model1 == '':
        model1 = '4'
        break
    elif model1 == '1' or model1 != '2' or model1 != '3' or model1 != '0':
        break
    elif model1 == 'cs':
        print("退出")
        input("按任意键结束")
        sys.exit()

while True:
    workstemp = input("输入压缩线程数（默认代码12,值小于32）：")
    if not str.isdigit(workstemp) and workstemp != '':
        if int(workstemp) > 32:
            print("请输入正确的数字")
    elif workstemp == '':
        work_num = 12
        break
    elif workstemp == 'cs':
        print("退出")
        input("按任意键结束")
        sys.exit()
    else:
        work_num = int(workstemp)
        break

while True:
    code2temp = input("是否删除原图片文件（y,n）（不删除不打包）（默认代码y）：")
    if code2temp != 'y' and code2temp != 'n' and code2temp != '' and code2temp != 'cs':
        print("请输入正确的代码")
    elif code2temp == 'y' or code2temp == '':
        code2 = 1
        break
    elif code2temp == 'n':
        code2 = 0
        break
    elif code2temp == 'cs':
        print("退出")
        input("按任意键结束")
        sys.exit()
code3 = 1
if code2 == 1:
    # code3=0 zip打包
    # code3=1 不打包
    while True:
        code3temp = input("是否进行打包（y,n）（默认代码y）：")
        if code3temp != 'y' and code3temp != 'n' and code3temp != '' and code3temp != 'cs':
            print("请输入正确的代码")
        elif code3temp == 'y' or code3temp == '':
            code3 = 0
            break
        elif code3temp == 'n':
            code3 = 1
            break
        elif code3temp == 'cs':
            print("退出")
            input("按任意键结束")
            sys.exit()
# 输入结束

# 解压处理
st = time.time()
zip_list1 = find_zip(mydir)  # 寻找压缩包
print("已找到 " + str(len(zip_list1)) + " 个待处理压缩包")
if len(zip_list1) != 0:
    print("开始进行解压")
    zip_numbers = unzips(zip_list1, code=int(code1), flag=scflag)  # 解压
    end = time.time()
    print("共" + str(len(zip_list1)) + " 个待处理压缩包；" + "已解压 " + str(zip_numbers) +
          " 个待处理压缩包" + ";共耗时 " + str(format(end - st, '.1f')) + ' s')
# 解压结束

# 进入图片处理
st = time.time()
compressSize = 0
move_count = move_and_rename_parent_directory_images(mydir)  # 先简单进行处理
print(f"总共进行了 {move_count} 次移动和重命名操作。")

dir_list_root = find_folders_with_images(mydir, had_ysdir)
dir_list_nroot = search_image_folders2(mydir, dir_list_root, had_ysdir)


if len(dir_list_root) != 0:
    compressSize = ys_rootdir(dir_list1=dir_list_root, rootdir=mydir, model=int(
        model1), works=int(work_num), code=int(code2), flag=scflag, code2=code3)
# 压缩最底层文件夹图片
if len(dir_list_nroot) != 0:
    compressSize += ys_nrootdir(dir_list1=dir_list_nroot, model=int(
        model1), works=int(work_num), code=int(code2), flag=scflag)
# 压缩上一层级文件夹内部图片


end = time.time()
print("全部图片处理完毕，共耗时 " + str(datetime.timedelta(seconds=end - st)))
print("全部图片共节省空间 " + str(StrOfSize(compressSize *1024 *1024)))
input("按任意键结束")

# pyinstaller --hidden-import=queue -i   favicon.ico -F main.py
# 123