#!/usr/bin/env python3
from argparse import ArgumentParser, Namespace
from pathlib import Path
from natsort import natsorted
import numpy as np
from vre import VRE, FFmpegVideo, FakeVideo
from vre.utils import image_read, abs_path, load_function_from_module
from vre.representations import build_representations_from_cfg, add_external_repositories
from vre.logger import vre_logger as logger
from vre_repository import get_vre_repository

def _get_raw_data(video_path: Path) -> dict[str, np.ndarray | list[str]]:
    all_files = natsorted([f for f in video_path.iterdir() if f.is_file()], key=lambda p: p.name)
    suffixes = set(p.suffix for p in all_files)
    assert len(suffixes) == 1, f"Expected a single type of files in '{video_path}' found {suffixes}"
    frames = [int(x.stem) for x in all_files]
    fn = {".png": image_read, ".npz": lambda p: np.load(p)["arr_0"], ".npy": np.load}[next(iter(suffixes))]
    raw_data = [fn(f) for f in all_files]
    assert all(x.shape == raw_data[0].shape for x in raw_data), f"Images shape differ in '{video_path}'"
    logger.info(f"--video_path is a directory. Assuming a directory of images. Found {len(raw_data)} images.")
    return {"data": np.array(raw_data, dtype=np.uint8), "frames": frames}

def get_args() -> Namespace:
    """cli args"""
    parser = ArgumentParser()
    parser.add_argument("video_path", help="Path to the scene video we are processing", type=abs_path)
    parser.add_argument("--output_path", "-o", required=True, type=abs_path,
                        help="Path to the output directory where representations are stored")
    parser.add_argument("--config_path", required=True, help="Path to global YAML cfg file", type=abs_path)
    parser.add_argument("--representations", nargs="+", help="Subset of representations to run VRE for")
    # run parameters
    parser.add_argument("--frames", nargs="+", help="The list of frames to be exported from the video or a l..r range")
    parser.add_argument("--output_dir_exists_mode", choices=["overwrite", "skip_computed", "raise"], default="raise")
    parser.add_argument("--exception_mode", type=str, choices=["skip_representation", "stop_execution"],
                        default="stop_execution")
    parser.add_argument("--n_threads_data_storer", type=int, default=0)
    parser.add_argument("--external_representations", "-I", nargs="+", default=[],
                        help="Path to external reprs. Format: /path/to/file.py:fn_name. fn -> [Representation]")
    parser.add_argument("--external_repositories", "-J", nargs="+", default=[],
                        help="Path to external reprs. Format: /path/to/file.py:fn_name. fn -> {str: Type[Repr]}")
    # collage parameters (calls vre_collage)
    parser.add_argument("--collage", action="store_true", help="If set, calls vre_collage on the output")
    parser.add_argument("--collage_video", action="store_true", help="if set, calls vre_collage on the output")
    parser.add_argument("--collage_fps", type=float, help="Required by --make_video_collage for the collage output")
    parser.add_argument("--collage_output_path", type=Path, help="If not set, uses 'args.output_path/.collage'")

    args = parser.parse_args()
    if args.collage is True:
        if args.collage_video:
            assert args.collage_fps is not None and args.collage_fps > 0, args.collage_fps
        if args.collage_output_path is None:
            args.collage_output_path = args.output_path / ".collage"
            logger.warning(f"--collage_output_path not set, defaulting to: '{args.collage_output_path}'")
    if args.frames is not None:
        assert len(args.frames) > 0, f"No frames provided, cannot run VRE"
        if len(args.frames) == 1 and args.frames[0].find("..") != -1:
            args.frames = list(range(*map(int, args.frames[0].split(".."))))
        args.frames = list(map(int, args.frames))

    return args

def main(args: Namespace):
    """Za main function"""

    if args.video_path.is_dir():
        video = FakeVideo(**_get_raw_data(args.video_path), fps=1)
    else:
        video = FFmpegVideo(args.video_path)

    representation_types = add_external_repositories(args.external_repositories, get_vre_repository())
    representations = build_representations_from_cfg(cfg=args.config_path, representation_types=representation_types,
                                                     external_representations=args.external_representations)

    vre = VRE(video, representations)
    vre.run(args.output_path, frames=args.frames or range(len(video)),
            output_dir_exists_mode=args.output_dir_exists_mode,
            n_threads_data_storer=args.n_threads_data_storer, exception_mode=args.exception_mode,
            exported_representations=args.representations)

    if args.collage:
        vre_collage_main = load_function_from_module(Path(__file__).parent / "vre_collage", "main")
        vre_collage_main(Namespace(vre_export_dir=args.output_path, output_path=args.collage_output_path,
                                   video=args.collage_video, config_path=args.config_path, fps=args.collage_fps,
                                   normalization=None, output_resolution=None,
                                   external_representations=args.external_representations,
                                   external_repositories=args.external_repositories, n_workers=1))

if __name__ == "__main__":
    main(get_args())
