#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import struct
import zlib
import time
import mmap
import shutil
import random  # ✅ 修复：补上random模块导入
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from collections import namedtuple
import multiprocessing

# ====================== 1:1复刻截图配色 =======================
class Color:
    YELLOW      = "\033[93m"
    CYAN        = "\033[96m"
    GREEN       = "\033[92m"
    RED         = "\033[91m"
    WHITE       = "\033[97m"
    RESET       = "\033[0m"

# ====================== 新增：自动复制PAK配置 + 逐字彩色 =======================
SRC_PAK_DIR = Path("/storage/emulated/0/勿辞制作区/提取的pak/人物范围pak")
DST_PAK_DIR = Path("/storage/emulated/0/勿辞制作区/pak")

COLOR_POOL = [
    "\033[1;31m", "\033[1;32m", "\033[1;33m", "\033[1;34m",
    "\033[1;35m", "\033[1;36m", "\033[1;91m", "\033[1;92m",
    "\033[1;93m", "\033[1;94m", "\033[1;95m", "\033[1;96m"
]
RESET = "\033[0m"
COLOR_IDX = 0

def print_rand(text):
    global COLOR_IDX
    out = ""
    i = 0
    while i < len(text):
        if text[i:i+2] == "\033[":
            j = i
            while j < len(text) and text[j] != 'm':
                j += 1
            if j < len(text):
                out += text[i:j+1]
                i = j + 1
                continue
        ch = text[i]
        if ch.isspace():
            out += ch
            i += 1
            continue
        c = COLOR_POOL[COLOR_IDX % len(COLOR_POOL)]
        out += c + ch + RESET
        COLOR_IDX += 1
        i += 1
    print(out)

def file_bin_equal(file1: Path, file2: Path) -> bool:
    if not file1.exists() or not file2.exists():
        return False
    try:
        with open(file1, 'rb') as f1, open(file2, 'rb') as f2:
            return f1.read() == f2.read()
    except:
        return False

def auto_copy_pak():
    DST_PAK_DIR.mkdir(parents=True, exist_ok=True)
    if not SRC_PAK_DIR.exists():
        print_rand("源目录不存在：提取的pak/人物范围pak")
        return
    src_paks = list(SRC_PAK_DIR.glob("*.pak"))
    if not src_paks:
        print_rand("源目录无任何pak文件")
        return
    copy_num = 0
    skip_num = 0
    cover_num = 0
    print_rand("开始比对并同步PAK文件...")
    for src_pak in src_paks:
        pak_name = src_pak.name
        dst_pak = DST_PAK_DIR / pak_name
        if dst_pak.exists():
            if file_bin_equal(src_pak, dst_pak):
                print_rand(f"文件完全一致，跳过：{pak_name}")
                skip_num += 1
            else:
                shutil.copy2(src_pak, dst_pak)
                print_rand(f"文件名相同内容不同，已覆盖：{pak_name}")
                cover_num += 1
        else:
            shutil.copy2(src_pak, dst_pak)
            print_rand(f"已复制新PAK：{pak_name}")
            copy_num += 1
    print_rand(f"同步完成：新增复制{copy_num}个 | 同名跳过{skip_num}个 | 内容不同覆盖{cover_num}个")
    time.sleep(0.8)

# ====================== 全局路径配置 =======================
PAK_DIR = "/storage/emulated/0/勿辞制作区/pak/"
UNPACK_DIR = "/storage/emulated/0/勿辞制作区/uexp解包"
PACK_DIR = "/storage/emulated/0/勿辞制作区/uexp打包"
ENCRYPT_KEY = 0x79
MAX_WORKERS = min(multiprocessing.cpu_count(), 8)

# ====================== 分隔线 =======================
LINE_DIVIDER = f"{Color.YELLOW}========================================{Color.RESET}"
SHORT_DIVIDER = f"{Color.YELLOW}----------------------------------------{Color.RESET}"

