import os
import shutil
import random


def get_all_dirs_with_files(root_dir):
    """递归获取 root_dir 下所有包含文件的目录"""
    result = []
    for dirpath, dirnames, filenames in os.walk(root_dir):
        if filenames:
            result.append(dirpath)
    return result


def process_directory(dir_path, seq_num, all_pairs):
    """处理单个目录：修改文本文件第一列为序列号，BMP 复制对应的 txt"""
    all_files = os.listdir(dir_path)

    bmp_files = [f for f in all_files if f.lower().endswith('.bmp')]
    txt_files = [f for f in all_files if f.lower().endswith('.txt')]
    other_files = [f for f in all_files
                   if not f.lower().endswith('.bmp') and not f.lower().endswith('.txt')]

    text_candidates = txt_files + other_files

    if not text_candidates:
        print(f"    [{os.path.basename(dir_path)}] 未找到文本文件，跳过")
        return

    text_filename = text_candidates[0]
    text_path = os.path.join(dir_path, text_filename)

    if not os.path.isfile(text_path):
        print(f"    [{os.path.basename(dir_path)}] 跳过非文件: {text_filename}")
        return

    short_name = os.path.basename(dir_path)
    print(f"    处理: {short_name} (seq={seq_num}), 文本: {text_filename}")

    # 读取文件
    try:
        with open(text_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
    except (UnicodeDecodeError, PermissionError, OSError):
        try:
            with open(text_path, 'r', encoding='gbk') as f:
                lines = f.readlines()
        except Exception as e:
            print(f"    读取失败: {e}")
            return

    modified_lines = []
    for line in lines:
        line = line.rstrip('\n').rstrip('\r')
        if not line.strip():
            modified_lines.append(line)
            continue
        parts = line.split()
        if parts:
            parts[0] = str(seq_num)
        modified_lines.append(' '.join(parts))

    # 写回
    with open(text_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(modified_lines))
    print(f"    已修改 {len([l for l in modified_lines if l.strip()])} 行")

    # 为每个 BMP 复制/关联文档
    if bmp_files:
        for bmp_filename in bmp_files:
            bmp_basename = os.path.splitext(bmp_filename)[0]
            copy_filename = f"{bmp_basename}.txt"
            copy_path = os.path.join(dir_path, copy_filename)

            if os.path.normpath(text_path) == os.path.normpath(copy_path):
                all_pairs.append((os.path.join(dir_path, bmp_filename), copy_path))
            else:
                shutil.copy2(text_path, copy_path)
                all_pairs.append((os.path.join(dir_path, bmp_filename), copy_path))
        print(f"    BMP 数量: {len(bmp_files)}")
    else:
        print(f"    无 BMP 图片")


def main():
    base_dir = r"D:\xwechat_files\wxid_ttkf0xgmyihv22_2388\msg\file\2026-07\抓图(1)\抓图\新建文件夹\1\1\新建文件夹"

    output_dir = os.path.join(os.path.dirname(base_dir.rstrip('\\/')), "yolo_output")
    train_img_dir = os.path.join(output_dir, "train", "images")
    train_lbl_dir = os.path.join(output_dir, "train", "labels")
    val_img_dir   = os.path.join(output_dir, "val", "images")
    val_lbl_dir   = os.path.join(output_dir, "val", "labels")

    for d in [train_img_dir, train_lbl_dir, val_img_dir, val_lbl_dir]:
        os.makedirs(d, exist_ok=True)

    if not os.path.isdir(base_dir):
        print(f"错误: 路径不存在: {base_dir}")
        return

    top_dirs = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d))]
    top_dirs.sort()

    if not top_dirs:
        print("没有找到子文件夹")
        return

    seq_start = 20
    all_pairs = []  # (bmp_path, txt_path)
    mapping = {}     # 顶层文件夹名 -> 序列号

    for idx, top_name in enumerate(top_dirs):
        top_path = os.path.join(base_dir, top_name)
        seq_num = seq_start + idx
        mapping[top_name] = seq_num

        print(f"\n{'='*50}")
        print(f"顶层文件夹: {top_name}  ->  序列号: {seq_num}")

        # 递归获取该顶层文件夹下所有含文件的目录
        leaf_dirs = get_all_dirs_with_files(top_path)

        if not leaf_dirs:
            print(f"  无文件，跳过")
            continue

        for leaf_dir in leaf_dirs:
            process_directory(leaf_dir, seq_num, all_pairs)

    # ---- 输出对应关系 ----
    print("\n" + "=" * 50)
    print("文件夹与序列号对应关系:")
    print("-" * 40)
    for name, num in mapping.items():
        print(f"  {num}  :  {name}")
    print("=" * 50)

    # ---- YOLO 格式 train/val 划分 ----
    if not all_pairs:
        print("\n没有 BMP 图片，跳过 YOLO 输出")
        return

    random.seed(42)
    random.shuffle(all_pairs)

    split_idx = int(len(all_pairs) * 0.8)
    train_pairs = all_pairs[:split_idx]
    val_pairs   = all_pairs[split_idx:]

    print(f"\n总图片数: {len(all_pairs)}")
    print(f"train: {len(train_pairs)} | val: {len(val_pairs)}")

    def copy_pairs(pairs, img_dir, lbl_dir):
        for bmp_path, txt_path in pairs:
            bmp_name = os.path.basename(bmp_path)
            txt_name = os.path.basename(txt_path)

            dst_bmp = bmp_name
            dst_txt = txt_name

            counter = 0
            while os.path.exists(os.path.join(img_dir, dst_bmp)):
                counter += 1
                base, ext = os.path.splitext(bmp_name)
                dst_bmp = f"{base}_{counter}{ext}"
                dst_txt = f"{base}_{counter}.txt"

            shutil.copy2(bmp_path, os.path.join(img_dir, dst_bmp))
            shutil.copy2(txt_path, os.path.join(lbl_dir, dst_txt))

    copy_pairs(train_pairs, train_img_dir, train_lbl_dir)
    copy_pairs(val_pairs, val_img_dir, val_lbl_dir)

    print(f"\nYOLO 格式输出: {output_dir}")
    print(f"  train/images/  ({len(train_pairs)} 张)")
    print(f"  train/labels/  ({len(train_pairs)} 个)")
    print(f"  val/images/    ({len(val_pairs)} 张)")
    print(f"  val/labels/    ({len(val_pairs)} 个)")
    print("\n全部完成!")


if __name__ == "__main__":
    main()