#!/usr/bin/env python3

"""
Extract diagnostic frames from a video for future WCX title detection.

This script never renames, moves, deletes, or modifies the input video.

Example:
    detect_wcx_title.py /path/to/video.mp4 \
        --output-dir /tmp/wcx-title-frames

    detect_wcx_title.py /path/to/video.mp4 \
        --output-dir /tmp/wcx-diagnostic \
        --diagnose-layout
"""

import argparse
import json
import os
import shlex
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Callable, Sequence


FRAME_TIMESTAMPS = tuple(range(6, 16))
DIAGNOSTIC_TIMESTAMPS = (8, 10, 12)
NORMALIZED_WIDTH = 640
NORMALIZED_HEIGHT = 360
DEFAULT_ENDINGS = ("mp4", "avi", "mkv", "mov", "m4v", "webm")
IGNORED_FILENAMES = {".ds_store", "thumbs.db", "desktop.ini"}


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Extract diagnostic PNG frames at 6 through 15 seconds "
            "without modifying the input video."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""Examples:
  scripts/detect_wcx_title.py /path/to/video.mp4 \
    --output-dir /tmp/wcx-diagnostic --diagnose-layout

  scripts/detect_wcx_title.py /storage/disk1/X --batch

  scripts/detect_wcx_title.py /storage/disk1/X \
    --batch --recursive --ending mp4,avi""",
    )

    parser.add_argument(
        "input_path",
        type=Path,
        help="Video file, or a directory when --batch is used.",
    )

    parser.add_argument(
        "--output-dir",
        type=Path,
        help="Directory in which the extracted PNG frames will be written.",
    )

    parser.add_argument(
        "--diagnose-layout",
        action="store_true",
        help=(
            "Extract normalized frames at 8.0, 10.0, and 12.0 seconds "
            "and report layout metrics."
        ),
    )

    parser.add_argument(
        "--batch",
        action="store_true",
        help="Classify video files in a directory without writing PNG files.",
    )

    parser.add_argument(
        "--recursive",
        action="store_true",
        help="Search subdirectories in batch mode.",
    )

    parser.add_argument(
        "--ending",
        action="append",
        help=(
            "Video extension for batch mode; may be comma-separated or "
            "repeated."
        ),
    )

    parser.add_argument(
        "--json-report",
        type=Path,
        help="Write a machine-readable JSON report after a batch run.",
    )

    return parser.parse_args()


def prepare_output_directory(output_dir: Path) -> None:
    try:
        output_dir.mkdir(parents=True, exist_ok=True)
    except OSError as error:
        raise RuntimeError(
            f"Could not create output directory {output_dir}: {error}"
        ) from error

    if not output_dir.is_dir():
        raise RuntimeError(
            f"Output path is not a directory: {output_dir}"
        )


def parse_endings(values: Sequence[str] | None) -> set[str]:
    if values is None:
        return set(DEFAULT_ENDINGS)

    endings = {
        ending.strip().lower().removeprefix(".")
        for value in values
        for ending in value.split(",")
        if ending.strip()
    }

    if not endings:
        raise ValueError("--ending must contain at least one extension")

    return endings


def is_ignored_file(path: Path) -> bool:
    return (
        path.name.startswith("._")
        or path.name.lower() in IGNORED_FILENAMES
    )


def is_video_file(path: Path, endings: set[str]) -> bool:
    return (
        path.is_file()
        and not is_ignored_file(path)
        and path.suffix.lower().removeprefix(".") in endings
    )


def find_video_files(
    directory: Path,
    endings: set[str],
    recursive: bool,
) -> list[Path]:
    directory = directory.resolve()
    candidates = directory.rglob("*") if recursive else directory.iterdir()
    return sorted(
        (
            path
            for path in candidates
            if is_video_file(path, endings)
        ),
        key=lambda path: str(path),
    )


def extract_frame(
    ffmpeg: str,
    video_file: Path,
    output_file: Path,
    timestamp: int,
) -> None:
    command = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "error",
        "-ss",
        f"{timestamp:.1f}",
        "-i",
        str(video_file),
        "-frames:v",
        "1",
        "-n",
        str(output_file),
    ]

    print(f"Running: {shlex.join(command)}")

    try:
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
        )
    except OSError as error:
        raise RuntimeError(
            f"Could not run ffmpeg for {timestamp:.1f} seconds: {error}"
        ) from error

    if result.returncode != 0:
        detail = result.stderr.strip() or "ffmpeg returned no error message"
        raise RuntimeError(
            f"ffmpeg could not extract the frame at "
            f"{timestamp:.1f} seconds: {detail}"
        )

    if not output_file.is_file():
        raise RuntimeError(
            f"ffmpeg reported success but did not create the frame at "
            f"{timestamp:.1f} seconds: {output_file}"
        )

    print(f"Created: {output_file}")


def is_nearly_black(red: int, green: int, blue: int) -> bool:
    return red <= 30 and green <= 30 and blue <= 30


def is_red(red: int, green: int, blue: int) -> bool:
    return (
        red >= 140
        and red >= green * 1.5
        and red >= blue * 1.5
    )


def is_bright(red: int, green: int, blue: int) -> bool:
    return red >= 180 and green >= 180 and blue >= 180


def calculate_region_ratio(
    rgb_data: bytes,
    width: int,
    x_start: int,
    x_end: int,
    y_start: int,
    y_end: int,
    predicate,
) -> float:
    matching_pixels = 0
    total_pixels = (x_end - x_start) * (y_end - y_start)

    for y_position in range(y_start, y_end):
        row_offset = y_position * width * 3
        for x_position in range(x_start, x_end):
            offset = row_offset + x_position * 3
            if predicate(
                rgb_data[offset],
                rgb_data[offset + 1],
                rgb_data[offset + 2],
            ):
                matching_pixels += 1

    return matching_pixels / total_pixels


def calculate_layout_metrics(
    rgb_data: bytes,
    width: int = NORMALIZED_WIDTH,
    height: int = NORMALIZED_HEIGHT,
) -> dict[str, float]:
    expected_bytes = width * height * 3
    if len(rgb_data) != expected_bytes:
        raise ValueError(
            f"Expected {expected_bytes} RGB bytes for {width}x{height}, "
            f"received {len(rgb_data)}"
        )

    edge_width = max(1, int(width * 0.10))
    edge_height = max(1, int(height * 0.10))
    band_height = max(1, int(height * 0.15))

    return {
        "black_ratio_total": calculate_region_ratio(
            rgb_data, width, 0, width, 0, height, is_nearly_black
        ),
        "black_ratio_top": calculate_region_ratio(
            rgb_data, width, 0, width, 0, edge_height, is_nearly_black
        ),
        "black_ratio_bottom": calculate_region_ratio(
            rgb_data,
            width,
            0,
            width,
            height - edge_height,
            height,
            is_nearly_black,
        ),
        "black_ratio_left": calculate_region_ratio(
            rgb_data, width, 0, edge_width, 0, height, is_nearly_black
        ),
        "black_ratio_right": calculate_region_ratio(
            rgb_data,
            width,
            width - edge_width,
            width,
            0,
            height,
            is_nearly_black,
        ),
        "red_ratio_top": calculate_region_ratio(
            rgb_data, width, 0, width, 0, band_height, is_red
        ),
        "bright_ratio_top": calculate_region_ratio(
            rgb_data, width, 0, width, 0, band_height, is_bright
        ),
        "bright_ratio_bottom": calculate_region_ratio(
            rgb_data,
            width,
            0,
            width,
            height - band_height,
            height,
            is_bright,
        ),
    }


def calculate_wcx_layout_score(metrics: dict[str, float]) -> int:
    conditions = (
        0.20 <= metrics["black_ratio_total"] <= 0.60,
        metrics["black_ratio_top"] >= 0.45,
        metrics["black_ratio_bottom"] >= 0.50,
        metrics["black_ratio_left"] >= 0.75,
        metrics["black_ratio_right"] >= 0.75,
        metrics["red_ratio_top"] >= 0.05,
        metrics["bright_ratio_top"] >= 0.02,
        metrics["bright_ratio_bottom"] >= 0.03,
    )
    return sum(conditions)


def select_best_timestamp(
    timestamp_scores: Sequence[tuple[float, int]],
) -> tuple[float, int]:
    if not timestamp_scores:
        raise ValueError("At least one timestamp and score is required")

    return min(
        timestamp_scores,
        key=lambda item: (-item[1], item[0]),
    )


def classify_wcx_score(score: int) -> str:
    if not 0 <= score <= 8:
        raise ValueError(f"WCX layout score must be between 0 and 8: {score}")

    if score == 8:
        return "wcx"

    if score >= 6:
        return "uncertain"

    return "not_wcx"


def split_rgb24_frames(
    rgb_data: bytes,
    width: int = NORMALIZED_WIDTH,
    height: int = NORMALIZED_HEIGHT,
    frame_count: int = len(DIAGNOSTIC_TIMESTAMPS),
) -> tuple[bytes, ...]:
    frame_size = width * height * 3
    expected_bytes = frame_size * frame_count
    actual_bytes = len(rgb_data)

    if actual_bytes < expected_bytes:
        raise ValueError(
            f"RGB24 buffer is too short: expected {expected_bytes} bytes, "
            f"received {actual_bytes}"
        )

    if actual_bytes > expected_bytes:
        raise ValueError(
            f"RGB24 buffer is too long: expected {expected_bytes} bytes, "
            f"received {actual_bytes}"
        )

    return tuple(
        rgb_data[offset:offset + frame_size]
        for offset in range(0, expected_bytes, frame_size)
    )


def analyze_rgb24_frames(
    rgb_data: bytes,
) -> tuple[
    tuple[tuple[float, dict[str, float], int], ...],
    float,
    int,
    str,
]:
    frames = split_rgb24_frames(rgb_data)
    analyses = tuple(
        (
            float(timestamp),
            metrics,
            calculate_wcx_layout_score(metrics),
        )
        for timestamp, frame in zip(
            DIAGNOSTIC_TIMESTAMPS,
            frames,
        )
        for metrics in (calculate_layout_metrics(frame),)
    )
    best_timestamp, best_score = select_best_timestamp(
        tuple(
            (timestamp, score)
            for timestamp, _metrics, score in analyses
        )
    )
    return (
        analyses,
        best_timestamp,
        best_score,
        classify_wcx_score(best_score),
    )


def extract_batch_rgb24(
    ffmpeg: str,
    video_file: Path,
) -> bytes:
    filter_graph = (
        "[0:v]setpts=PTS-STARTPTS,"
        "fps=1/2:start_time=0,"
        "scale=640:360:force_original_aspect_ratio=decrease,"
        "pad=640:360:(ow-iw)/2:(oh-ih)/2:black"
    )
    command = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "error",
        "-ss",
        f"{DIAGNOSTIC_TIMESTAMPS[0]:.1f}",
        "-t",
        "5.0",
        "-i",
        str(video_file),
        "-filter_complex",
        filter_graph,
        "-frames:v",
        str(len(DIAGNOSTIC_TIMESTAMPS)),
        "-c:v",
        "rawvideo",
        "-pix_fmt",
        "rgb24",
        "-f",
        "rawvideo",
        "pipe:1",
    ]

    try:
        result = subprocess.run(
            command,
            capture_output=True,
        )
    except OSError as error:
        raise RuntimeError(f"Could not run ffmpeg: {error}") from error

    if result.returncode != 0:
        detail = (
            result.stderr.decode("utf-8", errors="replace").strip()
            or "ffmpeg returned no error message"
        )
        raise RuntimeError(f"ffmpeg could not analyze the video: {detail}")

    try:
        split_rgb24_frames(result.stdout)
    except ValueError as error:
        raise RuntimeError(f"Invalid RGB24 output from ffmpeg: {error}") from error

    return result.stdout


def analyze_video_for_batch(
    ffmpeg: str,
    video_file: Path,
) -> tuple[float, int, str]:
    _analyses, best_timestamp, best_score, classification = (
        analyze_rgb24_frames(
            extract_batch_rgb24(ffmpeg, video_file)
        )
    )
    return best_timestamp, best_score, classification


def process_batch(
    video_files: Sequence[Path],
    analyzer: Callable[[Path], tuple[float, int, str]],
) -> tuple[list[dict[str, object]], dict[str, int]]:
    summary = {
        "processed": 0,
        "wcx": 0,
        "uncertain": 0,
        "not_wcx": 0,
        "errors": 0,
    }
    files = []

    for video_file in video_files:
        absolute_path = str(video_file.resolve())
        summary["processed"] += 1
        try:
            timestamp, score, classification = analyzer(video_file)
        except (OSError, RuntimeError, ValueError) as error:
            summary["errors"] += 1
            files.append(
                {
                    "path": absolute_path,
                    "classification": "error",
                    "score": None,
                    "best_timestamp": None,
                    "error": str(error),
                }
            )
            print(f"error       -     -   {video_file}: {error}")
            continue

        summary[classification] += 1
        files.append(
            {
                "path": absolute_path,
                "classification": classification,
                "score": score,
                "best_timestamp": timestamp,
            }
        )
        print(
            f"{classification:<10} {score}/8  "
            f"{timestamp:>4.1f}  {video_file}"
        )

    print()
    print(f"processed: {summary['processed']}")
    print(f"wcx: {summary['wcx']}")
    print(f"uncertain: {summary['uncertain']}")
    print(f"not_wcx: {summary['not_wcx']}")
    print(f"errors: {summary['errors']}")

    return files, summary


def build_json_report(
    source_directory: Path,
    recursive: bool,
    endings: set[str],
    summary: dict[str, int],
    files: list[dict[str, object]],
) -> dict[str, object]:
    return {
        "schema_version": 1,
        "source_directory": str(source_directory.resolve()),
        "recursive": recursive,
        "endings": sorted(endings),
        "summary": dict(summary),
        "files": list(files),
    }


def write_json_report(
    report_path: Path,
    report: dict[str, object],
) -> None:
    if report_path.exists():
        raise FileExistsError(
            f"JSON report already exists: {report_path}"
        )

    report_path.parent.mkdir(parents=True, exist_ok=True)
    temporary_path = None

    try:
        with tempfile.NamedTemporaryFile(
            mode="w",
            encoding="utf-8",
            dir=report_path.parent,
            prefix=f".{report_path.name}.",
            suffix=".tmp",
            delete=False,
        ) as temporary_file:
            temporary_path = Path(temporary_file.name)
            json.dump(
                report,
                temporary_file,
                indent=2,
                ensure_ascii=False,
            )
            temporary_file.write("\n")
            temporary_file.flush()
            os.fsync(temporary_file.fileno())

        if report_path.exists():
            raise FileExistsError(
                f"JSON report already exists: {report_path}"
            )

        temporary_path.replace(report_path)
        temporary_path = None
    finally:
        if temporary_path is not None:
            temporary_path.unlink(missing_ok=True)


def write_optional_json_report(
    report_path: Path | None,
    report: dict[str, object],
    writer: Callable[
        [Path, dict[str, object]],
        None,
    ] = write_json_report,
) -> int:
    if report_path is None:
        return 0

    try:
        writer(report_path, report)
    except (OSError, TypeError, ValueError) as error:
        print(
            f"Error: Could not write JSON report {report_path}: {error}",
            file=sys.stderr,
        )
        return 1

    return 0


def diagnose_layout(
    ffmpeg: str,
    video_file: Path,
    output_dir: Path,
) -> None:
    normalized_files = tuple(
        output_dir / f"normalized-{timestamp:04.1f}.png"
        for timestamp in DIAGNOSTIC_TIMESTAMPS
    )

    for output_file in normalized_files:
        if output_file.exists():
            raise RuntimeError(
                f"Output file already exists and will not be overwritten: "
                f"{output_file}"
            )

    filter_graph = (
        "[0:v]setpts=PTS-STARTPTS,"
        "fps=1/2:start_time=0,"
        "scale=640:360:force_original_aspect_ratio=decrease,"
        "pad=640:360:(ow-iw)/2:(oh-ih)/2:black,"
        "split=4[normalized_raw][frame_8_input]"
        "[frame_10_input][frame_12_input];"
        "[frame_8_input]select=eq(n\\,0)[frame_8];"
        "[frame_10_input]select=eq(n\\,1)[frame_10];"
        "[frame_12_input]select=eq(n\\,2)[frame_12]"
    )

    command = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "error",
        "-n",
        "-ss",
        f"{DIAGNOSTIC_TIMESTAMPS[0]:.1f}",
        "-t",
        "5.0",
        "-i",
        str(video_file),
        "-filter_complex",
        filter_graph,
        "-map",
        "[frame_8]",
        "-frames:v",
        "1",
        "-c:v",
        "png",
        str(normalized_files[0]),
        "-map",
        "[frame_10]",
        "-frames:v",
        "1",
        "-c:v",
        "png",
        str(normalized_files[1]),
        "-map",
        "[frame_12]",
        "-frames:v",
        "1",
        "-c:v",
        "png",
        str(normalized_files[2]),
        "-map",
        "[normalized_raw]",
        "-frames:v",
        str(len(DIAGNOSTIC_TIMESTAMPS)),
        "-c:v",
        "rawvideo",
        "-pix_fmt",
        "rgb24",
        "-f",
        "rawvideo",
        "pipe:1",
    ]

    print(f"Running: {shlex.join(command)}")

    try:
        result = subprocess.run(
            command,
            capture_output=True,
        )
    except OSError as error:
        raise RuntimeError(
            f"Could not run ffmpeg for diagnostic frames: {error}"
        ) from error

    if result.returncode != 0:
        detail = (
            result.stderr.decode("utf-8", errors="replace").strip()
            or "ffmpeg returned no error message"
        )
        raise RuntimeError(
            f"ffmpeg could not extract the diagnostic frames: {detail}"
        )

    for output_file in normalized_files:
        if not output_file.is_file():
            raise RuntimeError(
                f"ffmpeg reported success but did not create: {output_file}"
            )
        print(f"Created: {output_file}")

    try:
        analyses, best_timestamp, best_score, classification = (
            analyze_rgb24_frames(result.stdout)
        )
    except ValueError as error:
        raise RuntimeError(f"Invalid RGB24 output from ffmpeg: {error}") from error

    for index, (timestamp, metrics, score) in enumerate(analyses):
        if index:
            print()
        print(f"timestamp: {timestamp:.1f}")
        print(f"wcx_layout_score: {score}/8")
        print(f"normalized_size: {NORMALIZED_WIDTH}x{NORMALIZED_HEIGHT}")
        for name, ratio in metrics.items():
            print(f"{name}: {ratio:.6f}")

    print()
    print(f"best_timestamp: {best_timestamp:.1f}")
    print(f"best_wcx_layout_score: {best_score}/8")
    print(f"classification: {classification}")


def main() -> int:
    args = parse_arguments()

    if args.batch:
        if args.diagnose_layout:
            print(
                "Error: --batch cannot be combined with --diagnose-layout.",
                file=sys.stderr,
            )
            return 2

        if args.output_dir is not None:
            print(
                "Error: --batch cannot be combined with --output-dir.",
                file=sys.stderr,
            )
            return 2

        if (
            args.json_report is not None
            and args.json_report.exists()
        ):
            print(
                f"Error: JSON report already exists: {args.json_report}",
                file=sys.stderr,
            )
            return 2

        if not args.input_path.is_dir():
            print(
                f"Error: Batch directory does not exist: {args.input_path}",
                file=sys.stderr,
            )
            return 2

        try:
            endings = parse_endings(args.ending)
            video_files = find_video_files(
                args.input_path,
                endings=endings,
                recursive=args.recursive,
            )
        except (OSError, ValueError) as error:
            print(f"Error: {error}", file=sys.stderr)
            return 2

        if not video_files:
            print(
                f"No matching video files found in: "
                f"{args.input_path.resolve()}"
            )
            analyzer = lambda _path: (0.0, 0, "not_wcx")
        else:
            ffmpeg = shutil.which("ffmpeg")
            if ffmpeg is None:
                print("Error: ffmpeg was not found in PATH.", file=sys.stderr)
                return 1

            analyzer = lambda video_file: analyze_video_for_batch(
                ffmpeg,
                video_file,
            )

        files, summary = process_batch(
            video_files,
            analyzer,
        )
        report = build_json_report(
            source_directory=args.input_path,
            recursive=args.recursive,
            endings=endings,
            summary=summary,
            files=files,
        )
        report_exit_code = write_optional_json_report(
            args.json_report,
            report,
        )
        return 1 if summary["errors"] or report_exit_code else 0

    if args.json_report is not None:
        print(
            "Error: --json-report requires --batch.",
            file=sys.stderr,
        )
        return 2

    if args.recursive:
        print(
            "Error: --recursive requires --batch.",
            file=sys.stderr,
        )
        return 2

    if args.ending is not None:
        print(
            "Error: --ending requires --batch.",
            file=sys.stderr,
        )
        return 2

    if args.output_dir is None:
        print(
            "Error: --output-dir is required unless --batch is used.",
            file=sys.stderr,
        )
        return 2

    if not args.input_path.is_file():
        print(
            f"Error: Video file does not exist: {args.input_path}",
            file=sys.stderr,
        )
        return 1

    ffmpeg = shutil.which("ffmpeg")
    if ffmpeg is None:
        print(
            "Error: ffmpeg was not found in PATH.",
            file=sys.stderr,
        )
        return 1

    try:
        prepare_output_directory(args.output_dir)

        if args.diagnose_layout:
            diagnose_layout(
                ffmpeg=ffmpeg,
                video_file=args.input_path,
                output_dir=args.output_dir,
            )
        else:
            for timestamp in FRAME_TIMESTAMPS:
                output_file = (
                    args.output_dir
                    / f"frame-{timestamp:04.1f}.png"
                )
                extract_frame(
                    ffmpeg=ffmpeg,
                    video_file=args.input_path,
                    output_file=output_file,
                    timestamp=timestamp,
                )
    except RuntimeError as error:
        print(f"Error: {error}", file=sys.stderr)
        return 1

    return 0


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