"""Generate voiceover audio + word-timed ASS subtitles from story text."""
import asyncio
import re
from pathlib import Path

import edge_tts

from config import (
    AUDIO_DIR,
    SUB_COLOR_HIGHLIGHT,
    SUB_COLOR_NORMAL,
    SUB_FONT,
    SUB_FONT_SIZE,
    SUB_MARGIN_V,
    SUB_OUTLINE_COLOR,
    SUB_OUTLINE_SIZE,
    SUB_WORDS_PER_GROUP,
    TTS_VOICE,
)

# edge-tts reports offsets in 100-nanosecond units
_TICKS_PER_SEC = 10_000_000


def _ticks_to_seconds(ticks: int) -> float:
    return ticks / _TICKS_PER_SEC


def _seconds_to_ass(t: float) -> str:
    """Convert float seconds → ASS timestamp  H:MM:SS.cc"""
    t = max(0.0, t)
    h = int(t // 3600)
    m = int((t % 3600) // 60)
    s = t % 60
    return f"{h}:{m:02d}:{s:05.2f}"


async def _synthesise(text: str, voice: str) -> tuple[bytes, list[dict]]:
    """Stream edge-tts and collect audio bytes + per-word timing events."""
    communicate = edge_tts.Communicate(text, voice, rate="+8%", boundary="WordBoundary")
    audio_chunks: list[bytes] = []
    words: list[dict] = []

    async for chunk in communicate.stream():
        if chunk["type"] == "audio":
            audio_chunks.append(chunk["data"])
        elif chunk["type"] == "WordBoundary":
            words.append(
                {
                    "word": chunk["text"],
                    "start": _ticks_to_seconds(chunk["offset"]),
                    "end": _ticks_to_seconds(chunk["offset"] + chunk["duration"]),
                }
            )

    return b"".join(audio_chunks), words


def _build_ass(words: list[dict], total_duration: float) -> str:
    """
    Group words into fixed-size groups and emit one ASS event per group.
    Within each event the current word is highlighted in yellow; the rest
    are white.  Each word gets its own event overlapping the group so the
    highlight moves word-by-word.
    """
    header = f"""\
[Script Info]
ScriptType: v4.00+
PlayResX: 1080
PlayResY: 1920
WrapStyle: 1
ScaledBorderAndShadow: yes

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,{SUB_FONT},{SUB_FONT_SIZE},{SUB_COLOR_NORMAL},&H000000FF,{SUB_OUTLINE_COLOR},&H64000000,-1,0,0,0,100,100,0,0,1,{SUB_OUTLINE_SIZE},0,2,40,40,{SUB_MARGIN_V},1

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
"""
    # Split words into groups of SUB_WORDS_PER_GROUP
    groups: list[list[dict]] = []
    for i in range(0, len(words), SUB_WORDS_PER_GROUP):
        groups.append(words[i : i + SUB_WORDS_PER_GROUP])

    events: list[str] = []
    for g_idx, group in enumerate(groups):
        group_start = group[0]["start"]
        # group stays visible until the next group starts (or audio ends)
        if g_idx + 1 < len(groups):
            group_end = groups[g_idx + 1][0]["start"]
        else:
            group_end = total_duration + 0.3

        for w_idx, word in enumerate(group):
            # This event spans the word's own timing
            w_start = word["start"] if w_idx else group_start
            w_end = group[w_idx + 1]["start"] if w_idx + 1 < len(group) else group_end
            # Pop-in: the group scales from 80% to 100% when it first appears
            pop = "{\\fscx80\\fscy80\\t(0,90,\\fscx100\\fscy100)}" if w_idx == 0 else ""

            # Build the line with the highlighted word coloured
            parts: list[str] = []
            for j, w in enumerate(group):
                clean = re.sub(r"[^\w\'\-\,\.\!\?]", "", w["word"]).upper()
                if j == w_idx:
                    # Uppercase only the word; ASS override tags are case-sensitive
                    parts.append(
                        f"{{\\c{SUB_COLOR_HIGHLIGHT}&\\fscx112\\fscy112}}{clean}"
                        f"{{\\c{SUB_COLOR_NORMAL}&\\fscx100\\fscy100}}"
                    )
                else:
                    parts.append(clean)

            line_text = pop + " ".join(parts)
            events.append(
                f"Dialogue: 0,{_seconds_to_ass(w_start)},{_seconds_to_ass(w_end)},"
                f"Default,,0,0,0,,{line_text}"
            )

    return header + "\n".join(events) + "\n"


def generate_tts(text: str, stem: str) -> tuple[Path, Path, float]:
    """
    Synthesise `text` and write:
      AUDIO_DIR/<stem>.mp3
      AUDIO_DIR/<stem>.ass

    Returns (mp3_path, ass_path, audio_duration_seconds).
    """
    mp3_path = AUDIO_DIR / f"{stem}.mp3"
    ass_path = AUDIO_DIR / f"{stem}.ass"

    audio_bytes, words = asyncio.run(_synthesise(text, TTS_VOICE))

    mp3_path.write_bytes(audio_bytes)

    if not words:
        raise RuntimeError("edge-tts returned no word-boundary events")

    total_duration = words[-1]["end"] + 0.5
    ass_path.write_text(_build_ass(words, total_duration), encoding="utf-8")

    return mp3_path, ass_path, total_duration


if __name__ == "__main__":
    import sys

    sample = (
        "I never expected my coworker to steal my lunch for three months. "
        "So I put a tiny tracker inside a sealed container of leftovers. "
        "At noon the signal led straight to my manager's office. "
        "He had been taking food because he thought nobody would say a word. "
        "I showed HR the location history and suddenly he had a lot to explain. "
        "The next day he brought in sandwiches for everyone. Nobody touched them."
    )
    mp3, ass, dur = generate_tts(sample, "test_tts")
    print(f"Audio: {mp3}  ({dur:.1f}s)")
    print(f"ASS:   {ass}")
    print(ass.read_text()[:600])