import argparse
import os
import platform
import re
import shutil
import subprocess
import tempfile
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 resolve_path(value: str) -> Path:
    return Path(value).expanduser().resolve()


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 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 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 ps_quote(value: str) -> str:
    return "'" + value.replace("'", "''") + "'"


def speak_macos(text: str, voice: str, rate: int, output: Path) -> None:
    run(["say", "-v", voice, "-r", str(rate), "-o", str(output), text])


def speak_windows(text: str, voice: str, rate: int, volume: int, output: Path) -> None:
    ps_script = f"""
Add-Type -AssemblyName System.Speech
$synth = New-Object System.Speech.Synthesis.SpeechSynthesizer
$synth.SelectVoice({ps_quote(voice)})
$synth.Rate = {rate}
$synth.Volume = {volume}
$synth.SetOutputToWaveFile({ps_quote(str(output))})
$synth.Speak({ps_quote(text)})
$synth.Dispose()
"""
    run(["powershell", "-NoProfile", "-ExecutionPolicy", "Bypass", "-Command", ps_script])


def default_voice() -> str:
    if platform.system() == "Darwin":
        return "Daniel"
    if platform.system() == "Windows":
        return "Microsoft David Desktop"
    raise RuntimeError("One-step SRT voiceover generation is supported only on macOS and Windows.")


def default_rate() -> int:
    if platform.system() == "Darwin":
        return 185
    if platform.system() == "Windows":
        return 0
    raise RuntimeError("One-step SRT voiceover generation is supported only on macOS and Windows.")


def speak(text: str, voice: str, rate: int, volume: float, output: Path) -> None:
    system = platform.system()
    if system == "Darwin":
        speak_macos(text, voice, rate, output)
        return
    if system == "Windows":
        windows_volume = max(0, min(100, round(volume * 100)))
        speak_windows(text, voice, rate, windows_volume, output)
        return
    raise RuntimeError("One-step SRT voiceover generation is supported only on macOS and Windows.")


def create_voiceover_from_srt(
    srt_path: Path,
    wav_output: Path,
    voice: str,
    rate: int,
    work_dir: Path,
    volume: float,
    silence_padding: float,
    sample_rate: int,
    channels: int,
) -> None:
    captions = read_srt(srt_path)
    segment_files = []

    for i, (start, end, text) in enumerate(captions, 1):
        duration = end - start
        raw_audio = work_dir / f"raw_{i:03d}.aiff" if platform.system() == "Darwin" else work_dir / f"raw_{i:03d}.wav"
        segment_wav = work_dir / f"seg_{i:03d}.wav"

        speak(text, voice, rate, volume, raw_audio)

        spoken_duration = media_duration(raw_audio)
        padding = min(silence_padding, max(duration - spoken_duration, 0))
        effective_duration = duration if 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={volume},apad,atrim=0:{effective_duration}"

        run([
            "ffmpeg", "-y",
            "-i", str(raw_audio),
            "-af", audio_filter,
            "-ar", str(sample_rate),
            "-ac", str(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)])


def combine_video_audio(video_file: Path, audio_file: Path, output_file: Path, reencode_video: bool, volume: float) -> None:
    video_codec_args = ["-c:v", "libx264", "-preset", "medium", "-crf", "18"] if reencode_video else ["-c:v", "copy"]

    command = [
        "ffmpeg", "-y",
        "-i", str(video_file),
        "-i", str(audio_file),
        "-map", "0:v:0",
        "-map", "1:a:0",
        *video_codec_args,
        "-af", f"volume={volume}",
        "-c:a", "aac",
        "-b:a", "192k",
        "-movflags", "+faststart",
        str(output_file),
    ]
    run(command)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Replace a video's original audio with a voiceover audio file, or generate the voiceover from an SRT first."
    )
    parser.add_argument("--video", required=True, help="Input video file, for example input/tutorial.mp4")
    parser.add_argument("--audio", default=None, help="Existing voiceover audio file, for example output/voiceover.wav")
    parser.add_argument("--srt", default=None, help="Input SRT file for one-step voiceover generation")
    parser.add_argument("--output", required=True, help="Output MP4 file, for example output/final_video.mp4")
    parser.add_argument("--voice", default=None, help="Voice name for one-step SRT mode")
    parser.add_argument("--rate", type=int, default=None, help="Speech rate for one-step SRT mode")
    parser.add_argument("--audio-output", default=None, help="Optional copy of the generated voiceover WAV in one-step SRT mode")
    parser.add_argument("--keep-temp", action="store_true", help="Keep temporary voice parts for debugging in one-step SRT mode")
    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 in one-step SRT mode, if the SRT slot allows it. Default: 0.0")
    parser.add_argument("--sample-rate", type=int, default=44100, help="Generated voiceover sample rate in one-step SRT mode. Default: 44100")
    parser.add_argument("--channels", type=int, default=2, choices=[1, 2], help="Generated voiceover channels in one-step SRT mode. Default: 2")
    parser.add_argument(
        "--reencode-video",
        action="store_true",
        help="Re-encode video with H.264 instead of copying. Use only if the copied output has playback issues.",
    )
    args = parser.parse_args()

    if not args.audio and not args.srt:
        parser.error("Provide either --audio for the current two-step flow or --srt for one-step generation.")
    if args.audio and args.srt:
        parser.error("Use either --audio or --srt, not both.")
    if args.volume <= 0:
        parser.error("--volume must be greater than 0.")
    if args.silence_padding < 0:
        parser.error("--silence-padding must be 0 or greater.")

    video_file = resolve_path(args.video)
    output_file = resolve_path(args.output)

    if not video_file.exists():
        raise FileNotFoundError(f"Video not found: {video_file}")

    output_file.parent.mkdir(parents=True, exist_ok=True)

    if args.audio:
        audio_file = resolve_path(args.audio)
        if not audio_file.exists():
            raise FileNotFoundError(f"Audio not found: {audio_file}")
        combine_video_audio(video_file, audio_file, output_file, args.reencode_video, args.volume)
    else:
        srt_path = resolve_path(args.srt)
        if not srt_path.exists():
            raise FileNotFoundError(f"SRT not found: {srt_path}")

        voice = args.voice or default_voice()
        rate = args.rate if args.rate is not None else default_rate()
        audio_output = resolve_path(args.audio_output) if args.audio_output else None
        if audio_output:
            audio_output.parent.mkdir(parents=True, exist_ok=True)

        temp_dir = Path(tempfile.mkdtemp(prefix="video_generator_"))
        try:
            temp_wav = temp_dir / "voiceover.wav"
            create_voiceover_from_srt(
                srt_path,
                temp_wav,
                voice,
                rate,
                temp_dir,
                args.volume,
                args.silence_padding,
                args.sample_rate,
                args.channels,
            )
            if audio_output:
                run(["ffmpeg", "-y", "-i", str(temp_wav), "-c:a", "pcm_s16le", str(audio_output)])
                print(f"Done: {audio_output}")
            combine_video_audio(video_file, temp_wav, output_file, args.reencode_video, 1.0)
        finally:
            if not args.keep_temp:
                remove_temp_dir(temp_dir)
            else:
                print(f"Kept temporary voice parts: {temp_dir}")

    print(f"Done: {output_file}")


if __name__ == "__main__":
    main()