#!/usr/bin/env python3
"""Batch-convert all MP4 files in a folder to GIF using ffmpeg.

Example:
    python3 convert_all_mp4_to_gif.py . --fps 12 --width 500
"""

from __future__ import annotations

import argparse
import shutil
import subprocess
import tempfile
from pathlib import Path


MB = 1024 * 1024


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Convert all MP4 files in a folder to GIF files."
    )
    parser.add_argument(
        "folder",
        nargs="?",
        default=".",
        help="Folder containing MP4 files (default: current folder).",
    )
    parser.add_argument(
        "--fps",
        type=int,
        default=12,
        help="GIF frame rate (default: 12).",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=500,
        help="Output GIF width in pixels (default: 500). Height is auto.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing GIF files.",
    )
    parser.add_argument(
        "--max-size-mb",
        type=float,
        default=2.0,
        help="Maximum GIF size in MB (default: 2.0).",
    )
    parser.add_argument(
        "--disable-size-target",
        action="store_true",
        help="Disable automatic tuning to max-size-mb and use direct conversion.",
    )
    parser.add_argument(
        "--min-width",
        type=int,
        default=160,
        help="Minimum width used when tuning size (default: 160).",
    )
    return parser


def ffmpeg_exists() -> bool:
    return shutil.which("ffmpeg") is not None


def build_vf(fps: int, width: int, max_colors: int) -> str:
    return (
        f"fps={fps},scale={width}:-1:flags=lanczos,"
        "split[s0][s1];"
        f"[s0]palettegen=max_colors={max_colors}[p];"
        "[s1][p]paletteuse=dither=bayer:bayer_scale=5"
    )


def run_ffmpeg(mp4_path: Path, output_path: Path, fps: int, width: int, max_colors: int, overwrite: bool) -> bool:
    vf = build_vf(fps=fps, width=width, max_colors=max_colors)

    cmd = [
        "ffmpeg",
        "-y" if overwrite else "-n",
        "-i",
        str(mp4_path),
        "-vf",
        vf,
        str(output_path),
    ]

    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if result.returncode != 0:
        return False
    return True


def convert_one_direct(mp4_path: Path, gif_path: Path, fps: int, width: int, overwrite: bool) -> bool:
    if gif_path.exists() and not overwrite:
        print(f"Skipping existing GIF: {gif_path.name}")
        return True

    print(f"Converting: {mp4_path.name} -> {gif_path.name}")
    ok = run_ffmpeg(
        mp4_path=mp4_path,
        output_path=gif_path,
        fps=fps,
        width=width,
        max_colors=128,
        overwrite=overwrite,
    )
    if not ok:
        print(f"Failed: {mp4_path.name}")
    return ok


def score_candidate(size_bytes: int, target_bytes: int) -> int:
    return abs(target_bytes - size_bytes)


def search_width_for_combo(
    mp4_path: Path,
    tmp_dir_path: Path,
    target_bytes: int,
    fps: int,
    max_colors: int,
    min_width: int,
    max_width: int,
) -> tuple[tuple[int, int] | None, tuple[int, int] | None]:
    low = min_width
    high = max_width
    best_under: tuple[int, int] | None = None
    best_over: tuple[int, int] | None = None

    while low <= high:
        candidate_width = (low + high) // 2
        if candidate_width % 2 != 0:
            candidate_width -= 1
        if candidate_width < min_width:
            candidate_width = min_width

        temp_gif = tmp_dir_path / "trial.gif"
        if temp_gif.exists():
            temp_gif.unlink()

        ok = run_ffmpeg(
            mp4_path=mp4_path,
            output_path=temp_gif,
            fps=fps,
            width=candidate_width,
            max_colors=max_colors,
            overwrite=True,
        )
        if not ok or not temp_gif.exists():
            high = candidate_width - 2
            continue

        size_bytes = temp_gif.stat().st_size
        if size_bytes <= target_bytes:
            if best_under is None or size_bytes > best_under[1]:
                best_under = (candidate_width, size_bytes)
            low = candidate_width + 2
        else:
            if best_over is None or size_bytes < best_over[1]:
                best_over = (candidate_width, size_bytes)
            high = candidate_width - 2

    return best_under, best_over


def clamp_even_width(width: int, floor: int = 64) -> int:
    width = max(floor, width)
    if width % 2 != 0:
        width -= 1
    return max(floor, width)


