#!/usr/bin/env python3
"""Simple flat-folder photo/video movie maker using ffmpeg.

Usage:
    python movmaker.py input --output movie.mp4

Put images, videos, audio files, and optionally data.txt in one directory.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import shlex
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}
VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".avi", ".webm", ".m4v"}
AUDIO_EXTS = {".mp3", ".wav", ".m4a", ".aac", ".flac", ".ogg"}
DATA_NAMES = {"data.txt", "info.txt", "title.txt"}


@dataclass
class MovieData:
    title: str = ""
    date: str = ""
    place: str = ""
    captions: dict[str, str] | None = None


@dataclass
class MediaItem:
    path: Path
    kind: str  # image or video
    duration: float


def run(cmd: list[str]) -> None:
    print("+", shlex.join(cmd))
    subprocess.run(cmd, check=True)


def capture(cmd: list[str]) -> str:
    return subprocess.check_output(cmd, text=True).strip()


def require_tool(name: str) -> None:
    try:
        subprocess.run([name, "-version"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
    except Exception:
        raise SystemExit(f"Error: required tool not found: {name}")


def parse_resolution(value: str) -> tuple[int, int]:
    try:
        w, h = value.lower().split("x", 1)
        return int(w), int(h)
    except Exception:
        raise argparse.ArgumentTypeError("resolution must look like 1920x1080")


def ffprobe_duration(path: Path) -> float:
    try:
        out = capture([
            "ffprobe", "-v", "error",
            "-show_entries", "format=duration",
            "-of", "json",
            str(path),
        ])
        data = json.loads(out)
        return max(0.1, float(data["format"]["duration"]))
    except Exception:
        return 3.0


def natural_key(path: Path) -> list[object]:
    # Simple natural-ish sort: IMG_2 before IMG_10.
    import re
    parts = re.split(r"(\d+)", path.name.lower())
    return [int(p) if p.isdigit() else p for p in parts]


def scan_input(input_dir: Path, image_duration: float) -> tuple[list[MediaItem], list[Path], MovieData]:
    if not input_dir.is_dir():
        raise SystemExit(f"Input directory does not exist: {input_dir}")

    media: list[MediaItem] = []
    audio: list[Path] = []
    data_file: Path | None = None

    for path in sorted(input_dir.iterdir(), key=natural_key):
        if not path.is_file():
            continue
        ext = path.suffix.lower()
        if path.name.lower() in DATA_NAMES:
            data_file = path
        elif ext in IMAGE_EXTS:
            media.append(MediaItem(path=path, kind="image", duration=image_duration))
        elif ext in VIDEO_EXTS:
            media.append(MediaItem(path=path, kind="video", duration=ffprobe_duration(path)))
        elif ext in AUDIO_EXTS:
            audio.append(path)

    if not media:
        raise SystemExit(f"No images or videos found in {input_dir}")

    return media, audio, parse_data_file(data_file)


def parse_data_file(path: Path | None) -> MovieData:
    if not path or not path.exists():
        return MovieData(captions={})

    lines = [line.strip() for line in path.read_text(encoding="utf-8", errors="replace").splitlines()]
    lines = [line for line in lines if line and not line.startswith("#")]

    data = MovieData(captions={})
    plain: list[str] = []

    for line in lines:
        if ":" in line:
            key, value = line.split(":", 1)
            key = key.strip()
            value = value.strip()
            low = key.lower()
            if low in {"title", "name"}:
                data.title = value
            elif low in {"date", "dates", "when"}:
                data.date = value
            elif low in {"place", "location", "where"}:
                data.place = value
            elif Path(key).suffix.lower() in IMAGE_EXTS | VIDEO_EXTS:
                data.captions[key] = value
            else:
                plain.append(line)
        else:
            plain.append(line)

    # Simple mode: first three non-key lines are title/date/place.
    if plain:
        data.title = data.title or plain[0]
    if len(plain) > 1:
        data.date = data.date or plain[1]
    if len(plain) > 2:
        data.place = data.place or plain[2]

    return data


def drawtext_escape(text: str) -> str:
    # Escape for ffmpeg drawtext text='...'
    return text.replace("\\", "\\\\").replace(":", "\\:").replace("'", "\\'").replace("%", "\\%")


def normalize_clip(item: MediaItem, out: Path, width: int, height: int, fps: int) -> None:
    vf = (
        f"scale={width}:{height}:force_original_aspect_ratio=decrease,"
        f"pad={width}:{height}:(ow-iw)/2:(oh-ih)/2:black,"
        f"fps={fps},format=yuv420p,setsar=1"
    )

    if item.kind == "image":
        cmd = [
            "ffmpeg", "-y",
            "-loop", "1", "-t", f"{item.duration:.3f}", "-i", str(item.path),
            "-an", "-vf", vf,
            "-c:v", "libx264", "-preset", "veryfast", "-crf", "20",
            str(out),
        ]
    else:
        cmd = [
            "ffmpeg", "-y",
            "-i", str(item.path),
            "-an", "-vf", vf,
            "-c:v", "libx264", "-preset", "veryfast", "-crf", "20",
            str(out),
        ]
    run(cmd)


def concat_or_mix_audio(audio_files: list[Path], out: Path) -> Path | None:
    if not audio_files:
        return None
    if len(audio_files) == 1:
        return audio_files[0]

    list_file = out.with_suffix(".txt")
    def esc(p: Path) -> str:
        return str(p.resolve()).replace("'", "'\\''")
    list_file.write_text("".join(f"file '{esc(p)}'\n" for p in audio_files), encoding="utf-8")
    run([
        "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(list_file),
        "-vn", "-c:a", "aac", "-b:a", "192k", str(out),
    ])
    return out


def build_video(
    clips: list[Path],
    durations: list[float],
    audio: Path | None,
    output: Path,
    data: MovieData,
    fade: float,
    width: int,
    height: int,
    fps: int,
) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)

    cmd = ["ffmpeg", "-y"]
    for clip in clips:
        cmd += ["-i", str(clip)]
    if audio:
        cmd += ["-stream_loop", "-1", "-i", str(audio)]

    filters: list[str] = []
    n = len(clips)

    if n == 1:
        current = "0:v"
    else:
        current = "v01"
        offset = max(0.0, durations[0] - fade)
        filters.append(f"[0:v][1:v]xfade=transition=fade:duration={fade:.3f}:offset={offset:.3f}[{current}]")
        combined_duration = durations[0] + durations[1] - fade
        for i in range(2, n):
            next_label = f"v{i:02d}"
            offset = max(0.0, combined_duration - fade)
            filters.append(f"[{current}][{i}:v]xfade=transition=fade:duration={fade:.3f}:offset={offset:.3f}[{next_label}]")
            current = next_label
            combined_duration += durations[i] - fade

    title_lines = [x for x in [data.title, data.date, data.place] if x]
    final_label = "vout"
    if title_lines:
        text = drawtext_escape("\\n".join(title_lines))
        # Draw centered title for first 5 seconds with a translucent box.
        filters.append(
            f"[{current}]drawtext=text='{text}':x=(w-text_w)/2:y=(h-text_h)/2:"
            f"fontsize={max(32, height // 18)}:fontcolor=white:box=1:boxcolor=black@0.45:boxborderw=24:"
            f"enable='between(t,0,5)'[{final_label}]"
        )
    else:
        filters.append(f"[{current}]copy[{final_label}]")

    cmd += ["-filter_complex", ";".join(filters), "-map", f"[{final_label}]"]

    if audio:
        audio_index = len(clips)
        cmd += ["-map", f"{audio_index}:a:0", "-c:a", "aac", "-b:a", "192k", "-shortest"]
    else:
        cmd += ["-an"]

    cmd += ["-c:v", "libx264", "-preset", "veryfast", "-crf", "20", "-r", str(fps), "-pix_fmt", "yuv420p", str(output)]
    run(cmd)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Make a video from a flat folder of images/videos/audio.")
    parser.add_argument("input", type=Path, help="flat input directory")
    parser.add_argument("--output", "-o", type=Path, default=Path("movie.mp4"), help="output mp4 path")
    parser.add_argument("--image-duration", type=float, default=3.0, help="seconds each image is visible")
    parser.add_argument("--fade", type=float, default=1.0, help="crossfade duration in seconds")
    parser.add_argument("--resolution", type=parse_resolution, default=(1920, 1080), help="output resolution, e.g. 1920x1080")
    parser.add_argument("--fps", type=int, default=30, help="output frames per second")
    args = parser.parse_args(argv)

    require_tool("ffmpeg")
    require_tool("ffprobe")

    width, height = args.resolution
    media, audio_files, data = scan_input(args.input, args.image_duration)
    fade = min(args.fade, max(0.0, args.image_duration - 0.1))

    print(f"Found {len(media)} media file(s), {len(audio_files)} audio file(s).")
    if data.title or data.date or data.place:
        print("Title data:", " | ".join(x for x in [data.title, data.date, data.place] if x))

    with tempfile.TemporaryDirectory(prefix="movmaker-") as tmp_name:
        tmp = Path(tmp_name)
        clips: list[Path] = []
        durations: list[float] = []
        for idx, item in enumerate(media):
            out = tmp / f"clip_{idx:04d}.mp4"
            print(f"Preparing {item.kind}: {item.path.name}")
            normalize_clip(item, out, width, height, args.fps)
            clips.append(out)
            # Use probed normalized duration for accuracy.
            durations.append(ffprobe_duration(out))

        audio = concat_or_mix_audio(audio_files, tmp / "music.m4a")
        build_video(clips, durations, audio, args.output, data, fade, width, height, args.fps)

    print(f"Done: {args.output}")
    return 0


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