import os
import glob
import soundfile as sf
from tqdm import tqdm

def get_audio_info(file_path):
    """
    获取音频的时长和采样率
    """
    try:
        # sf.info 读取头部信息，速度快，不需要加载整个音频
        info = sf.info(file_path)
        return info.duration, info.samplerate
    except Exception as e:
        print(f"[Warning] Error reading {file_path}: {e}")
        return None, None

def process_seedtts_data(root_dir, output_file):
    # 用于去重，存储已经处理过的 (文件绝对路径)
    seen_paths = set()
    
    # 结果列表
    results = []
    
    # 定义需要遍历的子目录 (en, zh)
    sub_dirs = ['en', 'zh']
    
    print(f"Scanning root directory: {root_dir}")

    # 1. 收集所有任务文件 (.lst)
    lst_files = []
    for lang in sub_dirs:
        lang_dir = os.path.join(root_dir, lang)
        if not os.path.exists(lang_dir):
            continue
        
        # 查找该目录下所有的 .lst 文件
        curr_lst_files = glob.glob(os.path.join(lang_dir, "*.lst"))
        lst_files.extend(curr_lst_files)

    print(f"Found {len(lst_files)} .lst files.")

    # 2. 遍历处理每个 .lst 文件
    for lst_path in tqdm(lst_files, desc="Processing LST files"):
        # 获取当前语言目录的根路径 (例如 .../ntu_enzh_speech/en)
        # 假设 .lst 文件直接在语言目录下
        base_dir = os.path.dirname(lst_path)
        
        with open(lst_path, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                
                parts = line.split('|')
                if len(parts) < 4:
                    continue
                
                # 解析 .lst 结构
                # 格式: id | prompt_text | prompt_wav_rel_path | target_text
                utt_id = parts[0]
                prompt_text = parts[1]
                prompt_rel_path = parts[2]
                target_text = parts[3]

                # -------------------------------------------------
                # 1. 处理 Prompt 音频
                # -------------------------------------------------
                # 拼接绝对路径
                prompt_abs_path = os.path.join(base_dir, prompt_rel_path)
                
                if prompt_abs_path not in seen_paths:
                    if os.path.exists(prompt_abs_path):
                        dur, sr = get_audio_info(prompt_abs_path)
                        if dur is not None:
                            results.append(f"{prompt_abs_path}\t{prompt_text}\t{dur}\t{sr}")
                            seen_paths.add(prompt_abs_path)
                    else:
                        # 仅在第一次遇到缺失时打印警告
                        # print(f"[Missing] Prompt not found: {prompt_abs_path}")
                        pass

                # -------------------------------------------------
                # 2. 处理 Target 音频 (wavs 目录下)
                # -------------------------------------------------
                # 推断 Target 路径: 通常在 wavs/ 目录下，文件名就是 id.wav
                target_abs_path = os.path.join(base_dir, "wavs", f"{utt_id}.wav")
                
                if target_abs_path not in seen_paths:
                    if os.path.exists(target_abs_path):
                        dur, sr = get_audio_info(target_abs_path)
                        if dur is not None:
                            results.append(f"{target_abs_path}\t{target_text}\t{dur}\t{sr}")
                            seen_paths.add(target_abs_path)
                    else:
                        # 尝试另一种命名逻辑 (防御性编程)
                        # 有时候 id 是 A-B，target 文件名可能是 B.wav? 
                        # 但根据 seedtts 惯例通常就是 id.wav，这里保留上述逻辑
                        # print(f"[Missing] Target not found: {target_abs_path}")
                        pass

    # 3. 写入结果
    print(f"Writing {len(results)} unique entries to {output_file}...")
    with open(output_file, 'w', encoding='utf-8') as f_out:
        for line in results:
            f_out.write(line + "\n")
    
    print("Done.")

if __name__ == "__main__":
    # 配置输入根目录
    ROOT_DIR = "/workspace/echoloc/datas/ntu_enzh_speech/"
    
    # 配置输出文件路径
    OUTPUT_FILE = "/workspace/echoloc/datas/ntu_enzh_speech/all_audio_meta.tsv"
    
    process_seedtts_data(ROOT_DIR, OUTPUT_FILE)