"""
Forced-alignment: 把 control2instruct.json 里 segments 的 txt 在合成音频里定位起止时间（秒），
回写到每个 segment item 的 'start' / 'end' 字段。

依赖 torchaudio 自带的 MMS_FA bundle（多语言强制对齐器，1100+ 语言），
配合 uroman 进行非拉丁字符到拉丁字符的转写。

用法:
    python align_segments.py --tsv  /path/to/test.tsv  \
                             --out-json /path/to/test_aligned.json \
                             [--device cuda]

产物 (test_aligned.json) 的 schema：
    {
      "items": [
        {
          "row_id": 0,
          "json_path": "...",
          "wav_path": "...",
          "duration": 5.64,
          "segments": [
            {"instruct": "...", "txt": "...", "start": 0.08, "end": 2.51,
             "norm_text": "okay mom i should ..."},
            ...
          ]
        },
        ...
      ]
    }
"""
import argparse
import csv
import json
import os
import sys
import time
from typing import Dict, List, Optional, Tuple

import numpy as np
import soundfile as sf
import torch
import torchaudio
import torchaudio.functional as Faudio

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from utils_common import load_wav_mono_16k


_DEFAULT_TORCH_HOME = "/workspace/echoloc/hf_cache/torch_home"
os.environ.setdefault("TORCH_HOME", _DEFAULT_TORCH_HOME)


# ----------------- normalization (uroman + filter) -----------------
class TextNormalizer:
    def __init__(self):
        try:
            import uroman as _uroman
        except ImportError:
            raise RuntimeError("uroman not installed. pip install uroman")
        self._ur = _uroman.Uroman()

    def __call__(self, text: str, allowed: set) -> str:
        s = self._ur.romanize_string(text)
        s = s.lower()
        out_chars = []
        for ch in s:
            if ch in allowed:
                out_chars.append(ch)
            elif ch.isspace():
                out_chars.append(" ")
            else:
                out_chars.append(" ")
        s = "".join(out_chars)
        while "  " in s:
            s = s.replace("  ", " ")
        return s.strip()


# ----------------- alignment pipeline -----------------
class MMSAligner:
    BLANK = 0
    STAR_LABEL = "*"

    def __init__(self, device: str = "cuda"):
        bundle = torchaudio.pipelines.MMS_FA
        self.bundle = bundle
        self.sr = bundle.sample_rate
        self.device = device
        self.model = bundle.get_model().to(device).eval()
        labels = list(bundle.get_labels())
        # labels[0] is '-' (blank), labels[-1] is '*' (skip)
        self.labels = labels
        self.dict = {c: i for i, c in enumerate(labels)}
        self.allowed = set(c for c in labels if c not in ("-", self.STAR_LABEL))
        self.norm = TextNormalizer()

    @torch.inference_mode()
    def emission(self, wav_np: np.ndarray) -> Tuple[torch.Tensor, int]:
        """Return log-prob emission (1, T, C) and audio length in samples."""
        wav = torch.from_numpy(wav_np).float().unsqueeze(0).to(self.device)
        em, _ = self.model(wav)
        return em, wav.shape[-1]

    def align(self, wav_np: np.ndarray, segment_texts: List[str]) -> List[Dict]:
        """
        Returns one dict per segment: {"start": float, "end": float, "norm_text": str}
        On failure returns start/end = None.
        """
        em, n_samples = self.emission(wav_np)

        # build words list with seg id mapping
        words: List[str] = []
        word_seg: List[int] = []
        normed_per_seg: List[str] = []
        for sid, txt in enumerate(segment_texts):
            n = self.norm(txt, self.allowed)
            normed_per_seg.append(n)
            for w in n.split():
                words.append(w)
                word_seg.append(sid)

        if not words:
            return [
                {"start": None, "end": None, "norm_text": normed_per_seg[i]}
                for i in range(len(segment_texts))
            ]

        # token ids per word
        tok_ids: List[int] = []
        word_spans: List[Tuple[int, int]] = []
        for w in words:
            s = len(tok_ids)
            for ch in w:
                if ch in self.dict and ch != "-" and ch != self.STAR_LABEL:
                    tok_ids.append(self.dict[ch])
            e = len(tok_ids)
            word_spans.append((s, e))

        if not tok_ids:
            return [
                {"start": None, "end": None, "norm_text": normed_per_seg[i]}
                for i in range(len(segment_texts))
            ]

        targets = torch.tensor([tok_ids], dtype=torch.int32, device=self.device)
        try:
            aligned, scores = Faudio.forced_align(em, targets, blank=self.BLANK)
        except RuntimeError as ex:
            # If forced alignment fails (e.g. emission too short), fall back to None
            print(f"[WARN] forced_align failed: {ex}", flush=True)
            return [
                {"start": None, "end": None, "norm_text": normed_per_seg[i]}
                for i in range(len(segment_texts))
            ]

        spans = Faudio.merge_tokens(aligned[0].cpu(), scores[0].cpu())
        nonblank = [s for s in spans if int(s.token) != self.BLANK]
        if len(nonblank) != len(tok_ids):
            # rare mismatch -> degrade to proportional split
            print(
                f"[WARN] nonblank({len(nonblank)}) != tok_ids({len(tok_ids)}); "
                f"falling back to proportional time split",
                flush=True,
            )
            return self._proportional_split(segment_texts, normed_per_seg, n_samples)

        ratio = n_samples / em.shape[1] / self.sr
        seg_times: Dict[int, List[float]] = {}
        for wi, (s, e) in enumerate(word_spans):
            if e <= s:
                continue
            sid = word_seg[wi]
            sp = nonblank[s:e]
            t0 = float(sp[0].start) * ratio
            t1 = float(sp[-1].end) * ratio
            if sid not in seg_times:
                seg_times[sid] = [t0, t1]
            else:
                seg_times[sid][0] = min(seg_times[sid][0], t0)
                seg_times[sid][1] = max(seg_times[sid][1], t1)

        results = []
        dur_total = n_samples / self.sr
        for i in range(len(segment_texts)):
            if i in seg_times:
                t0, t1 = seg_times[i]
                t0 = max(0.0, t0 - 0.02)
                t1 = min(dur_total, t1 + 0.02)
                results.append(
                    {"start": round(t0, 3), "end": round(t1, 3), "norm_text": normed_per_seg[i]}
                )
            else:
                results.append({"start": None, "end": None, "norm_text": normed_per_seg[i]})
        return results

    @staticmethod
    def _proportional_split(
        segment_texts: List[str], normed_per_seg: List[str], n_samples: int, sr: int = 16000
    ) -> List[Dict]:
        # Fallback split based on text length ratio
        lens = [max(1, len(t.replace(" ", ""))) for t in normed_per_seg]
        total = sum(lens)
        dur = n_samples / sr
        out = []
        acc = 0.0
        for i, L in enumerate(lens):
            d = dur * L / total
            t0 = acc
            t1 = acc + d
            acc = t1
            out.append({"start": round(t0, 3), "end": round(t1, 3), "norm_text": normed_per_seg[i]})
        return out


