"""Build trimmed_transcripts.xlsx — one row per audio file.

Trimmed-text logic: include every scenario turn EXCEPT those listed as
dropped in concat_report.json. Turns whose audio was never rendered
(neither kept nor dropped in the report) are included with a [NO AUDIO]
marker so the reviewer can decide whether to keep or strip them.
"""
from __future__ import annotations

import json
from pathlib import Path

from openpyxl import Workbook
from openpyxl.styles import Alignment, Font, PatternFill

REPORT = Path(
    "/dataset/kemix-engine/package/face/animasync-face-v3/data/audio_preview_monologue_trim/concat_report.json"
)
EMO_DIR = Path(
    "/dataset/kemix-engine/package/face/animasync-face-v3/data/emotion"
)
OUT = Path(
    "/dataset/kemix-engine/package/face/animasync-face-v3/data/trimmed_transcripts.xlsx"
)

SEED_PRIORITY = [
    "seed_scenarios_final.jsonl",
    "seed_train_final.jsonl",
    "seed_train.jsonl",
    "seed_val.jsonl",
    "seed_test.jsonl",
    "seed_generated.jsonl",
    "seed_augmented.jsonl",
]

SET_LIST = [
    "long_001", "long_002", "long_003", "long_004", "long_005",
    "long_006", "long_007", "long_008", "long_009", "long_010",
    "long_011", "long_055_p1", "long_048", "long_042", "long_102_p1",
    "long_055", "long_gen_028", "long_012", "long_055_p0", "long_081",
    "long_049", "long_gen_026", "long_063", "long_032", "long_108_p1",
    "long_102_p0", "long_gen_050", "long_078", "long_013", "long_014",
    "long_001_p0", "long_002_p0", "long_018_p1", "long_116_p0", "long_111",
    "long_074", "long_003_p0", "long_108", "long_116_p1", "long_002_p1",
    "long_gen_025", "long_106_p1", "long_060", "long_065_p1", "long_050_p1",
    "long_087", "long_041_p1", "long_107_p1", "long_065", "long_050_p0",
    "long_066", "long_gen_042", "long_gen_035_p0", "long_024_p1", "long_084",
    "long_003_p1", "long_041", "long_041_p0", "long_gen_056", "long_001_p1",
]


def load_scenarios() -> dict[str, dict]:
    out: dict[str, dict] = {}
    for fn in SEED_PRIORITY:
        p = EMO_DIR / fn
        if not p.exists():
            continue
        with p.open() as f:
            for line in f:
                try:
                    obj = json.loads(line)
                except json.JSONDecodeError:
                    continue
                sid = obj.get("scenario_id")
                if sid and sid not in out:
                    out[sid] = obj
    return out


def main() -> None:
    report = json.loads(REPORT.read_text())
    scenarios = load_scenarios()

    wb = Workbook()
    ws = wb.active
    ws.title = "trimmed_transcripts"

    headers = [
        "set_no",
        "audio_file",
        "duration_s",
        "trimmed_text (per-turn lines)",
        "dropped_turns",
    ]
    ws.append(headers)

    head_font = Font(bold=True)
    head_fill = PatternFill("solid", fgColor="DDDDDD")
    for col_idx in range(1, len(headers) + 1):
        cell = ws.cell(row=1, column=col_idx)
        cell.font = head_font
        cell.fill = head_fill
        cell.alignment = Alignment(vertical="center")

    wrap_top = Alignment(wrap_text=True, vertical="top")

    for set_no, sid in enumerate(SET_LIST, start=1):
        filename = f"{sid}.mp3"
        entry = report.get(sid, {})
        scenario = scenarios.get(sid, {})
        turns = scenario.get("turns", [])

        kept_by_idx = {k["turn_idx"]: k.get("duration", "") for k in entry.get("kept", [])}
        dropped_indices = {d["turn_idx"] for d in entry.get("dropped", [])}

        kept_lines = []
        for ti, t in enumerate(turns):
            if ti in dropped_indices:
                continue
            emo = t.get("emotion", "")
            txt = t.get("text", "")
            if ti in kept_by_idx:
                dur = kept_by_idx[ti]
                kept_lines.append(f"[t{ti} · {emo} · {dur}s] {txt}")
            else:
                # Turn exists in script but has no audio in the trimmed mp3
                kept_lines.append(f"[t{ti} · {emo} · NO AUDIO] {txt}")
        trimmed_text = "\n".join(kept_lines)

        dropped_lines = []
        for d in entry.get("dropped", []):
            di = d["turn_idx"]
            if di < len(turns):
                t = turns[di]
                emo = t.get("emotion", "")
                txt = t.get("text", "")
            else:
                emo = ""
                txt = "<out of range>"
            ddur = d.get("duration", "")
            dropped_lines.append(f"[t{di} · {emo} · {ddur}s] {txt}")
        dropped_text = "\n".join(dropped_lines) if dropped_lines else "(none)"

        row = [
            set_no,
            filename,
            entry.get("concat_duration", ""),
            trimmed_text,
            dropped_text,
        ]
        ws.append(row)

        r = ws.max_row
        for col_idx in range(1, len(headers) + 1):
            ws.cell(row=r, column=col_idx).alignment = wrap_top

    ws.column_dimensions["A"].width = 7
    ws.column_dimensions["B"].width = 22
    ws.column_dimensions["C"].width = 11
    ws.column_dimensions["D"].width = 90
    ws.column_dimensions["E"].width = 60
    ws.freeze_panes = "A2"

    OUT.parent.mkdir(parents=True, exist_ok=True)
    wb.save(OUT)
    print(f"Wrote {OUT}")


if __name__ == "__main__":
    main()