#!/usr/bin/env python3
from argparse import ArgumentParser, Namespace
from pathlib import Path
from natsort import natsorted
import imp
import numpy as np
from vre import VRE, FFmpegVideo, FakeVideo
from vre.utils import image_read, abs_path
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
from omegaconf import OmegaConf

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("--make_video_collage", 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")

    args = parser.parse_args()
    if args.make_video_collage:
        assert args.collage_fps is not None and args.collage_fps > 0, args.collage_fps
    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"""
    cfg = OmegaConf.to_container(OmegaConf.load(args.config_path), resolve=True)

    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=cfg, representation_types=representation_types,
                                                     external_representations=args.external_representations)

    vre = VRE(video, representations)
    vre_metadata = 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)
    print(vre_metadata.pretty_format())

    if args.make_video_collage:
        vre_collage = imp.load_source("vre_collage", (Path(__file__).parent / "vre_collage").__str__())
        vre_collage.main(Namespace(vre_export_dir=args.output_path, output_path=args.output_path / "collage",
                                   video=True, config_path=args.config_path, fps=args.collage_fps, normalization=None,
                                   output_resolution=None, external_representations=[], external_repositories=[],
                                   n_workers=1))

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