def read_tsv(path: str) -> List[Tuple[int, str, str]]:
    rows = []
    with open(path, "r", encoding="utf-8") as f:
        rdr = csv.reader(f, delimiter="\t")
        for line in rdr:
            if len(line) < 3:
                continue
            rid, jp, wp = line[0], line[1], line[2]
            try:
                rid_i = int(rid)
            except ValueError:
                continue
            rows.append((rid_i, jp, wp))
    return rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tsv", required=True, help="input tsv with columns row_id, json, wav")
    ap.add_argument("--out-json", required=True, help="output aligned json")
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--seg-key", default="segments", help="json key holding segments")
    ap.add_argument("--limit", type=int, default=-1, help="debug: only process first N rows")
    args = ap.parse_args()

    rows = read_tsv(args.tsv)
    if args.limit > 0:
        rows = rows[: args.limit]
    print(f"[align] {len(rows)} rows from {args.tsv}", flush=True)

    aligner = MMSAligner(device=args.device)

    items = []
    t0 = time.time()
    for k, (rid, jp, wp) in enumerate(rows):
        if not os.path.exists(jp) or not os.path.exists(wp):
            print(f"[skip row={rid}] missing file: {jp} | {wp}", flush=True)
            continue
        try:
            with open(jp, "r", encoding="utf-8") as f:
                meta = json.load(f)
        except Exception as ex:
            print(f"[skip row={rid}] bad json: {ex}", flush=True)
            continue

        segs = meta.get(args.seg_key, [])
        if not segs:
            continue
        seg_texts = [s.get("txt", "") for s in segs]

        try:
            wav = load_wav_mono_16k(wp, target_sr=aligner.sr)
        except Exception as ex:
            print(f"[skip row={rid}] load_wav: {ex}", flush=True)
            continue

        try:
            timed = aligner.align(wav, seg_texts)
        except Exception as ex:
            print(f"[fallback row={rid}] align fail: {ex}", flush=True)
            timed = MMSAligner._proportional_split(seg_texts, seg_texts, len(wav), aligner.sr)

        new_segs = []
        for orig, t in zip(segs, timed):
            d = dict(orig)
            d["start"] = t["start"]
            d["end"] = t["end"]
            d["norm_text"] = t["norm_text"]
            new_segs.append(d)

        items.append(
            {
                "row_id": rid,
                "json_path": jp,
                "wav_path": wp,
                "duration": round(len(wav) / aligner.sr, 3),
                "segments": new_segs,
                "combined": meta.get("combined"),
                "combined_no_speaker": meta.get("combined_no_speaker"),
            }
        )

        if (k + 1) % 50 == 0:
            elapsed = time.time() - t0
            print(
                f"[align] {k + 1}/{len(rows)}  ~{elapsed:.1f}s  "
                f"({(k + 1) / max(elapsed, 1e-6):.2f} it/s)",
                flush=True,
            )

    os.makedirs(os.path.dirname(args.out_json), exist_ok=True)
    with open(args.out_json, "w", encoding="utf-8") as f:
        json.dump({"items": items}, f, ensure_ascii=False, indent=2)
    print(f"[align] done. saved {len(items)} items -> {args.out_json}", flush=True)


if __name__ == "__main__":
    main()