import argparse
import re
from pathlib import Path

import imageio.v2 as imageio
import numpy as np
from PIL import Image


def _natural_key(text):
    parts = re.split(r"(\d+)", text)
    return [int(p) if p.isdigit() else p.lower() for p in parts]


def _iter_images(folder, extensions, natural_sort):
    paths = [
        p
        for p in Path(folder).iterdir()
        if p.is_file() and p.suffix.lower() in extensions
    ]
    if natural_sort:
        return sorted(paths, key=lambda p: _natural_key(p.name))
    return sorted(paths, key=lambda p: p.name.lower())


def images_to_mp4(
    input_dir,
    output_path,
    fps,
    extensions,
    natural_sort,
    resize_to_first,
    codec,
):
    images = _iter_images(input_dir, extensions, natural_sort)
    if not images:
        raise ValueError(f"No images found in {input_dir}")

    first = Image.open(images[0]).convert("RGB")
    target_size = first.size

    with imageio.get_writer(output_path, fps=fps, codec=codec) as writer:
        writer.append_data(np.array(first))
        for path in images[1:]:
            img = Image.open(path).convert("RGB")
            if img.size != target_size:
                if resize_to_first:
                    img = img.resize(target_size, Image.Resampling.LANCZOS)
                else:
                    raise ValueError(
                        f"Size mismatch: {path.name} is {img.size}, expected {target_size}"
                    )
            writer.append_data(np.array(img))


def main():
    parser = argparse.ArgumentParser(
        description="Create an MP4 video from images in a folder."
    )
    parser.add_argument("input_dir", help="Folder containing images.")
    parser.add_argument("output", help="Output MP4 path, e.g. out.mp4.")
    parser.add_argument(
        "--fps", type=float, default=24, help="Frames per second (default: 24)."
    )
    parser.add_argument(
        "--ext",
        action="append",
        default=[".jpg", ".jpeg", ".png"],
        help="Allowed extensions (repeatable). Default: .jpg .jpeg .png",
    )
    parser.add_argument(
        "--natural",
        action="store_true",
        help="Use natural sort for filenames (e.g. img2 before img10).",
    )
    parser.add_argument(
        "--no-resize",
        action="store_true",
        help="Fail if image sizes differ instead of resizing to the first image.",
    )
    parser.add_argument(
        "--codec",
        default="libx264",
        help="Video codec for MP4 (default: libx264).",
    )
    args = parser.parse_args()

    images_to_mp4(
        args.input_dir,
        args.output,
        fps=args.fps,
        extensions={e.lower() for e in args.ext},
        natural_sort=args.natural,
        resize_to_first=not args.no_resize,
        codec=args.codec,
    )


if __name__ == "__main__":
    main()