import os
import glob
import torch
import torchaudio
import torchaudio.transforms as T
from tqdm import tqdm

def merge_audio_data(root_dir, output_wav_root, output_tsv_path):
    """
    遍历 root_dir 下的 lst 文件，将 prompt 和 target 音频拼接。
    输出音频到 output_wav_root，输出元数据到 output_tsv_path。
    """
    
    # 1. 准备工作
    sub_dirs = ['en', 'zh']
    silence_duration = 0.1  # 拼接中间插入 0.3秒 静音
    
    # 创建输出目录
    os.makedirs(output_wav_root, exist_ok=True)
    
    tsv_lines = []
    
    print(f"Scanning root directory: {root_dir}")

    # 2. 收集 .lst 文件
    lst_files = []
    for lang in sub_dirs:
        lang_dir = os.path.join(root_dir, lang)
        if not os.path.exists(lang_dir):
            continue
        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.")

    success_count = 0
    
    for lst_path in tqdm(lst_files, desc="Merging Audios"):
        # 获取当前是 en 还是 zh
        # 假设路径结构是 .../en/xxx.lst，所以 dirname 是 .../en，再 basename 就是 en
        lang_category = os.path.basename(os.path.dirname(lst_path))
        
        # 对应的输出子目录
        save_dir = os.path.join(output_wav_root, lang_category)
        os.makedirs(save_dir, exist_ok=True)
        
        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
                
                utt_id = parts[0]
                prompt_text = parts[1]
                prompt_rel_path = parts[2]
                target_text = parts[3]

                # -------------------------------------------------
                # ✅ 文本拼接逻辑修改
                # -------------------------------------------------
                if lang_category == 'zh':
                    # 中文模式：直接拼接，不加空格
                    full_text = f"{prompt_text}{target_text}"
                else:
                    # 英文模式：中间加空格
                    full_text = f"{prompt_text} {target_text}"

                # -------------------------------------------------
                # 获取路径
                # -------------------------------------------------
                prompt_abs_path = os.path.join(base_dir, prompt_rel_path)
                target_abs_path = os.path.join(base_dir, "wavs", f"{utt_id}.wav")
                
                new_filename = f"{utt_id}_merged.wav"
                new_file_path = os.path.join(save_dir, new_filename)
                
                if not (os.path.exists(prompt_abs_path) and os.path.exists(target_abs_path)):
                    continue

                try:
                    # -------------------------------------------------
                    # 音频处理
                    # -------------------------------------------------
                    # load 默认返回 float32 类型
                    wav_p, sr_p = torchaudio.load(prompt_abs_path)
                    wav_t, sr_t = torchaudio.load(target_abs_path)
                    
                    # 统一通道数 (转为单声道)
                    if wav_p.shape[0] > 1: wav_p = torch.mean(wav_p, dim=0, keepdim=True)
                    if wav_t.shape[0] > 1: wav_t = torch.mean(wav_t, dim=0, keepdim=True)

                    # 统一采样率 (以 prompt 的 SR 为基准，通常是 16k, 22k 或 24k)
                    target_sr = sr_p 
                    
                    if sr_t != target_sr:
                        resampler = T.Resample(sr_t, target_sr)
                        wav_t = resampler(wav_t)
                    
                    # 创建静音片段
                    silence_samples = int(silence_duration * target_sr)
                    silence_wav = torch.zeros(1, silence_samples)

                    # 拼接
                    merged_wav = torch.cat([wav_p, silence_wav, wav_t], dim=1)
                    
                    # 保存
                    torchaudio.save(new_file_path, merged_wav, target_sr)
                    
                    # 计算总时长
                    total_duration = merged_wav.shape[1] / target_sr
                    
                    # 记录 (路径 \t 文本 \t 时长 \t 采样率)
                    tsv_lines.append(f"{new_file_path}\t{full_text}\t{total_duration:.4f}\t{target_sr}")
                    success_count += 1
                    
                except Exception as e:
                    print(f"[Error] Failed to process {utt_id}: {e}")
                    continue

    # 4. 写入 TSV
    print(f"Saving metadata to {output_tsv_path}...")
    with open(output_tsv_path, 'w', encoding='utf-8') as f:
        for line in tsv_lines:
            f.write(line + "\n")
    
    print(f"Done. Processed {success_count} files.")

if __name__ == "__main__":
    # 配置你的路径
    INPUT_ROOT = "/workspace/echoloc/datas/ntu_enzh_speech/"
    OUTPUT_WAV_ROOT = "/workspace/echoloc/datas/ntu_enzh_speech/merged_reconstruction/wavs"
    OUTPUT_TSV_PATH = "/workspace/echoloc/datas/ntu_enzh_speech/merged_reconstruction/test_merged.tsv"
    
    merge_audio_data(INPUT_ROOT, OUTPUT_WAV_ROOT, OUTPUT_TSV_PATH)