#!/usr/bin/env python3
"""Bundle video sketches into ~40-minute episodes with short/long alternation."""

from __future__ import annotations

import argparse
import shutil
import statistics
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

VIDEO_EXTENSIONS = {".mp4", ".mkv", ".mov", ".avi", ".webm", ".m4v"}
TARGET_WIDTH = 1920
TARGET_HEIGHT = 1080


@dataclass
class Sketch:
    path: Path
    duration: float  # seconds
    label: str = ""

    @property
    def name(self) -> str:
        return self.label or self.path.name


def format_duration(seconds: float) -> str:
    total = int(round(seconds))
    minutes, secs = divmod(total, 60)
    hours, minutes = divmod(minutes, 60)
    if hours:
        return f"{hours}:{minutes:02d}:{secs:02d}"
    return f"{minutes}:{secs:02d}"


def require_binaries() -> None:
    missing = [name for name in ("ffprobe", "ffmpeg") if shutil.which(name) is None]
    if missing:
        print(
            f"Error: missing required binary/binaries: {', '.join(missing)}. "
            "Install ffmpeg and ensure it is on PATH.",
            file=sys.stderr,
        )
        sys.exit(1)


def scan_videos(folder: Path) -> list[Path]:
    files = sorted(
        p
        for p in folder.rglob("*")
        if p.is_file() and p.suffix.lower() in VIDEO_EXTENSIONS
    )
    if not files:
        print(f"Error: no video files found in {folder}", file=sys.stderr)
        sys.exit(1)
    return files