def enforce_size_cap(
    mp4_path: Path,
    gif_path: Path,
    start_fps: int,
    start_width: int,
    start_colors: int,
    target_bytes: int,
    overwrite: bool,
) -> tuple[bool, int, int, int, int]:
    fps = max(3, start_fps)
    width = clamp_even_width(start_width)
    colors = max(16, start_colors)

    color_steps = [128, 96, 64, 48, 32, 24, 16]

    for _ in range(12):
        ok = run_ffmpeg(
            mp4_path=mp4_path,
            output_path=gif_path,
            fps=fps,
            width=width,
            max_colors=colors,
            overwrite=overwrite,
        )
        if not ok or not gif_path.exists():
            return False, fps, width, colors, -1

        size_bytes = gif_path.stat().st_size
        if size_bytes <= target_bytes:
            return True, fps, width, colors, size_bytes

        ratio = target_bytes / size_bytes
        proposed_width = clamp_even_width(int(width * (ratio ** 0.5) * 0.97))
        if proposed_width < width:
            width = proposed_width
            continue

        if fps > 4:
            fps -= 1
            continue

        if colors in color_steps:
            idx = color_steps.index(colors)
            if idx < len(color_steps) - 1:
                colors = color_steps[idx + 1]
                continue
        else:
            colors = max(16, colors - 16)
            continue

        if width > 64:
            width = clamp_even_width(width - 20)
            continue

        return False, fps, width, colors, size_bytes

    final_size = gif_path.stat().st_size if gif_path.exists() else -1
    return False, fps, width, colors, final_size


def convert_one_target_size(
    mp4_path: Path,
    gif_path: Path,
    fps: int,
    width: int,
    max_size_mb: float,
    min_width: int,
    overwrite: bool,
) -> bool:
    if gif_path.exists() and not overwrite:
        print(f"Skipping existing GIF: {gif_path.name}")
        return True

    target_bytes = int(max_size_mb * MB)
    if target_bytes <= 0:
        print("Invalid max-size-mb value. It must be > 0.")
        return False

    min_width = max(64, min_width)
    width = max(min_width, width)

    # Keep the trial count low: try a small set of quality combos in priority order.
    combo_candidates: list[tuple[int, int]] = []
    for item in [
        (fps, 96),
        (fps, 128),
        (max(4, fps - 2), 96),
        (8, 96),
        (6, 64),
        (5, 48),
    ]:
        candidate_fps, candidate_colors = item
        if candidate_fps >= 3 and item not in combo_candidates:
            combo_candidates.append(item)

    best_under: tuple[int, int, int, int] | None = None
    best_over: tuple[int, int, int, int] | None = None

    print(
        f"Converting with size target: {mp4_path.name} -> {gif_path.name} "
        f"(max {max_size_mb:.2f} MB)"
    )

    with tempfile.TemporaryDirectory(prefix="gif_tune_") as tmp_dir:
        tmp_dir_path = Path(tmp_dir)

        for candidate_fps, candidate_colors in combo_candidates:
            local_under, local_over = search_width_for_combo(
                mp4_path=mp4_path,
                tmp_dir_path=tmp_dir_path,
                target_bytes=target_bytes,
                fps=candidate_fps,
                max_colors=candidate_colors,
                min_width=min_width,
                max_width=width,
            )

            if local_under is not None:
                under_width, under_size = local_under
                candidate = (candidate_fps, under_width, candidate_colors, under_size)
                if best_under is None or score_candidate(under_size, target_bytes) < score_candidate(best_under[3], target_bytes):
                    best_under = candidate

                # Good enough threshold to avoid extra slow trials.
                if under_size >= int(target_bytes * 0.96):
                    break

            if local_over is not None:
                over_width, over_size = local_over
                candidate = (candidate_fps, over_width, candidate_colors, over_size)
                if best_over is None or over_size < best_over[3]:
                    best_over = candidate

    chosen = best_under if best_under is not None else best_over
    if chosen is None:
        print(f"Failed: {mp4_path.name}")
        return False

    chosen_fps, chosen_width, chosen_colors, _ = chosen
    ok, final_fps, final_width, final_colors, final_size_bytes = enforce_size_cap(
        mp4_path=mp4_path,
        gif_path=gif_path,
        start_fps=chosen_fps,
        start_width=chosen_width,
        start_colors=chosen_colors,
        target_bytes=target_bytes,
        overwrite=overwrite,
    )
    if not ok or not gif_path.exists() or final_size_bytes < 0:
        print(f"Failed final encode: {mp4_path.name}")
        return False

    final_size_mb = final_size_bytes / MB
    print(
        "Selected settings "
        f"fps={final_fps}, width={final_width}, colors={final_colors} -> "
        f"{final_size_mb:.3f} MB"
    )
    return True


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    folder = Path(args.folder).expanduser().resolve()
    if not folder.exists() or not folder.is_dir():
        print(f"Folder not found: {folder}")
        return 1

    if not ffmpeg_exists():
        print("ffmpeg was not found in PATH. Please install ffmpeg first.")
        return 1

    mp4_files = sorted([p for p in folder.iterdir() if p.is_file() and p.suffix.lower() == ".mp4"])
    if not mp4_files:
        print(f"No MP4 files found in: {folder}")
        return 0

    success = 0
    failed = 0

    for mp4_file in mp4_files:
        gif_file = mp4_file.with_suffix(".gif")
        if args.disable_size_target:
            ok = convert_one_direct(mp4_file, gif_file, args.fps, args.width, args.overwrite)
        else:
            ok = convert_one_target_size(
                mp4_file,
                gif_file,
                args.fps,
                args.width,
                args.max_size_mb,
                args.min_width,
                args.overwrite,
            )

        if ok:
            success += 1
        else:
            failed += 1

    print(f"Done. Success: {success}, Failed: {failed}")
    return 0 if failed == 0 else 2


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