"""FFmpeg front-end for NCMR song videos.

Wraps the ffmpeg invocations this project actually needs, so the long command
lines live in one place instead of in shell history and notes files.

Modes:

``mux``     attach an audio track to a video-only render, without re-encoding
            the video (the Remotion output is video-only; this is step 1)
``bake``    burn an ASS subtitle file into a video (step 2)
``still``   build a video from a single image plus audio, for the no-lyrics case
``strip``   pull an MP3 out of a finished video
``probe``   print a compact stream summary

Both ``mux`` and ``bake`` can be chained with ``--then-bake``/``--then-mux`` so a
render goes to finished file in one command.

Usage::

    python3 scripts/ffmpeg_tools.py probe out/pulseraw.mp4
    python3 scripts/ffmpeg_tools.py mux  out/pulseraw.mp4 song.wav -o raw.mkv --format flac
    python3 scripts/ffmpeg_tools.py bake raw.mkv subs.ass -o final.mp4
    python3 scripts/ffmpeg_tools.py still cover.png song.wav -o pv.mkv --fps 60 --format flac
    python3 scripts/ffmpeg_tools.py strip final.mp4 -o song.mp3
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from dataclasses import dataclass

if __package__ in (None, ""):
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from _common import (  # type: ignore[import-not-found]
        Die,
        Report,
        add_common_arguments,
        add_ffmpeg_argument,
        find_ffmpeg,
        main_wrapper,
        run,
        text_table,
    )
else:  # pragma: no cover
    from ._common import (
        Die,
        Report,
        add_common_arguments,
        add_ffmpeg_argument,
        find_ffmpeg,
        main_wrapper,
        run,
        text_table,
    )


# --------------------------------------------------------------------------- #
# Audio/codec presets
# --------------------------------------------------------------------------- #

#: ``--format`` choices, mapping to the ffmpeg audio arguments to use.
AUDIO_FORMATS = {
    # name:     (extension, container args, codec args)
    "aac": ("mp4", [], ["-c:a", "aac", "-b:a", "320k"]),
    "mp3": ("mp4", [], ["-c:a", "libmp3lame", "-b:a", "320k"]),
    "flac": ("mkv", [], ["-c:a", "flac", "-compression_level", "12"]),
    "opus": ("mkv", [], ["-c:a", "libopus", "-b:a", "192k"]),
    "copy": (None, [], ["-c:a", "copy"]),
}

#: ``--video`` choices. ``auto`` picks a hardware encoder when one works.
VIDEO_ENCODERS = {
    "auto": None,
    "x264": ["-c:v", "libx264"],
    "x265": ["-c:v", "libx265"],
    "copy": ["-c:v", "copy"],
}

VAAPI_ENCODERS = {
    "h264": ["-c:v", "h264_vaapi"],
    "hevc": ["-c:v", "hevc_vaapi"],
    "av1": ["-c:v", "av1_vaapi"],
}

#: Software fallback quality knob, applied only for libx264/libx265.
X264_QUALITY = ["-preset", "medium", "-crf", "18"]

DEFAULT_DEVICE = "/dev/dri/renderD128"


@dataclass
class HwContext:
    """Hardware acceleration state: which encoder to use and how to init it."""

    encoder: str
    device: str | None = None
    is_vaapi: bool = False
    note: str = ""

    def input_args(self) -> list[str]:
        """Args that go before ``-i``."""
        if self.is_vaapi and self.device:
            return ["-init_hw_device", f"vaapi=va:{self.device}", "-filter_hw_device", "va"]

        return []

    def filter_suffix(self) -> str:
        """Filters appended to a user filter chain to move frames onto the GPU."""
        if self.is_vaapi:
            return "format=nv12,hwupload"

        return ""


# --------------------------------------------------------------------------- #
# Encoder selection
# --------------------------------------------------------------------------- #


def _vaapi_device_readable(device: str) -> bool:
    try:
        fd = os.open(device, os.O_RDWR)
    except OSError:
        return False

    os.close(fd)

    return True


_DEVICE_CACHE: dict[str, bool] = {}


def _encoder_works(ffmpeg: str, encoder: str, extra: list[str], device_key: str = "") -> bool:
    """Try a two-frame encode. Results are memoised per process.

    VA-API encoders only accept GPU-resident frames, so the probe must upload the
    test pattern first. Without ``format=nv12,hwupload`` ffmpeg fails with
    "Function not implemented" even though the encoder works fine — which would
    make this fall back to software on a machine that has no problem at all.
    """
    cache_key = f"{encoder}|{device_key}"
    if cache_key in _DEVICE_CACHE:
        return _DEVICE_CACHE[cache_key]

    vf = "format=nv12,hwupload" if encoder.endswith("_vaapi") else None

    cmd = [
        ffmpeg, "-hide_banner", "-loglevel", "error", *extra,
        "-f", "lavfi", "-i", "color=c=black:s=256x256:r=30:d=1",
    ]
    if vf:
        cmd += ["-vf", vf]
    cmd += ["-c:v", encoder, "-frames:v", "2", "-f", "null", "-"]

    try:
        ok = subprocess.run(cmd, capture_output=True, timeout=60).returncode == 0
    except (OSError, subprocess.TimeoutExpired):
        ok = False

    _DEVICE_CACHE[cache_key] = ok

    return ok


#: Hardware encoders worth trying for each codec family, in preference order.
HW_ENCODER_CHAIN: dict[str, tuple[tuple[str, str], ...]] = {
    # codec: ((encoder, init-family), ...)
    "h264": (("h264_vaapi", "vaapi"), ("h264_nvenc", "nvenc"), ("h264_qsv", "qsv")),
    "hevc": (("hevc_vaapi", "vaapi"), ("hevc_nvenc", "nvenc"), ("hevc_qsv", "qsv")),
    "av1": (("av1_vaapi", "vaapi"), ("av1_nvenc", "nvenc"), ("av1_qsv", "qsv")),
}


def pick_acceleration(
    ffmpeg: str,
    report: Report,
    *,
    video: str = "auto",
    codec: str = "h264",
    device: str | None = None,
    allow_software: bool = True,
) -> HwContext:
    """Decide how to encode, honouring an explicit ``--video`` choice."""
    device = device or os.environ.get("HW_DEVICE") or DEFAULT_DEVICE

    if video == "copy":
        return HwContext(encoder="copy")

    if video == "auto":
        chain = HW_ENCODER_CHAIN.get(codec)
        if chain is None:
            raise Die(f"unknown --codec {codec!r}")

        attempted: list[str] = []

        for encoder, family in chain:
            if family == "vaapi":
                if not _vaapi_device_readable(device):
                    attempted.append(f"{encoder} (unreadable {device})")
                    continue
                init = ["-init_hw_device", f"vaapi=va:{device}", "-filter_hw_device", "va"]
                if _encoder_works(ffmpeg, encoder, init, device):
                    return HwContext(
                        encoder=encoder,
                        device=device,
                        is_vaapi=True,
                        note="hardware (VA-API)",
                    )
                attempted.append(f"{encoder} (device rejects it)")
            else:
                init = []
                if _encoder_works(ffmpeg, encoder, init, family):
                    return HwContext(encoder=encoder, note=f"hardware ({family})")
                attempted.append(f"{encoder} (unsupported)")

        if not allow_software:
            raise Die(
                "no working hardware encoder found and software fallback is disabled; "
                f"tried {', '.join(attempted)}. Diagnose with: python3 scripts/hwprobe.py"
            )

        report.verbose(f"no hardware encoder for {codec}: {', '.join(attempted)}")

        return HwContext(encoder="libx264" if codec == "h264" else "libx265", note="software")

    if video in ("x264", "x265"):
        return HwContext(encoder=VIDEO_ENCODERS[video][1], note=f"software ({video})")

    raise Die(f"unknown --video {video!r}")


def video_args(context: HwContext, *, quality: bool = True) -> list[str]:
    """Encoder arguments for the selected context."""
    if context.encoder == "copy":
        return ["-c:v", "copy"]

    args = ["-c:v", context.encoder]

    if context.encoder.startswith("libx") and quality:
        args += X264_QUALITY

    return args


def pix_fmt_args(context: HwContext) -> list[str]:
    # Other pixel formats come from the VA-API upload filter; without a filter
    # chain, hardware encoders still need 4:2:0 input.
    if context.is_vaapi:
        return []

    return ["-pix_fmt", "yuv420p"]


def audio_args(fmt: str) -> tuple[list[str], str]:
    if fmt not in AUDIO_FORMATS:
        raise Die(f"unknown --format {fmt!r}; choose from {', '.join(AUDIO_FORMATS)}")

    extension, container, codec = AUDIO_FORMATS[fmt]

    return container + codec, extension or ("")


def resolve_output(output: str | None, stem: str, fmt: str, suffix: str) -> str:
    """Work out the output filename, fixing up the extension for the container."""
    if output:
        return output

    extension = AUDIO_FORMATS[fmt][0]
    if extension and not stem.lower().endswith(f".{extension}"):
        return f"{stem}{suffix}.{extension}"

    return f"{stem}{suffix}"


def build_filter(context: HwContext, *filters: str) -> str | None:
    """Join filter fragments, appending the VA-API upload step when needed."""
    parts = [f for f in filters if f]
    upload = context.filter_suffix()

    if upload:
        parts.append(upload)

    return ",".join(parts) if parts else None


# --------------------------------------------------------------------------- #
# Commands
# --------------------------------------------------------------------------- #


def cmd_mux(args: argparse.Namespace, context: HwContext, report: Report) -> int:
    """Attach audio to a video-only file. Video stream is stream-copied."""
    codec_args, _ = audio_args(args.format)
    output = resolve_output(args.output, os.path.splitext(args.input_video)[0], args.format, "_mux")

    filter_chain = build_filter(context, *args.vf) if context.is_vaapi else (",".join(args.vf) or None)
    video = video_args(context, quality=False) if context.is_vaapi else ["-c:v", "copy"]

    cmd = [
        args.ffmpeg, "-y",
        *context.input_args(),
        "-i", args.input_video,
        "-i", args.audio,
        "-map", "0:v:0", "-map", "1:a:0",
        *video,
        *codec_args,
    ]

    if filter_chain:
        cmd += ["-vf", filter_chain]

    cmd += ["-shortest", output]

    run(cmd, report=report, dry_run=args.dry_run)
    if not args.dry_run:
        report.ok(f"muxed -> {output}")

    return 0


def cmd_bake(args: argparse.Namespace, context: HwContext, report: Report) -> int:
    """Burn ASS subtitles into the video stream."""
    if not os.path.isfile(args.subtitles):
        raise Die(f"no such subtitle file: {args.subtitles}")

    # libass needs the filename escaped for the filtergraph; ':' and '\' and "'"
    # all have meaning there.
    escaped = (
        args.subtitles.replace("\\", "\\\\")
        .replace(":", "\\:")
        .replace("'", r"\'")
        .replace("[", r"\[")
        .replace("]", r"\]")
        .replace(",", r"\,")
    )

    # Burning subtitles means changing pixels, so the video stream cannot be copied.
    # Fail early with a clear message instead of letting ffmpeg refuse the filter.
    if context.encoder == "copy":
        raise Die(
            "cannot bake subtitles with --video copy: burning them in requires "
            "re-encoding. Use --video auto (default) or --video x264."
        )

    filter_chain = build_filter(context, f"ass={escaped}", *args.vf)
    output = args.output or _append_suffix(args.input_video, "_sub")

    codec_args, _ = audio_args(args.format) if args.audio_format != "keep" else (["-c:a", "copy"], None)

    cmd = [
        args.ffmpeg, "-y",
        *context.input_args(),
        "-i", args.input_video,
        "-vf", filter_chain or "null",
        *video_args(context),
        *pix_fmt_args(context),
        *codec_args,
        "-shortest" if args.shortest else "-map_metadata", "-1",
        output,
    ]

    run(cmd, report=report, dry_run=args.dry_run)
    if not args.dry_run:
        report.ok(f"subtitles baked -> {output}")

    return 0


def cmd_still(args: argparse.Namespace, context: HwContext, report: Report) -> int:
    """Build a video from one image and one audio file (the no-lyrics render)."""
    codec_args, _ = audio_args(args.format)
    output = resolve_output(args.output, os.path.splitext(args.image)[0], args.format, "_pv")

    filters = [f"fps={args.fps}", f"scale={args.width}:{args.height}:force_original_aspect_ratio=decrease", "pad=ceil(iw/2)*2:ceil(ih/2)*2"]

    if context.is_vaapi:
        filter_chain = build_filter(context, *filters)
        cmd = [args.ffmpeg, "-y", "-loop", "1", "-i", args.image, "-i", args.audio, *context.input_args(), "-vf", filter_chain]
    else:
        cmd = [args.ffmpeg, "-y", "-loop", "1", "-i", args.image, "-i", args.audio, "-vf", ",".join(filters)]

    cmd += ["-r", str(args.fps), *video_args(context), *pix_fmt_args(context), *codec_args, "-shortest", output]

    run(cmd, report=report, dry_run=args.dry_run)
    if not args.dry_run:
        report.ok(f"still-image video -> {output}")

    return 0


def cmd_strip(args: argparse.Namespace, context: HwContext | None, report: Report) -> int:
    """Extract the audio track as MP3."""
    output = args.output or os.path.splitext(args.input_video)[0] + ".mp3"

    cmd = [
        args.ffmpeg, "-y", "-i", args.input_video, "-vn",
        "-ar", str(args.sample_rate), "-ac", str(args.channels),
        "-c:a", "libmp3lame", "-b:a", args.bitrate,
        output,
    ]

    run(cmd, report=report, dry_run=args.dry_run)
    if not args.dry_run:
        report.ok(f"audio -> {output}")

    return 0


def cmd_probe(args: argparse.Namespace, context: HwContext | None, report: Report) -> int:
    """Print a compact stream summary for a media file."""
    if not os.path.isfile(args.media[0]):
        raise Die(f"no such file: {args.media[0]}")

    cmd = [
        "ffprobe" if args.ffmpeg.endswith("ffmpeg") else args.ffmpeg.replace("ffmpeg", "ffprobe"),
        "-hide_banner", "-loglevel", "error",
        "-show_entries", "stream=index,codec_type,codec_name,width,height,r_frame_rate,channels,sample_rate,bit_rate",
        "-show_entries", "format=duration,size,format_name",
        "-of", "json",
        *args.media,
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
    except OSError as exc:
        raise Die(f"could not run ffprobe: {exc}") from exc

    if result.returncode != 0:
        raise Die(f"ffprobe failed: {(result.stderr or '').strip()}")

    data = json.loads(result.stdout or "{}")

    rows = []
    for stream in data.get("streams", []):
        kind = stream.get("codec_type", "?")
        if kind == "video":
            detail = f"{stream.get('width')}x{stream.get('height')} @ {stream.get('r_frame_rate')}"
        elif kind == "audio":
            detail = f"{stream.get('channels')}ch @ {stream.get('sample_rate')} Hz"
        else:
            detail = ""
        rows.append((str(stream.get("index", "")), kind, stream.get("codec_name", "?"), detail))

    report.info(text_table(("stream", "type", "codec", "detail"), rows))

    fmt = data.get("format", {})
    duration = fmt.get("duration")
    size = fmt.get("size")
    if duration:
        report.info(f"\ncontainer: {fmt.get('format_name', '?')}")
        report.info(f"duration:  {float(duration):.2f}s")
    if size:
        report.info(f"size:      {int(size):,} bytes")

    return 0


def _append_suffix(path: str, suffix: str) -> str:
    stem, extension = os.path.splitext(path)

    return f"{stem}{suffix}{extension}"


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #


def _add_video_options(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--video",
        default=None,
        choices=("auto", "x264", "x265", "copy"),
        help="video encoder: auto (hardware if possible), x264, x265, or copy (default: auto)",
    )
    parser.add_argument(
        "--codec",
        default=None,
        choices=("h264", "hevc", "av1"),
        help="hardware codec family to use when --video auto (default: h264)",
    )
    parser.add_argument(
        "--device",
        default=None,
        metavar="PATH",
        help=f"VA-API render node (default: $HW_DEVICE or {DEFAULT_DEVICE})",
    )
    parser.add_argument(
        "--vf",
        action="append",
        default=[],
        metavar="FILTER",
        help="extra video filter, repeatable (appended after the built-in chain)",
    )


def _add_audio_options(parser: argparse.ArgumentParser, *, default: str = "flac") -> None:
    parser.add_argument(
        "-f",
        "--format",
        default=default,
        choices=tuple(AUDIO_FORMATS),
        help=f"audio format preset (default: {default})",
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="ffmpeg_tools",
        description="FFmpeg front-end for NCMR song videos (mux, bake, still, strip, probe).",
        epilog="Run `%(prog)s hwprobe` for a separate encoder report, or see scripts/hwprobe.py.",
    )
    add_common_arguments(parser)
    add_ffmpeg_argument(parser)

    subparsers = parser.add_subparsers(dest="mode", metavar="MODE")

    mux = subparsers.add_parser("mux", help="attach audio to a video-only file, video copied as-is")
    add_common_arguments(mux, suppress_defaults=True)
    mux.add_argument("input_video", help="video-only input (e.g. the Remotion render)")
    mux.add_argument("audio", help="audio input (e.g. the exported song)")
    mux.add_argument("-o", "--output", help="output path")
    _add_audio_options(mux)
    _add_video_options(mux)
    mux.set_defaults(func=cmd_mux)

    bake = subparsers.add_parser("bake", help="burn an ASS subtitle file into a video")
    add_common_arguments(bake, suppress_defaults=True)
    bake.add_argument("input_video", help="input video")
    bake.add_argument("subtitles", help="ASS subtitle file")
    bake.add_argument("-o", "--output", help="output path (default: <video>_sub.<ext>)")
    _add_audio_options(bake, default="aac")
    bake.add_argument(
        "--audio-format",
        default="keep",
        choices=("keep", *AUDIO_FORMATS),
        help="re-encode audio to this preset, or keep it as-is (default: keep)",
    )
    bake.add_argument("--shortest", action="store_true", help="stop at the shortest stream")
    _add_video_options(bake)
    bake.set_defaults(func=cmd_bake)

    still = subparsers.add_parser("still", help="one image + audio -> video")
    add_common_arguments(still, suppress_defaults=True)
    still.add_argument("image", help="cover art image")
    still.add_argument("audio", help="audio file")
    still.add_argument("-o", "--output", help="output path")
    still.add_argument("--fps", type=int, default=60, help="frame rate (default: 60)")
    still.add_argument("--width", type=int, default=1920, help="target width (default: 1920)")
    still.add_argument("--height", type=int, default=1080, help="target height (default: 1080)")
    _add_audio_options(still)
    _add_video_options(still)
    still.set_defaults(func=cmd_still)

    strip = subparsers.add_parser("strip", help="extract audio from a video as MP3")
    add_common_arguments(strip, suppress_defaults=True)
    strip.add_argument("input_video", help="input video")
    strip.add_argument("-o", "--output", help="output MP3 path")
    strip.add_argument("--sample-rate", type=int, default=44100, help="(default: 44100)")
    strip.add_argument("--channels", type=int, default=2, help="(default: 2)")
    strip.add_argument("--bitrate", default="320k", help="(default: 320k)")
    strip.set_defaults(func=cmd_strip)

    probe = subparsers.add_parser("probe", help="print a compact stream summary")
    add_common_arguments(probe, suppress_defaults=True)
    probe.add_argument("media", nargs=1, help="media file to inspect")
    probe.set_defaults(func=cmd_probe)

    return parser


@main_wrapper
def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    report = Report(quiet=args.quiet)

    if not getattr(args, "func", None):
        parser.print_help()
        return 2

    args.ffmpeg = find_ffmpeg(args.ffmpeg)

    # strip and probe never encode video, so they have no --video/--codec/--device
    # and no acceleration context to work out.
    if not hasattr(args, "video"):
        return args.func(args, None, report)

    # --video/--codec/--device live on the subparsers, so the namespace already
    # carries whatever was parsed. Supply the documented defaults here rather than
    # duplicating them across every subparser.
    video = args.video if args.video is not None else "auto"
    codec = args.codec if args.codec is not None else "h264"
    device = args.device or os.environ.get("HW_DEVICE")

    context = pick_acceleration(
        args.ffmpeg,
        report,
        video=video,
        codec=codec,
        device=device,
    )

    if context.note and args.mode != "probe":
        report.info(f"encoder: {context.encoder} — {context.note}")

    return args.func(args, context, report)


def _supplied_options(argv: list[str]) -> dict[str, str]:
    """Map ``--flag`` to its value for flags present on the command line.

    Retained for tests and external callers; ``main`` no longer needs it now that
    the encoder options are parsed by argparse rather than sniffed out of argv.
    """
    supplied: dict[str, str] = {}
    i = 0
    while i < len(argv):
        token = argv[i]
        if token.startswith("--") and "=" in token:
            flag, _, value = token.partition("=")
            supplied[flag] = value
        elif token.startswith("--"):
            if i + 1 < len(argv):
                supplied[token] = argv[i + 1]
            i += 1
        i += 1

    return supplied


def _decide(flag: str, supplied: dict[str, str], namespace_value, fallback):
    """Prefer a flag supplied on the command line, then the namespace, then default."""
    if flag in supplied:
        return supplied[flag]

    return namespace_value if namespace_value is not None else fallback


if __name__ == "__main__":
    raise SystemExit(main())