import subprocess
import tyro


import re
import subprocess
from pathlib import Path
from dataclasses import dataclass
from benchmark import get_output_dir, get_version_name, run_command


@dataclass
class VideoGeneratorConfig:

    scene_name : str
    """Name of the scene"""
    
    dataset_path: Path
    """Path to dataset"""


def generate_video_from_images(output_path: Path, mode, output_video="output.mp4"):
    # Step 1: Collect and parse filenames
    files = []
    for directory in output_path.iterdir():
        if directory.is_dir():
            for file in directory.joinpath(mode).glob("*.jpg"):
                match = re.search(r"(\d+)", file.name)
                if match:
                    files.append((int(match.group(1)), file))

    sorted_files = sorted(files, key=lambda x: x[0])
    assert len(sorted_files) > 0, f"Error processing directory: {output_path}"

    filelist_path = Path("concat.txt")
    with filelist_path.open("w") as f:
        for _, filepath in sorted_files:
            f.write(f"file '{filepath.as_posix()}'\n")
            print(filepath)

    assert filelist_path.exists(), "Well, something went wrong writing the file."

    ffmpeg_command = [
        "ffmpeg", "-y",
        "-f", "concat",
        "-i", str(filelist_path), 
        "-filter_complex", "[0:v]scale=iw*2:ih*2:flags=lanczos,fps=24",
        "-c:v", "libx264",
        "-crf", "23",
        "-preset", "veryfast",
        output_video
    ]
    subprocess.run(ffmpeg_command)

    filelist_path.unlink(missing_ok=True)
    return output_video


def main(config: VideoGeneratorConfig):
    video_dir = Path("./rendered_video")
    if not video_dir.exists():
        video_dir.mkdir()

    scene = config.scene_name
    dataset = config.dataset_path.name
    dataset_path = config.dataset_path

    assert dataset in ["mipnerf360", "internal"], "Dataset is not currently supported."

    video_file = dataset_path.joinpath(scene, "processed", "camera_paths", f"wiggle_{scene}.json")
    if not video_file.exists():
        raise FileNotFoundError(f"{video_file}. Please export a file called `{video_file.name}` from nerfstudio into {video_file.parent}.")

    output_dir = get_output_dir(scene)
    models = {"qr": "raytracingfacto", "nerfstudio": "dn-splatter"}
    modes = ["rgb", "depth"]
    for key, model in models.items():
        version_name = get_version_name(output_dir, dataset, model, False)
        config_file = output_dir.joinpath(dataset, model, version_name, "config.yml")
        print(config_file)
        for mode in modes:
            video_path = video_dir.joinpath(f"{model}_wiggle_{scene}_{mode}")
            video_filename = video_dir.joinpath(f"{video_path.name}.mp4")
            if not video_path.exists():
                video_path.mkdir()
            commad = [
                "ns-render", "interpolate"
                , "--load_config", config_file.as_posix()
                , "--rendered_output_names", mode
                , "--output_path", video_filename.as_posix()
                , "--pose-source", "train"
            ]
            run_command({key: commad})

            # video_filename = video_dir.joinpath(f"{video_path.name}.mp4")
            # output_file = generate_video_from_images(video_path, mode, video_filename)
            # assert output_file.exists(), "Generation failed."


if __name__ == "__main__":
    config = tyro.cli(VideoGeneratorConfig)
    main(config)



