import argparse
import os
import re
import shutil
import subprocess
from pathlib import Path

SRT_BLOCK_RE = re.compile(
    r"\s*(\d+)\s+([\d:,]+)\s+-->\s+([\d:,]+)\s+(.+?)(?=\n\s*\d+\s+\d\d:|\Z)",
    re.S,
)


def to_seconds(timestamp: str) -> float:
    h, m, rest = timestamp.split(":")
    s, ms = rest.split(",")
    return int(h) * 3600 + int(m) * 60 + int(s) + int(ms) / 1000


def read_srt(path: Path) -> list[tuple[float, float, str]]:
    text = path.read_text(encoding="utf-8-sig")
    captions = []
    for _, start, end, body in SRT_BLOCK_RE.findall(text):
        clean_text = " ".join(line.strip() for line in body.strip().splitlines())
        captions.append((to_seconds(start), to_seconds(end), clean_text))
    if not captions:
        raise ValueError(f"No captions found in {path}")
    return captions


def resolve_path(value: str) -> Path:
    return Path(value).expanduser().resolve()


def media_duration(path: Path) -> float:
    result = subprocess.run(
        [
            "ffprobe",
            "-v", "error",
            "-show_entries", "format=duration",
            "-of", "default=noprint_wrappers=1:nokey=1",
            str(path),
        ],
        check=True,
        capture_output=True,
        text=True,
    )
    return float(result.stdout.strip())


def run(command: list[str]) -> None:
    print("Running:", " ".join(command))
    subprocess.run(command, check=True)


def remove_temp_dir(path: Path) -> None:
    try:
        shutil.rmtree(path, ignore_errors=True)
    except PermissionError:
        pass
    except OSError:
        pass


def main() -> None:
    parser = argparse.ArgumentParser(description="Create a timed voiceover WAV/MP3 from an SRT file on macOS.")
    parser.add_argument("--srt", required=True, help="Input SRT file")
    parser.add_argument("--voice", default="Daniel", help="macOS voice name, e.g. Daniel or Reed (English (US))")
    parser.add_argument("--rate", type=int, default=185, help="Speech rate. Higher is faster. Example: 180-200")
    parser.add_argument("--wav", required=True, help="Output WAV file")
    parser.add_argument("--mp3", default=None, help="Optional output MP3 file")
    parser.add_argument("--audio-output", default=None, help="Optional extra copy of the generated WAV file")
    parser.add_argument("--keep-temp", action="store_true", help="Keep temporary voice parts for debugging")
    parser.add_argument("--volume", type=float, default=1.0, help="Voiceover volume multiplier. Default: 1.0")
    parser.add_argument("--silence-padding", type=float, default=0.0, help="Seconds of silence after each caption, if the SRT slot allows it. Default: 0.0")
    parser.add_argument("--sample-rate", type=int, default=44100, help="Output sample rate. Default: 44100")
    parser.add_argument("--channels", type=int, default=2, choices=[1, 2], help="Output audio channels. Default: 2")
    args = parser.parse_args()

    srt_path = resolve_path(args.srt)
    wav_output = resolve_path(args.wav)
    mp3_output = resolve_path(args.mp3) if args.mp3 else None
    audio_output = resolve_path(args.audio_output) if args.audio_output else None
    work_dir = Path(__file__).resolve().parent / "voice_parts"

    if not srt_path.exists():
        raise FileNotFoundError(f"SRT not found: {srt_path}")
    if args.volume <= 0:
        raise ValueError("--volume must be greater than 0")
    if args.silence_padding < 0:
        raise ValueError("--silence-padding must be 0 or greater")

    work_dir.mkdir(parents=True, exist_ok=True)
    wav_output.parent.mkdir(parents=True, exist_ok=True)
    if mp3_output:
        mp3_output.parent.mkdir(parents=True, exist_ok=True)
    if audio_output:
        audio_output.parent.mkdir(parents=True, exist_ok=True)

    captions = read_srt(srt_path)
    segment_files = []

    try:
        for i, (start, end, text) in enumerate(captions, 1):
            duration = end - start
            raw_aiff = work_dir / f"raw_{i:03d}.aiff"
            segment_wav = work_dir / f"seg_{i:03d}.wav"

            run(["say", "-v", args.voice, "-r", str(args.rate), "-o", str(raw_aiff), text])

            spoken_duration = media_duration(raw_aiff)
            padding = min(args.silence_padding, max(duration - spoken_duration, 0))
            effective_duration = duration if args.silence_padding == 0 else min(duration, spoken_duration + padding)

            if spoken_duration > duration:
                print(f"Warning: caption {i} speech is {spoken_duration:.2f}s but slot is {duration:.2f}s; it will be trimmed.")

            audio_filter = f"volume={args.volume},apad,atrim=0:{effective_duration}"

            run([
                "ffmpeg", "-y",
                "-i", str(raw_aiff),
                "-af", audio_filter,
                "-ar", str(args.sample_rate),
                "-ac", str(args.channels),
                str(segment_wav),
            ])
            segment_files.append(segment_wav)

        concat_list = work_dir / "list.txt"
        concat_list.write_text("".join(f"file '{os.path.abspath(file)}'\n" for file in segment_files), encoding="utf-8")

        run(["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list), "-c:a", "pcm_s16le", str(wav_output)])
        print(f"Done: {wav_output}")

        if audio_output:
            run(["ffmpeg", "-y", "-i", str(wav_output), "-c:a", "pcm_s16le", str(audio_output)])
            print(f"Done: {audio_output}")

        if mp3_output:
            run(["ffmpeg", "-y", "-i", str(wav_output), "-codec:a", "libmp3lame", "-b:a", "192k", str(mp3_output)])
            print(f"Done: {mp3_output}")
    finally:
        if not args.keep_temp:
            remove_temp_dir(work_dir)


if __name__ == "__main__":
    main()