# ====================== 目标文件列表 =======================
FLASH_FILES = {
    "M_Shanguangdan.uasset",
    "M_Concussive_PostProcess_02_Inst.uasset",
    "M_Screen_White.uasset",
    "M_Concussive_PostProcess_02.uasset",
    "M_Screen_D_01.uasset",
    "M_Shanguangdan2.uasset"
}
SMOKE_FILES = {
    "P_smoke_vehicle_dirt_01.uasset", "M_NearDeath_001.uasset",
    "P_smoke_vehicle_Grass_01.uasset", "M_Screen_D_01_Inst.uasset",
    "P_smoke_vehicle_sand_01.uasset", "P_smoke_Lobby_LIAOJI.uasset",
    "P_MF_SpentCartridge_45mm_nosmoke.uasset", "M_MiniMapFX_01_Inst.uasset",
    "M_Smoke_PostProcess_SwitchWeather.uasset", "P_BJ_smoke_01.uasset",
    "M_Smoke_09.uasset", "P_smoke_vehicle_Exhaust_03.uasset",
    "M_smoke_vehicle_dirt_01.uasset", "M_ScoutUAV_Smoke_01.uasset",
    "P_Dead_Smoke_10.uasset", "P_Dead_Smoke_02.uasset",
    "P_smoke_vehicle_snow_03.uasset", "M_red_Smoke_02.uasset",
    "P_zombiebox_dead_smoke_01.uasset", "M_Smoke_14.uasset",
    "P_Smoke_RadarTower_01.uasset", "P_Grenade_Smoke_test3.uasset",
    "P_Red_Smoke_01.uasset", "P_Smoke_Gun_01.uasset",
    "M_Smoke_27.uasset", "M_Smoke_28_002.uasset",
    "P_Dead_Smoke_03.uasset", "M_red_Smoke_01.uasset",
    "P_smoke_vehicle_Exhaust_Lobby_01.uasset", "M_Smoke_32.uasset",
    "M_Smoke_30_001.uasset", "M_smoke_logo.uasset",
    "M_NearDeath_PostProcess_02_Inst.uasset", "M_Screen_black.uasset",
    "P_MF_SpentCartridge_45mm_nosmokeFFP.uasset", "P_Grenade_Smoke.uasset",
    "M_Screen_D_01_w_Pain_Inst1.uasset", "M_MiniMapFX_01.uasset",
    "M_smoketrail_Int.uasset", "M_Lobby_fx_01.uasset",
    "M_Blizzard_SM_01.uasset", "P_GlideSmoke_011_show.uasset",
    "M_MF_Smoke_02.uasset", "M_Smoke_30.uasset",
    "M_smoke_katong_02.uasset", "P_smoke_aircraft_09.uasset",
    "M_Blizzard_SM_ceshi.uasset", "M_NearDeath_001_Inst.uasset",
    "P_smoke_vehicle_road_01.uasset", "M_Screen_D_01_LOD.uasset",
    "P_smoke_vehicle_Lobby_snow_01.uasset", "M_Smoke_35.uasset",
    "M_Smoke_050_01.uasset", "P_smoke_vehicle_Lobby_01.uasset",
    "P_smoke_vehicle_sand_02.uasset", "Smoke_011_tex.uasset",
    "M_Smoke_31.uasset", "P_zombiebox_Destory_smoke_01.uasset",
    "M_smoke_djf03.uasset", "M_NearDeath_PostProcess_02.uasset",
    "M_Smoke_052.uasset", "M_Blizzard_SM_main.uasset",
    "P_yellow_Smoke_02.uasset", "M_MF_Smoke_01.uasset",
    "P_Shooting_park_Chick_smoke_01.uasset", "M_Smoke_30_002.uasset",
    "M_Smoke_PostProcess_01.uasset", "M_Smoke_25.uasset",
    "M_Tyre_vehicle_Explosion_Smoke_001.uasset", "P_White_Smoke_02.uasset",
    "M_smoke_katong_01.uasset", "P_smoke_03.uasset",
    "M_smoketrail_Int_Inst10.uasset"
}

