#!/usr/bin/env python3
"""Inventory local media files for the VideoReDo comparison matrix."""

from __future__ import annotations

import argparse
import concurrent.futures
import json
import os
from pathlib import Path
import subprocess
import sys
import tempfile
from typing import Any


MEDIA_EXTENSIONS = {
    ".avi", ".m2ts", ".m4v", ".mkv", ".mov", ".mp4", ".mpeg",
    ".mpg", ".mts", ".ts", ".vob", ".webm", ".wmv",
}


def integer(value: Any) -> int:
    try:
        return int(value)
    except (TypeError, ValueError):
        return 0


def number(value: Any) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return 0.0


def bit_depth(stream: dict[str, Any]) -> int:
    explicit = integer(stream.get("bits_per_raw_sample")) or integer(stream.get("bits_per_sample"))
    if explicit:
        return explicit
    pixel_format = str(stream.get("pix_fmt", "")).lower()
    for depth in (16, 14, 12, 10, 9):
        if str(depth) in pixel_format:
            return depth
    return 8 if pixel_format else 0


def probe_file(ffprobe: str, path: Path, timeout: int) -> dict[str, Any]:
    command = [
        ffprobe, "-v", "error", "-show_entries",
        "format=format_name,duration,size,bit_rate:"
        "stream=index,codec_type,codec_name,profile,pix_fmt,width,height,"
        "bits_per_raw_sample,bits_per_sample,channels,channel_layout,sample_rate,"
        "avg_frame_rate,r_frame_rate,nb_frames:stream_disposition=attached_pic",
        "-of", "json", str(path),
    ]
    try:
        completed = subprocess.run(
            command, capture_output=True, text=True, timeout=timeout,
            encoding="utf-8", errors="replace", check=False,
        )
    except subprocess.TimeoutExpired:
        return {"path": str(path), "status": "timeout"}
    if completed.returncode != 0:
        return {
            "path": str(path), "status": "probe-failed",
            "error": completed.stderr.strip()[-1000:],
        }
    try:
        document = json.loads(completed.stdout)
    except json.JSONDecodeError as error:
        return {"path": str(path), "status": "invalid-json", "error": str(error)}
    streams = document.get("streams", [])
    videos = [
        stream for stream in streams
        if stream.get("codec_type") == "video"
        and not integer(stream.get("disposition", {}).get("attached_pic"))
    ]
    audios = [stream for stream in streams if stream.get("codec_type") == "audio"]
    if not videos:
        return {"path": str(path), "status": "no-video"}
    video = videos[0]
    audio_rows = [
        {
            "codec": str(stream.get("codec_name", "unknown")),
            "profile": str(stream.get("profile", "")),
            "channels": integer(stream.get("channels")),
            "layout": str(stream.get("channel_layout", "")),
            "sample_rate": integer(stream.get("sample_rate")),
        }
        for stream in audios
    ]
    media_format = document.get("format", {})
    return {
        "path": str(path),
        "status": "ok",
        "extension": path.suffix.lower(),
        "size": integer(media_format.get("size")) or path.stat().st_size,
        "duration": number(media_format.get("duration")),
        "format": str(media_format.get("format_name", "")),
        "video": {
            "codec": str(video.get("codec_name", "unknown")),
            "profile": str(video.get("profile", "")),
            "pixel_format": str(video.get("pix_fmt", "")),
            "bit_depth": bit_depth(video),
            "width": integer(video.get("width")),
            "height": integer(video.get("height")),
            "average_frame_rate": str(video.get("avg_frame_rate", "")),
            "real_frame_rate": str(video.get("r_frame_rate", "")),
            "declared_frames": integer(video.get("nb_frames")),
        },
        "audio": audio_rows,
        "stream_count": len(streams),
    }


def media_files(roots: list[Path]) -> list[Path]:
    result: list[Path] = []
    for root in roots:
        if not root.exists():
            continue
        for directory, _, names in os.walk(root):
            for name in names:
                path = Path(directory, name)
                if path.suffix.lower() in MEDIA_EXTENSIONS:
                    result.append(path)
    return sorted(result, key=lambda item: str(item).lower())


def write_atomic(path: Path, document: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        "w", encoding="utf-8", newline="\n", delete=False, dir=path.parent,
    ) as temporary:
        json.dump(document, temporary, indent=2, ensure_ascii=False)
        temporary.write("\n")
        temporary_path = Path(temporary.name)
    temporary_path.replace(path)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", action="append", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--ffprobe", default=r"C:\ffmpeg_latest\ffprobe.exe")
    parser.add_argument("--workers", type=int, default=8)
    parser.add_argument("--timeout", type=int, default=30)
    parser.add_argument("--max-files", type=int, default=0)
    arguments = parser.parse_args()
    files = media_files(arguments.root)
    if arguments.max_files > 0:
        files = files[: arguments.max_files]
    rows: list[dict[str, Any]] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=max(1, arguments.workers)) as executor:
        futures = {
            executor.submit(probe_file, arguments.ffprobe, path, arguments.timeout): path
            for path in files
        }
        for completed, future in enumerate(concurrent.futures.as_completed(futures), 1):
            rows.append(future.result())
            if completed % 100 == 0 or completed == len(futures):
                print(f"probed {completed}/{len(futures)}", flush=True)
    rows.sort(key=lambda row: row["path"].lower())
    combinations: dict[str, int] = {}
    for row in rows:
        if row.get("status") != "ok":
            continue
        audio = "+".join(sorted({track["codec"] for track in row["audio"]})) or "none"
        video = row["video"]
        key = f"{row['extension']}|{video['codec']}|{video['bit_depth']}bit|{audio}"
        combinations[key] = combinations.get(key, 0) + 1
    document = {
        "schema": "smartcut-media-inventory-v1",
        "roots": [str(root) for root in arguments.root],
        "files": len(rows),
        "successful": sum(row.get("status") == "ok" for row in rows),
        "combinations": dict(sorted(combinations.items())),
        "items": rows,
    }
    write_atomic(arguments.output, document)
    print(f"wrote {arguments.output} with {len(combinations)} combinations")
    return 0


if __name__ == "__main__":
    sys.exit(main())
