# Copyright 2026 Xiaohongshu. (Author: Kaituo Xu)

import glob
import logging
import os

logger = logging.getLogger(__name__)

from textgrid import IntervalTier, TextGrid


def get_wav_info(args):
    """
    Returns:
        wavs: list of (uttid, wav_path)
    """
    base = lambda p: os.path.basename(p).replace(".wav", "")
    if args.wav_path:
        wavs = [(base(args.wav_path), args.wav_path)]
    elif args.wav_paths and len(args.wav_paths) >= 1:
        wavs = [(base(p), p) for p in sorted(args.wav_paths)]
    elif args.wav_scp:
        wavs = [line.strip().split() for line in open(args.wav_scp)]
        if args.sort_wav_by_dur:
            logger.info("Sort wav by duration...")
            utt2dur = os.path.join(os.path.dirname(args.wav_scp), "utt2dur")
            if os.path.exists(utt2dur):
                utt2dur = [l.strip().split() for l in open(utt2dur)]
                utt2dur = {l[0]: float(l[1]) for l in utt2dur if len(l) == 2}
                wavs = sorted(wavs, key=lambda x: -utt2dur[x[0]])
                logger.info("Sort Done")
            else:
                logger.info(f"Not find {utt2dur}, un-sort")
    elif args.wav_dir:
        wavs = glob.glob(f"{args.wav_dir}/**/*.wav", recursive=True)
        wavs = [(base(p), p) for p in sorted(wavs)]
    else:
        raise ValueError("Please provide valid wav info")
    logger.info(f"#wavs={len(wavs)}")
    return wavs


def write_textgrid(wav_path, wav_dur, event):
    textgrid_file = wav_path.replace(".wav", ".TextGrid")
    logger.info(f"Write {textgrid_file}")
    textgrid = TextGrid(maxTime=wav_dur)
    tier = IntervalTier(name="token", maxTime=wav_dur)
    for token, start_s, end_s in event:
        if start_s == end_s:
            logger.info(f"Write TG, skip start=end {start_s}")
            continue
        start_s = max(start_s, 0)
        end_s = min(end_s, wav_dur)
        tier.add(minTime=start_s, maxTime=end_s, mark=token)
    textgrid.append(tier)
    textgrid.write(textgrid_file)
