# https://t.me/ZGQinc

import os
import shutil
import sys

os.makedirs("input", exist_ok=True)
os.makedirs("temp", exist_ok=True)
os.makedirs("output", exist_ok=True)

input_files = os.listdir("input")
if not input_files:
    print("请将要处理的文件放入 input 文件夹后重新运行本脚本。")
    sys.exit(0)

ctex_exts = [".bmp", ".dds", ".ktx", ".exr", ".hdr", ".jpg", ".jpeg", ".png", ".tga", ".svg", ".webp"]
stex_exts = [".jpg", ".jpeg", ".png", ".bmp", ".svg", ".webp"]

magic_map = {
    "ctex": b"RIFF",
    "stex": b"\x89PNG",
    "oggstr": b"OggS"
}

def find_real_ext(filename, ext_list):
    for ext in ext_list:
        if ext in filename:
            return ext
    return None

def process_file(filepath, filetype):
    filename = os.path.basename(filepath)
    with open(filepath, "rb") as f:
        data = f.read()

    real_ext = None
    header = magic_map.get(filetype)
    if filetype == "ctex":
        real_ext = find_real_ext(filename, ctex_exts)
    elif filetype == "stex":
        real_ext = find_real_ext(filename, stex_exts)
    elif filetype == "oggstr":
        real_ext = ".ogg"
    elif filetype == "mp3str":
        real_ext = ".mp3"

    if not real_ext:
        print(f"未找到扩展名: {filename}")
        return

    if filetype in ["ctex", "stex", "oggstr"]:
        index = data.find(header)
        if index == -1:
            print(f"找不到头部: {filename}")
            return
        data = data[index:]

    temp_path = os.path.join("temp", filename)
    with open(temp_path, "wb") as f:
        f.write(data)

    base = filename.split(real_ext)[0] + real_ext
    if filetype == "ctex":
        base += ".webp"
    output_path = os.path.join("output", base)

    shutil.copy(temp_path, output_path)
    print(f"处理完成: {filename} → {base}")

for fname in input_files:
    fpath = os.path.join("input", fname)
    if not os.path.isfile(fpath):
        continue

    lower = fname.lower()
    if lower.endswith(".ctex"):
        process_file(fpath, "ctex")
    elif lower.endswith(".stex"):
        process_file(fpath, "stex")
    elif lower.endswith(".oggstr"):
        process_file(fpath, "oggstr")
    elif lower.endswith(".mp3str"):
        process_file(fpath, "mp3str")

shutil.rmtree("temp", ignore_errors=True)
print("已删除临时文件夹")