# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import subprocess
from pathlib import Path

import numpy as np

import torch
from PIL import Image

from .utils import (
    AV1_VIDEO,
    H265_VIDEO,
    NASA_VIDEO,
    NASA_VIDEO_HDR,
    TEST_SRC_2_12BIT_HDR,
    TEST_SRC_2_720P_HDR,
    TestVideo,
)

# Run this script to update the resources used in unit tests. The resources are all derived
# from source media already checked into the repo.


def convert_image_to_tensor(image_path: str) -> None:
    image_path = Path(image_path)
    if not image_path.exists():
        return
    # Get base filename without extension
    base_filename = image_path.with_suffix("")
    pil_image = Image.open(image_path)
    img_tensor = torch.from_numpy(np.asarray(pil_image))
    # Save tensor to disk
    torch.save(
        img_tensor, str(base_filename) + ".pt", _use_new_zipfile_serialization=True
    )
    image_path.unlink()


def generate_frame_by_index(
    video: TestVideo,
    *,
    frame_index: int,
    stream_index: int,
    filters: str | None = None,
) -> None:
    # Note that we are using 0-based index naming. As a result, we are
    # generating files one-by-one, giving the actual file name that we want.
    # ffmpeg does have an option to generate multiple files for us, but it uses
    # 1-based indexing. We can't use 1-based indexing because we want to match
    # the 0-based indexing in our tests.
    base_path = video.get_base_path_by_index(
        frame_index, stream_index=stream_index, filters=filters
    )
    output_bmp = f"{base_path}.bmp"

    # Note that we have an exlicit format conversion to rgb24 in our filtergraph
    # specification, and we always place the user-supplied filters AFTER the
    # format conversion. We do this to ensure that the filters are applied in
    # RGB24 colorspace, which matches TorchCodec's behavior.
    select = f"select='eq(n\\,{frame_index})'"
    format = "format=rgb24"
    if filters is not None:
        filtergraph = ",".join([select, format, filters])
    else:
        filtergraph = ",".join([select, format])

    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        video.path,
        "-map",
        f"0:{stream_index}",
        "-vf",
        filtergraph,
        "-fps_mode",
        "passthrough",
        "-update",
        "1",
        output_bmp,
    ]
    subprocess.run(cmd, check=True)
    convert_image_to_tensor(output_bmp)


def generate_frame_by_index_rgb48(
    video: TestVideo,
    *,
    frame_index: int,
    stream_index: int,
) -> None:
    base_path = video.get_base_path_by_index(frame_index, stream_index=stream_index)
    output_raw = f"{base_path}.rgb48.raw"

    select = f"select='eq(n\\,{frame_index})'"
    filtergraph = ",".join([select, "format=rgb48"])

    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        video.path,
        "-map",
        f"0:{stream_index}",
        "-vf",
        filtergraph,
        "-fps_mode",
        "passthrough",
        "-f",
        "rawvideo",
        "-pix_fmt",
        "rgb48",
        "-frames:v",
        "1",
        output_raw,
    ]
    subprocess.run(cmd, check=True)

    height = video.get_height(stream_index=stream_index)
    width = video.get_width(stream_index=stream_index)
    data = np.fromfile(output_raw, dtype=np.uint16).reshape(height, width, 3)
    tensor = torch.from_numpy(data.copy())
    torch.save(tensor, f"{base_path}.rgb48.pt", _use_new_zipfile_serialization=True)
    Path(output_raw).unlink()


def generate_frame_by_timestamp(
    video_path: str, timestamp: float, output_path: str
) -> None:
    cmd = [
        "ffmpeg",
        "-y",
        "-ss",
        str(timestamp),
        "-i",
        video_path,
        "-frames:v",
        "1",
        output_path,
    ]
    subprocess.run(cmd, check=True)
    convert_image_to_tensor(output_path)


def generate_nasa_13013_references_by_index():
    # Note: The naming scheme used here must match the naming scheme used to load
    # tensors in ./utils.py.
    streams = [0, 3]
    frames = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 30, 35, 386, 387, 388, 389]
    for stream in streams:
        for frame in frames:
            generate_frame_by_index(NASA_VIDEO, frame_index=frame, stream_index=stream)