# ====================== 工具函数 =======================
class Utils:
    @staticmethod
    def read_file(filepath):
        try:
            with open(filepath, 'rb') as f:
                return f.read()
        except Exception:
            return b''
    @staticmethod
    def decrypt_data(data, encrypt_flag):
        if encrypt_flag:
            return bytes([b ^ ENCRYPT_KEY for b in data])
        return data
    @staticmethod
    def encrypt_data(data, encrypt_flag):
        return Utils.decrypt_data(data, encrypt_flag)
    @staticmethod
    def compress_data(data, max_size):
        for level in range(9, 4, -1):
            try:
                compressed = zlib.compress(data, level)
                if len(compressed) <= max_size:
                    return compressed
            except Exception:
                continue
        return data.ljust(max_size, b'\x00')
    @staticmethod
    def decompress_data(data):
        try:
            return zlib.decompress(data)
        except Exception:
            return data
    @staticmethod
    def clear_pack_unpack_dirs():
        for dir_path in [UNPACK_DIR, PACK_DIR]:
            try:
                if Path(dir_path).exists():
                    shutil.rmtree(dir_path, ignore_errors=True)
                Path(dir_path).mkdir(parents=True, exist_ok=True)
            except Exception:
                pass

# ====================== PAK解包核心类 =======================
CompressionInfo = namedtuple('CompressionInfo', [
    'offset', 'size', 'zip', 'zsize', 'encrypted', 'chunks', 'chunk_size'
])
class FastPakExtractor:
    def __init__(self):
        self.encrypt = 0
        self.mm = None
        self.selected_pak = None
        self.compression_info = {}
        self.base_path = "../../../"
        self.file_success = 0
        self.file_fail = 0
        self.target_files = set()
        self.valid_targets = {}
        self.file_size = 0

    def get_pak_files(self):
        pak_path = Path(PAK_DIR)
        if not pak_path.exists():
            print_rand(f"[错误] PAK目录不存在: {PAK_DIR}")
            return []
        return sorted(list(pak_path.glob("*.pak")), key=lambda x: x.stat().st_size, reverse=True)

    # ✅ 全自动选择PAK，无需手动输入
    def select_pak_file(self, pak_files):
        if pak_files:
            self.selected_pak = pak_files[0]
            print_rand(f"✅ 自动选择：{self.selected_pak.name}")
            return True
        return False

    def find_magic_offset(self):
        pattern1 = b"\x2E\x2E\x2F\x2E\x2E\x2F\x2E\x2E\x2F"
        pattern2 = b"\x57\x57\x56\x57\x57\x56\x57\x57\x56"
        pattern3 = b"\x2E\x2E\x2F"
        scan_len = min(5 * 1024 * 1024, self.file_size)
        data = self.mm[-scan_len:]
        offset1 = data.find(pattern1)
        offset2 = data.find(pattern2)
        offset3 = data.rfind(pattern3)
        if offset1 != -1:
            self.encrypt = 0
            return self.file_size - scan_len + offset1 - 4
        elif offset2 != -1:
            self.encrypt = 1
            return self.file_size - scan_len + offset2 - 4
        elif offset3 != -1:
            self.encrypt = 0
            return self.file_size - scan_len + offset3 - 4
        else:
            return self.file_size - 0x2C

    def get_base_path(self, offset):
        if offset + 4 > self.file_size:
            return "../../../", offset + 4
        try:
            name_size = struct.unpack('<I', self.mm[offset:offset+4])[0]
        except Exception:
            return "../../../", offset + 4
        if name_size == 0:
            return "../../../", offset + 4
        if name_size > 1024 or offset + 4 + name_size > self.file_size:
            return "../../../", offset + 4
        try:
            base_path = self.mm[offset+4:offset+4+name_size].decode('utf-8', errors='ignore').rstrip('\x00')
            if name_size != 0x0A and name_size < 0xFF:
                base_path = "../../../" + base_path
            self.base_path = base_path
            return base_path, offset + 4 + name_size
        except Exception:
            return "../../../", offset + 4 + name_size

    def parse_file_entry(self, offset):
        entry_start = offset
        if offset + 20 + 49 > self.file_size:
            return None, offset + 69
        try:
            hash_data = self.mm[offset:offset+20]
            offset += 20
            entry_data = self.mm[offset:offset+49]
            offset += 49
            entry = {
                'hash': hash_data,
                'offset': struct.unpack('<Q', entry_data[0:8])[0],
                'size': struct.unpack('<Q', entry_data[8:16])[0],
                'zip': struct.unpack('<I', entry_data[16:20])[0],
                'zsize': struct.unpack('<Q', entry_data[20:28])[0],
                'chunks': [],
                'chunk_size': 0x10000,
                'encrypted': 0
            }
            if entry['zip'] != 0:
                if offset + 4 > self.file_size:
                    return entry, offset
                chunk_count = struct.unpack('<I', self.mm[offset:offset+4])[0]
                offset += 4
                chunk_count = min(chunk_count, 1000)
                for _ in range(chunk_count):
                    if offset + 16 > self.file_size:
                        break
                    chunk_data = self.mm[offset:offset+16]
                    offset += 16
                    chunk_offset = struct.unpack('<Q', chunk_data[0:8])[0]
                    chunk_end = struct.unpack('<Q', chunk_data[8:16])[0]
                    entry['chunks'].append((chunk_offset, chunk_end))
            if offset + 5 > self.file_size:
                return entry, offset
            chunk_size_data = self.mm[offset:offset+5]
            offset += 5
            entry['chunk_size'] = struct.unpack('<I', chunk_size_data[0:4])[0]
            entry['encrypted'] = chunk_size_data[4]
            if self.encrypt:
                try:
                    encrypted_entry = self.mm[entry_start:offset]
                    decrypted_entry = self.decrypt_data(encrypted_entry)
                    if len(decrypted_entry) >= 69:
                        entry['offset'] = struct.unpack('<Q', decrypted_entry[20:28])[0]
                        entry['size'] = struct.unpack('<Q', decrypted_entry[28:36])[0]
                        entry['zip'] = struct.unpack('<I', decrypted_entry[36:40])[0]
                        entry['zsize'] = struct.unpack('<Q', decrypted_entry[40:48])[0]
                except Exception:
                    pass
            return entry, offset
        except Exception:
            return None, offset + 69

    def decrypt_data(self, data):
        if not self.encrypt:
            return data
        return bytes([b ^ ENCRYPT_KEY for b in data])

    def extract_file(self, entry, output_path):
        try:
            if entry['chunks']:
                total_size = entry['size']
                remaining = total_size
                output_path.parent.mkdir(parents=True, exist_ok=True)
                with open(output_path, 'wb') as out:
                    for chunk_offset, chunk_end in entry['chunks']:
                        if chunk_offset + (chunk_end - chunk_offset) > self.file_size:
                            continue
                        chunk_zsize = chunk_end - chunk_offset
                        chunk_data = self.mm[chunk_offset:chunk_offset+chunk_zsize]
                        if self.encrypt and entry['encrypted'] == 1:
                            chunk_data = self.decrypt_data(chunk_data)
                        if entry['zip'] != 0:
                            chunk_data = Utils.decompress_data(chunk_data)
                        write_size = min(len(chunk_data), remaining)
                        out.write(chunk_data[:write_size])
                        remaining -= write_size
                        if remaining <= 0:
                            break
            else:
                if entry['offset'] + entry['zsize'] > self.file_size:
                    return False
                if entry['zip'] != 0:
                    compressed_data = self.mm[entry['offset']:entry['offset']+entry['zsize']]
                    if self.encrypt and entry['encrypted'] == 1:
                        compressed_data = self.decrypt_data(compressed_data)
                    data = Utils.decompress_data(compressed_data)
                else:
                    data = self.mm[entry['offset']:entry['offset']+entry['size']]
                    if self.encrypt and entry['encrypted'] == 1:
                        data = self.decrypt_data(data)
                output_path.parent.mkdir(parents=True, exist_ok=True)
                with open(output_path, 'wb') as f:
                    f.write(data)
            self.file_success += 1
            return True
        except Exception:
            self.file_fail += 1
            return False

    def parse_toc(self, toc_offset, toc_size):
        toc_data = self.mm[toc_offset:toc_offset+toc_size]
        toc_len = len(toc_data)
        pos = 0
        entries = []
        stack = [(1, 0)]
        target_set = self.target_files
        while stack:
            if pos + 8 > toc_len:
                break
            flag, count = stack.pop()
            if flag == 1:
                try:
                    dir_count = struct.unpack('<Q', toc_data[pos:pos+8])[0]
                    pos += 8
                    stack.append((0, dir_count))
                except Exception:
                    break
                continue
            if count == 0:
                continue
            count -= 1
            if pos + 4 > toc_len:
                break
            try:
                name_size = struct.unpack('<i', toc_data[pos:pos+4])[0]
                pos += 4
            except Exception:
                break
            dir_name = ""
            if name_size >= 0:
                if pos + name_size > toc_len:
                    pos += name_size
                    continue
                try:
                    dir_name = toc_data[pos:pos+name_size].decode('utf-8', errors='ignore').rstrip('\x00')
                except Exception:
                    pass
                pos += name_size
            else:
                abs_name_size = abs(name_size) * 2
                if pos + abs_name_size > toc_len:
                    pos += abs_name_size
                    continue
                try:
                    dir_name = toc_data[pos:pos+abs_name_size].decode('utf-16-le', errors='ignore').rstrip('\x00')
                except Exception:
                    pass
                pos += abs_name_size
            if pos + 8 > toc_len:
                break
            try:
                file_count = struct.unpack('<Q', toc_data[pos:pos+8])[0]
                pos += 8
            except Exception:
                file_count = 0
            if file_count == 0:
                stack.append((0, count))
                continue
            file_count = min(file_count, 100000)
            for _ in range(file_count):
                if pos + 4 > toc_len:
                    break
                try:
                    name_size = struct.unpack('<i', toc_data[pos:pos+4])[0]
                    pos += 4
                except Exception:
                    break
                file_name = ""
                if name_size > 0:
                    if pos + name_size > toc_len:
                        pos += name_size
                        continue
                    try:
                        file_name = toc_data[pos:pos+name_size].decode('utf-8', errors='ignore').rstrip('\x00')
                    except Exception:
                        pass
                    pos += name_size
                else:
                    abs_name_size = abs(name_size) * 2
                    if pos + abs_name_size > toc_len:
                        pos += abs_name_size
                        continue
                    try:
                        file_name = toc_data[pos:pos+abs_name_size].decode('utf-16-le', errors='ignore').rstrip('\x00')
                    except Exception:
                        pass
                    pos += abs_name_size
                full_path = f"{dir_name}{file_name}"
                if pos + 4 > toc_len:
                    break
                try:
                    entry_index = struct.unpack('<I', toc_data[pos:pos+4])[0]
                    pos += 4
                except Exception:
                    break
                if file_name.lower() in target_set:
                    entries.append((entry_index, full_path, file_name))
            if count > 0:
                stack.append((0, count))
        return entries

    def unpack_pak(self, target_set):
        if not self.selected_pak:
            print_rand("[错误] 未选择PAK文件")
            return False
        start_time = time.time()
        self.target_files = {f.lower() for f in target_set}
        self.file_success = 0
        self.file_fail = 0
        self.valid_targets = {}
        self.compression_info = {}
        self.file_size = self.selected_pak.stat().st_size
        try:
            with open(self.selected_pak, 'rb') as f:
                self.mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
            magic_offset = self.find_magic_offset()
            base_path, pos = self.get_base_path(magic_offset)
            if pos + 4 > self.file_size:
                file_count = 0
            else:
                try:
                    file_count = struct.unpack('<I', self.mm[pos:pos+4])[0]
                    pos += 4
                except Exception:
                    file_count = 0
            file_count = min(file_count, 100000)
            entries = []
            for _ in range(file_count):
                entry, pos = self.parse_file_entry(pos)
                if entry:
                    entries.append(entry)
                else:
                    break
            if pos + 8 <= self.file_size:
                try:
                    struct.unpack('<Q', self.mm[pos:pos+8])[0]
                    pos += 8
                except Exception:
                    pass
            if self.encrypt:
                pos += 1
            toc_offset = pos
            toc_size = self.file_size - toc_offset
            toc_entries = self.parse_toc(toc_offset, toc_size)
            tasks = []
            for entry_idx, full_path, file_name in toc_entries:
                if entry_idx < len(entries):
                    entry = entries[entry_idx].copy()
                    rel_path = full_path
                    if base_path and base_path != "../../../":
                        rel_path = base_path + full_path
                    while rel_path.startswith('../'):
                        rel_path = rel_path[3:]
                    output_path = Path(UNPACK_DIR) / rel_path
                    self.compression_info[full_path] = CompressionInfo(
                        offset=entry['offset'], size=entry['size'], zip=entry['zip'],
                        zsize=entry['zsize'], encrypted=entry['encrypted'],
                        chunks=entry['chunks'], chunk_size=entry['chunk_size']
                    )
                    self.valid_targets[file_name.lower()] = full_path
                    tasks.append((entry, output_path))
            with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
                futures = {executor.submit(self.extract_file, e, o): (e,o) for e,o in tasks}
                for future in as_completed(futures):
                    future.result()
            self.unpack_elapsed = time.time() - start_time
            return True
        except Exception as e:
            print_rand(f"[错误] 解包失败: {str(e)}")
            return False
        finally:
            if self.mm:
                self.mm.close()

