import json
import os
from pathlib import Path
from tqdm import tqdm

def split_jsonl_to_jsons(jsonl_path, old_root, new_root, evaluator):
    """
    解析 JSONL 文件，替换路径，在文件名后追加 evaluator，并保存为独立的 JSON 文件
    """
    jsonl_file = Path(jsonl_path)
    if not jsonl_file.exists():
        print(f"❌ 找不到 JSONL 文件: {jsonl_file}")
        return

    # 确保根目录是字符串格式
    old_root = str(old_root)
    new_root = str(new_root)

    success_count = 0
    error_count = 0

    print("开始处理数据...")

    with open(jsonl_file, 'r', encoding='utf-8') as f:
        for line_num, line in tqdm(enumerate(f, 1)):
            line = line.strip()
            if not line:
                continue

            try:
                # 解析单行 JSON
                data = json.loads(line)
                audio_path_str = data.get('audio_path')

                if not audio_path_str:
                    print(f"⚠️ 第 {line_num} 行: 未找到 'audio_path' 字段，已跳过。")
                    error_count += 1
                    continue

                # 替换旧的根目录路径
                if old_root in audio_path_str:
                    new_audio_path_str = audio_path_str.replace(old_root, new_root)
                    # 同步更新 JSON 内部的 audio_path 字段为新路径
                    data['audio_path'] = new_audio_path_str
                else:
                    print(f"⚠️ 第 {line_num} 行: 'audio_path' 中不包含指定的旧根目录，已跳过 ({audio_path_str})")
                    error_count += 1
                    continue

                # 构建新的保存路径：拼接 evaluator 并修改后缀为 .json
                base_path = Path(new_audio_path_str)
                # base_path.stem 获取不带后缀的文件名 (例如: 884_indextts_spk-zh_emo-en_aiden)
                new_file_name = f"{base_path.stem}_{evaluator}.json"
                new_file_path = base_path.with_name(new_file_name)
                
                # 如果你想在控制台查看生成的路径，可以取消下面的注释
                # print(new_file_path)

                # 如果对应的子目录 (比如 /884/) 不存在，则自动创建
                new_file_path.parent.mkdir(parents=True, exist_ok=True)

                # 写入独立的 JSON 文件
                with open(new_file_path, 'w', encoding='utf-8') as out_f:
                    json.dump(data, out_f, ensure_ascii=False, indent=4)

                success_count += 1

            except json.JSONDecodeError:
                print(f"❌ 第 {line_num} 行: JSON 格式错误，无法解析。")
                error_count += 1
            except Exception as e:
                print(f"❌ 第 {line_num} 行: 发生意外错误 -> {e}")
                error_count += 1

    print("\n✅ 处理完成！")
    print(f"成功生成 JSON 文件数: {success_count}")
    if error_count > 0:
        print(f"跳过的异常数据数: {error_count}")


if __name__ == "__main__":
    # ================= 配置区域 =================
    
    evaluator = "geminiflash"
    
    # 你的 jsonl 文件绝对路径
    JSONL_FILE = "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/ds3_vstyle_controls_3_17/ds3_vstyle_controls_3_17_geminiflash_evaled.jsonl"
    
    # 需要被替换的旧的根目录
    OLD_ROOT = "/workspace/echoloc/datas/huawei_gqs/ds3_vstyle_controls_3_17"
    
    # 修改为你新电脑上的实际根目录路径
    NEW_ROOT = "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/indextts/ds3_vstyle_controls_3_17" 
    
    # ============================================

    # 记得把 evaluator 作为参数传进去
    split_jsonl_to_jsons(JSONL_FILE, OLD_ROOT, NEW_ROOT, evaluator)