#!/usr/bin/env python3
"""vre_collage tool"""
from argparse import ArgumentParser, Namespace
from pathlib import Path
from natsort import natsorted
from tqdm import tqdm
from moviepy.editor import ImageSequenceClip
import numpy as np

from vre.utils import image_read, image_write, image_resize, image_add_title, collage_fn
from vre.logger import vre_logger as logger

def get_args():
    """Cli args"""
    parser = ArgumentParser()
    parser.add_argument("vre_export_dir", type=lambda p: Path(p).absolute(), help="Path to the dir where VRE exported")
    parser.add_argument("--output_path", "-o", required=False, help="Path to the collage is stored")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--video", action="store_true")
    parser.add_argument("--fps", type=int)
    parser.add_argument("--output_resolution", type=int, nargs=2)
    # TODO: add order of collage
    args = parser.parse_args()

    if args.output_path is None:
        args.output_path = args.vre_export_dir / "collage"
        logger.warning(f"--output_path not provided, defaulting to '{args.output_path}'")
    args.output_path = Path(args.output_path).absolute()
    # assert not args.output_path.exists() or args.overwrite, f"Output path '{args.output_path}' exists. Use --overwrite."
    assert args.video is False or args.fps is not None, "If video is set, --fps must also be set."
    return args

def resolve_files(vre_export_dir: Path) -> dict[str, list[Path]]:
    """Gets a list of all images that are going to be used for the collage."""
    found_dirs = [x / "png" for x in vre_export_dir.iterdir() if (x / "png").exists()]
    repr_files = {}
    for repr_png_dir in found_dirs:
        files = natsorted([x for x in repr_png_dir.iterdir()], key=lambda p: p.name)
        repr_files[repr_png_dir.parent.name] = files

    name_files = np.array([[x.name for x in v] for v in repr_files.values()])
    assert (name_files == name_files[0]).all(), f"Not all png files are named identically: {name_files}"
    n_files = {k: len(v) for k, v in repr_files.items()}
    assert np.std(list(n_files.values())) == 0, f"Not all png files have the same number of images: {n_files}"
    logger.info(f"Returning {n_files[found_dirs[0].parent.name]} images from {len(n_files)} representations.")
    return repr_files

def main(args: Namespace):
    """main fn"""
    logger.info(f"VRE pngs dir: {args.vre_export_dir}")
    logger.info(f"Output dir: {args.output_path}")
    logger.info(f"Export video: {args.video} {args.output_path / 'collage.mp4' if args.video else ''}")

    found_files = resolve_files(args.vre_export_dir)
    keys = found_files.keys()
    # convert [[key1_0, key1_1], [key2_0, key2_1]] to [[key1_0, key2_0], [key1_1, key2_1]]
    values = np.array(list(found_files.values())).T
    args.output_path.mkdir(parents=True, exist_ok=True)
    for current_frame_paths in tqdm(values, desc="collage images"):
        # export the png with the same name as all the other png from which it is constructed
        current_frames = [image_read(x) for x in current_frame_paths]
        h, w = current_frames[0].shape[0:2]
        # keep only first 25 chars of vre representation, otherwise the title is too long
        titles = [k[0:25] for k in keys]
        h_out, w_out = args.output_resolution if args.output_resolution is not None else (h, w)
        rsz_frames = [image_resize(img, height=h_out, width=w_out) for img in current_frames]
        titled_frames = [image_add_title(img, title, size_px=w_out // 15, top_padding=w_out // 10)
                         for img, title in zip(rsz_frames, titles)]
        img = collage_fn(titled_frames, pad_right=w_out // 75)
        image_write(img, args.output_path / f"{current_frame_paths[0].stem}.png")

    if args.video:
        files = natsorted([x for x in args.output_path.iterdir() if x.suffix == ".png"], key=lambda p: p.name)
        video = ImageSequenceClip(list(map(str, files)), fps=args.fps)
        video.write_videofile(f"{args.output_path}/collage.mp4")

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