# ====================== PAK打包核心类 =======================
class FastPakPacker:
    def __init__(self, extractor):
        self.extractor = extractor
        self.encrypt = extractor.encrypt
        self.file_success = 0
        self.file_fail = 0

    def replace_non_chunked(self, pak_file, info, new_data):
        try:
            if info.zip != 0:
                compressed = Utils.compress_data(new_data, info.zsize)
                compressed = Utils.encrypt_data(compressed, self.encrypt and info.encrypted == 1)
                pak_file.seek(info.offset)
                pak_file.write(compressed.ljust(info.zsize, b'\x00'))
            else:
                write_data = new_data.ljust(info.size, b'\x00')
                write_data = Utils.encrypt_data(write_data, self.encrypt and info.encrypted == 1)
                pak_file.seek(info.offset)
                pak_file.write(write_data)
            self.file_success += 1
            return True
        except Exception:
            self.file_fail += 1
            return False

    def replace_chunked(self, pak_file, info, new_data):
        try:
            data_pos = 0
            total_size = info.size
            for chunk_offset, chunk_end in info.chunks:
                max_size = chunk_end - chunk_offset
                chunk_size = min(info.chunk_size, total_size - data_pos)
                chunk_data = new_data[data_pos:data_pos+chunk_size]
                data_pos += chunk_size
                if info.zip != 0:
                    chunk_data = Utils.compress_data(chunk_data, max_size)
                chunk_data = Utils.encrypt_data(chunk_data, self.encrypt and info.encrypted == 1)
                if len(chunk_data) < max_size:
                    chunk_data = chunk_data.ljust(max_size, b'\x00')
                pak_file.seek(chunk_offset)
                pak_file.write(chunk_data)
                if data_pos >= total_size:
                    break
            self.file_success += 1
            return True
        except Exception:
            self.file_fail += 1
            return False

    def pack_pak(self):
        original_pak = self.extractor.selected_pak
        if not original_pak:
            print_rand("[错误] 未选择PAK文件")
            return False
        start_time = time.time()
        self.file_success = 0
        self.file_fail = 0
        file_map = {}
        for root, _, files in os.walk(PACK_DIR):
            for file in files:
                file_map[file.lower()] = Path(root) / file
        try:
            with open(original_pak, 'r+b') as pak_file:
                for file_name, full_path in self.extractor.valid_targets.items():
                    if file_name not in file_map:
                        self.file_fail += 1
                        continue
                    info = self.extractor.compression_info.get(full_path)
                    if not info:
                        self.file_fail += 1
                        continue
                    new_data = Utils.read_file(file_map[file_name])
                    if not new_data:
                        self.file_fail += 1
                        continue
                    if info.chunks:
                        self.replace_chunked(pak_file, info, new_data)
                    else:
                        self.replace_non_chunked(pak_file, info, new_data)
            self.pack_elapsed = time.time() - start_time
            return True
        except Exception as e:
            print_rand(f"[错误] 打包失败: {str(e)}")
            return False