def probe_duration(path: Path) -> float:
    result = subprocess.run(
        [
            "ffprobe",
            "-v",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "csv=p=0",
            str(path),
        ],
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        raise RuntimeError(
            f"ffprobe failed for {path.name}: {result.stderr.strip() or 'unknown error'}"
        )
    value = result.stdout.strip()
    if not value:
        raise RuntimeError(f"ffprobe returned empty duration for {path.name}")
    return float(value)


def load_sketches(folder: Path) -> list[Sketch]:
    sketches: list[Sketch] = []
    for path in scan_videos(folder):
        try:
            duration = probe_duration(path)
        except (RuntimeError, ValueError) as exc:
            print(f"Warning: skipping {path}: {exc}", file=sys.stderr)
            continue
        try:
            label = str(path.relative_to(folder))
        except ValueError:
            label = path.name
        sketches.append(Sketch(path=path, duration=duration, label=label))
    if not sketches:
        print("Error: no sketches with readable duration", file=sys.stderr)
        sys.exit(1)
    return sketches


def split_short_long(sketches: list[Sketch]) -> tuple[list[Sketch], list[Sketch]]:
    median = statistics.median(s.duration for s in sketches)
    short = [s for s in sketches if s.duration <= median]
    long = [s for s in sketches if s.duration > median]
    # If everything landed in one bucket (e.g. identical lengths), keep packing workable.
    if not short or not long:
        by_duration = sorted(sketches, key=lambda s: s.duration)
        mid = len(by_duration) // 2
        short, long = by_duration[:mid], by_duration[mid:]
        if not short and long:
            short.append(long.pop(0))
        elif not long and short:
            long.append(short.pop())
    return short, long


def alternating_order(sketches: list[Sketch]) -> list[Sketch]:
    """Alternate short/long within a set of sketches (median split)."""
    if len(sketches) <= 1:
        return list(sketches)
    short, long = split_short_long(sketches)
    short_q = sorted(short, key=lambda s: s.duration)
    long_q = sorted(long, key=lambda s: s.duration, reverse=True)

    ordered: list[Sketch] = []
    prefer_short = True
    while short_q or long_q:
        primary, secondary = (short_q, long_q) if prefer_short else (long_q, short_q)
        if primary:
            ordered.append(primary.pop(0))
        elif secondary:
            ordered.append(secondary.pop(0))
        prefer_short = not prefer_short
    return ordered


def pack_into_n_episodes(
    sketches: list[Sketch], episode_count: int
) -> tuple[list[list[Sketch]], float]:
    """Pack into exactly episode_count balanced episodes.

    Uses longest-processing-time-first bin packing for even lengths, then
    reorders each episode for short/long alternation.
    """
    if episode_count < 1:
        raise ValueError("episode_count must be at least 1")
    if episode_count > len(sketches):
        raise ValueError(
            f"cannot make {episode_count} episodes from only {len(sketches)} sketches"
        )

    total = sum(s.duration for s in sketches)
    target_seconds = total / episode_count

    # LPT multiprocessor scheduling: place longest sketches first into the
    # currently shortest episode — minimizes length spread.
    by_length = sorted(sketches, key=lambda s: (-s.duration, s.name))
    episodes: list[list[Sketch]] = [[] for _ in range(episode_count)]
    durations = [0.0] * episode_count
    for sketch in by_length:
        index = min(range(episode_count), key=lambda i: (durations[i], i))
        episodes[index].append(sketch)
        durations[index] += sketch.duration

    episodes = [alternating_order(ep) for ep in episodes]
    return episodes, target_seconds


def pack_episodes(
    sketches: list[Sketch], target_seconds: float
) -> list[list[Sketch]]:
    """Pack into episodes of roughly target_seconds, balanced in length."""
    total = sum(s.duration for s in sketches)
    episode_count = max(1, round(total / target_seconds))
    episode_count = min(episode_count, len(sketches))
    episodes, _ = pack_into_n_episodes(sketches, episode_count)
    return episodes


def print_plan(episodes: list[list[Sketch]], target_seconds: float) -> None:
    print(f"Target episode length: {format_duration(target_seconds)}")
    print(f"Episodes: {len(episodes)}\n")
    for index, episode in enumerate(episodes, start=1):
        total = sum(s.duration for s in episode)
        print(f"=== Aflevering {index:02d} ({format_duration(total)}) ===")
        for sketch in episode:
            print(f"  {format_duration(sketch.duration):>8}  {sketch.name}")
        print()


def nvenc_available(codec: str) -> bool:
    encoder = "hevc_nvenc" if codec == "h265" else "h264_nvenc"
    result = subprocess.run(
        ["ffmpeg", "-hide_banner", "-encoders"],
        capture_output=True,
        text=True,
        check=False,
    )
    return result.returncode == 0 and encoder in result.stdout


def encoder_name(use_gpu: bool, codec: str) -> str:
    if codec == "h265":
        return "hevc_nvenc" if use_gpu else "libx265"
    return "h264_nvenc" if use_gpu else "libx264"


def video_encoder_args(use_gpu: bool, codec: str) -> list[str]:
    if codec == "h265":
        if use_gpu:
            return [
                "-c:v",
                "hevc_nvenc",
                "-preset",
                "p4",
                "-rc",
                "vbr",
                "-cq",
                "28",
                "-b:v",
                "0",
                "-tag:v",
                "hvc1",
            ]
        return [
            "-c:v",
            "libx265",
            "-preset",
            "medium",
            "-crf",
            "28",
            "-tag:v",
            "hvc1",
        ]

    if use_gpu:
        return [
            "-c:v",
            "h264_nvenc",
            "-preset",
            "p4",
            "-rc",
            "vbr",
            "-cq",
            "23",
            "-b:v",
            "0",
            "-profile:v",
            "high",
        ]
    return [
        "-c:v",
        "libx264",
        "-preset",
        "medium",
        "-crf",
        "23",
    ]


def merge_episode(
    episode: list[Sketch],
    output_path: Path,
    *,
    use_gpu: bool = False,
    codec: str = "h264",
) -> None:
    """Re-encode and concat sketches into one episode file."""
    n = len(episode)
    cmd: list[str] = ["ffmpeg", "-y"]
    for sketch in episode:
        cmd.extend(["-i", str(sketch.path)])

    # Scale/pad to a common canvas, then concat. Force consistent fps/sample rate.
    filter_parts: list[str] = []
    for i in range(n):
        filter_parts.append(
            f"[{i}:v]scale={TARGET_WIDTH}:{TARGET_HEIGHT}:"
            f"force_original_aspect_ratio=decrease,"
            f"pad={TARGET_WIDTH}:{TARGET_HEIGHT}:(ow-iw)/2:(oh-ih)/2,"
            f"setsar=1,fps=25,format=yuv420p[v{i}];"
            f"[{i}:a]aformat=sample_rates=48000:channel_layouts=stereo,"
            f"aresample=async=1[a{i}]"
        )
    concat_inputs = "".join(f"[v{i}][a{i}]" for i in range(n))
    filter_parts.append(f"{concat_inputs}concat=n={n}:v=1:a=1[outv][outa]")
    filter_complex = ";".join(filter_parts)

    encoder = encoder_name(use_gpu, codec)
    cmd.extend(
        [
            "-filter_complex",
            filter_complex,
            "-map",
            "[outv]",
            "-map",
            "[outa]",
            *video_encoder_args(use_gpu, codec),
            "-c:a",
            "aac",
            "-b:a",
            "192k",
            "-movflags",
            "+faststart",
            str(output_path),
        ]
    )

    print(f"Merging {output_path.name} ({n} sketches, {encoder})...")
    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    if result.returncode != 0:
        raise RuntimeError(
            f"ffmpeg failed for {output_path.name}:\n{result.stderr[-2000:]}"
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Bundle video sketches from a folder into episodes, "
            "alternating short and long sketches. "
            "Specify either a target length or a target episode count."
        )
    )
    parser.add_argument(
        "input",
        type=Path,
        help="Folder containing sketch video files (scanned recursively)",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=Path("episodes"),
        help="Output folder for merged episodes (default: ./episodes)",
    )
    target = parser.add_mutually_exclusive_group()
    target.add_argument(
        "-t",
        "--target-minutes",
        type=float,
        help="Target episode length in minutes (default: 40 if neither target is set)",
    )
    target.add_argument(
        "-n",
        "--target-episodes",
        type=int,
        help="Exact number of episodes to create (length derived from total duration)",
    )
    parser.add_argument(
        "--codec",
        choices=("h264", "h265"),
        default="h264",
        help="Video codec (default: h264). h265 is smaller; use with --gpu for speed",
    )
    parser.add_argument(
        "--gpu",
        action="store_true",
        help="Encode with NVIDIA NVENC; falls back to CPU on failure",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the episode plan without merging",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    require_binaries()

    input_dir = args.input.expanduser().resolve()
    if not input_dir.is_dir():
        print(f"Error: input is not a directory: {input_dir}", file=sys.stderr)
        sys.exit(1)

    if args.target_minutes is not None and args.target_minutes <= 0:
        print("Error: --target-minutes must be positive", file=sys.stderr)
        sys.exit(1)
    if args.target_episodes is not None and args.target_episodes < 1:
        print("Error: --target-episodes must be at least 1", file=sys.stderr)
        sys.exit(1)

    codec: str = args.codec
    use_gpu = False
    if args.gpu:
        if nvenc_available(codec):
            use_gpu = True
            print(f"GPU encoding enabled ({encoder_name(True, codec)})")
        else:
            print(
                f"Warning: --gpu requested but {encoder_name(True, codec)} "
                "is unavailable; using CPU.",
                file=sys.stderr,
            )

    print(f"Scanning {input_dir}...")
    sketches = load_sketches(input_dir)
    print(f"Found {len(sketches)} sketches\n")

    if args.target_episodes is not None:
        try:
            episodes, target_seconds = pack_into_n_episodes(
                sketches, args.target_episodes
            )
        except ValueError as exc:
            print(f"Error: {exc}", file=sys.stderr)
            sys.exit(1)
    else:
        target_minutes = (
            args.target_minutes if args.target_minutes is not None else 40.0
        )
        target_seconds = target_minutes * 60.0
        episodes = pack_episodes(sketches, target_seconds)

    print_plan(episodes, target_seconds)

    if args.dry_run:
        print("Dry run — no files written.")
        return

    output_dir = args.output.expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    for index, episode in enumerate(episodes, start=1):
        out_path = output_dir / f"aflevering-{index:02d}.mp4"
        try:
            merge_episode(episode, out_path, use_gpu=use_gpu, codec=codec)
        except RuntimeError as exc:
            if use_gpu:
                print(
                    f"Warning: GPU encode failed for {out_path.name}; "
                    "falling back to CPU for remaining episodes.\n"
                    f"{exc}",
                    file=sys.stderr,
                )
                use_gpu = False
                try:
                    merge_episode(episode, out_path, use_gpu=False, codec=codec)
                except RuntimeError as cpu_exc:
                    print(str(cpu_exc), file=sys.stderr)
                    sys.exit(1)
            else:
                print(str(exc), file=sys.stderr)
                sys.exit(1)
        print(f"Wrote {out_path}")

    print(f"\nDone. {len(episodes)} episode(s) in {output_dir}")


if __name__ == "__main__":
    main()