import os
import signal
import sys
import threading
from types import FrameType
from typing import Any, Literal, cast

import ffmpeg  # type: ignore[import-untyped]

_ffmpeg: Any = ffmpeg

WIDTH = 640
HEIGHT = 360
FRAME_SIZE = WIDTH * HEIGHT * 3
OUTPUT = "/tmp/video"
IMAGE_HOLD_SECONDS = 10

# Set when the user asks us to shut down (Ctrl-C / SIGTERM).
_shutdown = threading.Event()

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp", ".tiff", ".tif"}
VIDEO_EXTS = {".mkv", ".mp4", ".mov", ".avi", ".webm", ".m4v", ".flv", ".wmv", ".mpg", ".mpeg"}

Kind = Literal["image", "video"]


def classify(path: str) -> Kind | None:
    ext = os.path.splitext(path)[1].lower()
    if ext in IMAGE_EXTS:
        return "image"
    if ext in VIDEO_EXTS:
        return "video"
    return None


def collect_media(input_path: str) -> list[tuple[str, Kind]]:
    """Return a sorted list of (path, kind) for the given file or folder."""
    if os.path.isfile(input_path):
        kind = classify(input_path)
        if kind is None:
            print(f"unsupported file type: {input_path}")
            sys.exit(1)
        return [(input_path, kind)]

    if os.path.isdir(input_path):
        media: list[tuple[str, Kind]] = []
        for name in sorted(os.listdir(input_path)):
            full = os.path.join(input_path, name)
            if not os.path.isfile(full):
                continue
            kind = classify(full)
            if kind is not None:
                media.append((full, kind))
        if not media:
            print(f"no images or videos found in: {input_path}")
            sys.exit(1)
        return media

    print(f"no such file or folder: {input_path}")
    sys.exit(1)


def write_frame(fd: int, frame: bytes) -> None:
    os.lseek(fd, 0, os.SEEK_SET)
    os.write(fd, frame)
    os.lseek(fd, 0, os.SEEK_SET)


def show_image(fd: int, path: str) -> None:
    """Decode a single image, scale to the display, and hold it for a while."""
    print(f"image: {path}")
    stream: Any = (
        _ffmpeg
        .input(path)
        .filter("scale", WIDTH, HEIGHT)
        .output("pipe:", format="rawvideo", pix_fmt="rgb24", vframes=1)
    )
    result = cast("tuple[bytes, bytes]", stream.run(capture_stdout=True, quiet=True))
    out = result[0]
    if len(out) < FRAME_SIZE:
        print(f"skipping {path}: decoded {len(out)} bytes, expected {FRAME_SIZE}")
        return
    write_frame(fd, out[:FRAME_SIZE])
    # Interruptible hold: returns immediately once shutdown is requested.
    _shutdown.wait(IMAGE_HOLD_SECONDS)


def play_video(fd: int, path: str) -> None:
    """Stream a video through to the end at real-time pace."""
    print(f"video: {path}")
    stream: Any = (
        _ffmpeg
        .input(path, re=None)
        .filter("scale", WIDTH, HEIGHT)
        .output("pipe:", format="rawvideo", pix_fmt="rgb24")
    )
    process: Any = stream.run_async(pipe_stdout=True, quiet=False)
    stdout: Any = process.stdout
    try:
        while not _shutdown.is_set():
            frame = cast("bytes", stdout.read(FRAME_SIZE))
            if len(frame) != FRAME_SIZE:
                break
            write_frame(fd, frame)
    finally:
        # Stop ffmpeg so it doesn't keep decoding after we stop reading.
        process.terminate()
        stdout.close()
        process.wait()


def _request_shutdown(signum: int, frame: FrameType | None) -> None:
    print("\nshutting down...")
    _shutdown.set()


def main() -> None:
    if len(sys.argv) != 2:
        print(f"usage: {sys.argv[0]} <file-or-folder>")
        sys.exit(1)

    media = collect_media(sys.argv[1])

    signal.signal(signal.SIGINT, _request_shutdown)
    signal.signal(signal.SIGTERM, _request_shutdown)

    fd = os.open(OUTPUT, os.O_RDWR | os.O_CREAT)
    # Ensure the file is at least one frame long so the C side can mmap it.
    os.ftruncate(fd, FRAME_SIZE)

    print(f"{WIDTH}x{HEIGHT}x3={FRAME_SIZE}")
    print(OUTPUT)

    try:
        while not _shutdown.is_set():
            for path, kind in media:
                if _shutdown.is_set():
                    break
                if kind == "image":
                    show_image(fd, path)
                else:
                    play_video(fd, path)
    finally:
        os.close(fd)


if __name__ == "__main__":
    main()