# ====================== 资源修改函数 =======================
def process_target_files(target_set, process_name):
    print_rand("正在扫描文件...")
    found_files = []
    for root, _, files in os.walk(UNPACK_DIR):
        for file in files:
            if file.lower() in target_set:
                found_files.append(Path(root) / file)
    total = len(found_files)
    if total == 0:
        print_rand("未找到目标文件")
        return 0, 0
    print_rand("正在修改文件...")
    Path(PACK_DIR).mkdir(parents=True, exist_ok=True)
    success = 0
    current = 0
    def process_single(fpath):
        nonlocal success, current
        dest = Path(PACK_DIR) / fpath.name
        try:
            size = fpath.stat().st_size
            with open(dest, 'wb') as f:
                f.truncate(size)
            success += 1
        except Exception:
            pass
        current += 1
        # ✅ 修复：用print_rand打印进度，不直接用random.choice
        print(f"\r{Color.CYAN}进度: {current}/{total}{Color.RESET}", end="")
        sys.stdout.flush()
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = [executor.submit(process_single, f) for f in found_files]
        for future in as_completed(futures):
            future.result()
    print()
    print_rand(f"[{process_name}] 完成")
    return total, success

# ====================== 主流程 =======================
def main():
    os.environ['PYTHONIOENCODING'] = 'utf-8'
    while True:
        print(LINE_DIVIDER)
        print_rand("            勿辞云端-除烟防闪工具")
        print(LINE_DIVIDER)
        print_rand("1. 防闪处理")
        print_rand("2. 除烟处理")
        print_rand("3. 全部都修改")
        print_rand("0. 退出程序")
        print(LINE_DIVIDER)

        choice = input(f"{Color.CYAN}请选择 (1/2/3) : {Color.RESET}").strip()
        if choice == "0":
            print_rand("已退出程序")
            sys.exit(0)
        if choice not in ["1","2","3"]:
            print_rand("输入无效，请输入1-3之间的数字")
            input(f"{Color.CYAN}按回车返回...{Color.RESET}")
            continue

        print_rand("执行中...")
        Utils.clear_pack_unpack_dirs()
        auto_copy_pak()

        if choice == "1":
            target_set = FLASH_FILES
            process_name = "防闪修改"
        elif choice == "2":
            target_set = SMOKE_FILES
            process_name = "除烟修改"
        else:
            target_set = FLASH_FILES.union(SMOKE_FILES)
            process_name = "全部修改"

        print_rand("\n[1/3] 解包:")
        extractor = FastPakExtractor()
        pak_files = extractor.get_pak_files()
        if not pak_files:
            input(f"{Color.CYAN}按回车返回...{Color.RESET}")
            continue
        extractor.select_pak_file(pak_files)

        print_rand("\n[2/3] 解包中...")
        if not extractor.unpack_pak(target_set):
            input(f"{Color.CYAN}按回车返回...{Color.RESET}")
            continue
        print_rand(f"解包耗时: {extractor.unpack_elapsed:.3f}s")

        print_rand("\n[3/3] 修改数据")
        total, success = process_target_files({f.lower() for f in target_set}, process_name)
        print_rand(f"总计: {total} 个")
        print_rand(f"成功: {success} 个")

        print_rand("\n[4/3] 打包")
        print_rand(f"源文件: {extractor.selected_pak.name}")
        packer = FastPakPacker(extractor)
        packer.pack_pak()
        print_rand(f"打包耗时: {packer.pack_elapsed:.3f}s")
        print_rand("√ 全部完成!")
        print(LINE_DIVIDER)

        input(f"{Color.CYAN}按 Enter 退出...{Color.RESET}")
        break

if __name__ == "__main__":
    main()