def generate_nasa_13013_references_by_timestamp():
    # Extract individual frames at specific timestamps, including the last frame of the video.
    seek_timestamp = [6.0, 6.1, 10.0, 12.979633]
    timestamp_name = [f"{seek_timestamp:06f}" for seek_timestamp in seek_timestamp]
    for timestamp, name in zip(seek_timestamp, timestamp_name):
        output_bmp = f"{NASA_VIDEO.path}.time{name}.bmp"
        generate_frame_by_timestamp(NASA_VIDEO.path, timestamp, output_bmp)


def generate_nasa_13013_references_crop():
    # Extract frames with specific filters. We have tests that assume these exact filters.
    frames = [0, 15, 200, 389]
    crop_filter = "crop=300:200:50:35:exact=1"
    for frame in frames:
        generate_frame_by_index(
            NASA_VIDEO, frame_index=frame, stream_index=3, filters=crop_filter
        )


def generate_nasa_13013_references_resize():
    frames = [17, 230, 389]
    # Note that the resize algorithm passed to flags is exposed to users,
    # but bilinear is the default we use.
    resize_filter = "scale=240:135:flags=bilinear"
    for frame in frames:
        generate_frame_by_index(
            NASA_VIDEO, frame_index=frame, stream_index=3, filters=resize_filter
        )


def generate_nasa_13013_references():
    generate_nasa_13013_references_by_index()
    generate_nasa_13013_references_by_timestamp()
    generate_nasa_13013_references_crop()
    generate_nasa_13013_references_resize()


def generate_h265_video_references():
    # This video was generated by running the following:
    # conda install -c conda-forge x265
    # ./configure --enable-nonfree --enable-gpl --prefix=$(readlink -f ../bin) --enable-libx265  --enable-rpath --extra-ldflags=-Wl,-rpath=$CONDA_PREFIX/lib --enable-filter=drawtext --enable-libfontconfig --enable-libfreetype --enable-libharfbuzz
    # ffmpeg -f lavfi -i color=size=128x128:duration=1:rate=10:color=blue -vf "drawtext=fontsize=30:fontcolor=white:x=(w-text_w)/2:y=(h-text_h)/2:text='Frame %{frame_num}'" -vcodec libx265 -pix_fmt yuv420p -g 2 -crf 10 h265_video.mp4 -y
    # Note that this video only has 1 stream, at index 0.
    frames = [5]
    for frame in frames:
        generate_frame_by_index(H265_VIDEO, frame_index=frame, stream_index=0)


def generate_av1_video_references():
    # This video was generated by running the following:
    # ffmpeg -f lavfi -i testsrc=duration=5:size=640x360:rate=25,format=yuv420p -c:v libaom-av1 -crf 30 -colorspace bt709 -color_primaries bt709 -color_trc bt709 av1_video.mkv
    # Note that this video only has 1 stream, at index 0.
    frames = [10]
    for frame in frames:
        generate_frame_by_index(AV1_VIDEO, frame_index=frame, stream_index=0)


def generate_hdr_references_rgb48():
    frames = [0, 5, 10]
    for video in (
        NASA_VIDEO,
        NASA_VIDEO_HDR,
        TEST_SRC_2_720P_HDR,
        TEST_SRC_2_12BIT_HDR,
    ):
        for frame in frames:
            generate_frame_by_index_rgb48(
                video,
                frame_index=frame,
                stream_index=video.default_stream_index,
            )


def generate_hdr_references_rgb24():
    frames = [0, 5, 10]
    for video in (NASA_VIDEO_HDR, TEST_SRC_2_720P_HDR, TEST_SRC_2_12BIT_HDR):
        for frame in frames:
            generate_frame_by_index(video, frame_index=frame, stream_index=0)


def main():
    generate_nasa_13013_references()
    generate_h265_video_references()
    generate_av1_video_references()
    generate_hdr_references_rgb48()
    generate_hdr_references_rgb24()


if __name__ == "__main__":
    main()
