import gradio as gr
import os
import subprocess
import signal
import atexit
import platform

# Global variable to hold the process for the viewer
viser_process = None

def kill_process(proc):
    if proc is None:
        return
    if platform.system() == "Windows":
        proc.kill()
    else:
        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)


# Clean up function to kill the viser process when the app exits
def cleanup():
    print("Cleaning up and stopping viser server...")
    global viser_process
    if viser_process:
        kill_process(viser_process)
        viser_process = None

atexit.register(cleanup)


def get_experiment_ids():
    print("Getting experiment ids")
    video_files = os.listdir("video")
    print(f"Found {len(video_files)} video files")
    return sorted([os.path.splitext(f)[0] for f in video_files if f.endswith(".mp4")])

# Get absolute paths for all files
def get_absolute_path(relative_path):
    """Convert relative path to absolute path"""
    return os.path.abspath(relative_path)

def check_file_exists(file_path):
    """Check if file exists and return absolute path or None"""
    abs_path = get_absolute_path(file_path)
    if os.path.exists(abs_path):
        return abs_path
    else:
        print(f"Warning: File not found: {file_path}")
        return None
def get_file_paths(exp_id):
    paths = {
        "video": get_absolute_path(os.path.join("video", f"{exp_id}.mp4")),
        "yolo": get_absolute_path(os.path.join("yolo", f"yolo{exp_id}.mp4")),
        "optical_flow": get_absolute_path(os.path.join("optical_flow", f"optical_flow{exp_id}.mp4")),
        "pcd": get_absolute_path(os.path.join("pcd", f"vggt_{exp_id}.pcd")),
    }
    for key, path in paths.items():
        abs_path = check_file_exists(path)
        paths[key] = abs_path
    return paths

def show_results(exp_id):
    # global viser_process
    
    # if viser_process:
    #     print("Stopping previous viser server...")
        # kill_process(viser_process)
    #     viser_process = None
    print(f"Showing results for {exp_id}")
    paths = get_file_paths(exp_id)
    
    pcd_path = paths.get("pcd")
    if pcd_path and os.path.exists(pcd_path):
        print(f"Starting viser server for {pcd_path}...")
        # For cross-platform compatibility, use shell=True on Windows
        use_shell = platform.system() == "Windows"
        preexec_fn = os.setsid if platform.system() != "Windows" else None
        
        viser_process = subprocess.Popen(
            ["python", "vis.py", pcd_path],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            shell=use_shell,
            preexec_fn=preexec_fn
        )
        print(f"Viser server started with PID: {viser_process.pid}")
    
    iframe_html = '<iframe src="http://127.0.0.1:8088" style="width: 100%; height: 500px; border: none;"></iframe>'
    
    return (
        paths.get("video"),
        paths.get("yolo"),
        paths.get("optical_flow"),
        iframe_html
    )

with gr.Blocks() as demo:
    print("Creating blocks")
    gr.Markdown("# Experiment Results Gallery")
    
    experiment_ids = get_experiment_ids()
    print(f"Experiment ids: {experiment_ids}")
    if not experiment_ids:
        gr.Markdown("## No experiments found. Please add videos to the 'video' directory.")
    else:
        with gr.Row():
            exp_id_dropdown = gr.Dropdown(
                choices=experiment_ids, 
                label="Select Experiment ID",
                value=experiment_ids[0]
            )

        gr.Markdown("## Input")
        input_video = gr.Video(label="Input Video")

        gr.Markdown("## Outputs")
        with gr.Tabs():
            with gr.TabItem("YOLO"):
                yolo_video = gr.Video(label="YOLO Output")
            with gr.TabItem("Optical Flow"):
                optical_flow_video = gr.Video(label="Optical Flow Output")
            with gr.TabItem("Point Cloud"):
                pcd_vis = gr.HTML(label="Point Cloud Visualization")

        print("Creating dropdown")
        exp_id_dropdown.change(
            fn=show_results,
            inputs=exp_id_dropdown,
            outputs=[input_video, yolo_video, optical_flow_video, pcd_vis]
        )
        print("Creating load")
        # Trigger initial load
        demo.load(
            fn=show_results,
            inputs=exp_id_dropdown,
            outputs=[input_video, yolo_video, optical_flow_video, pcd_vis]
        )


if __name__ == "__main__":
    print("Launching demo")
    demo.launch(
        server_name="127.0.0.1",
        server_port=7862,
        show_error=True,
        inbrowser=True
    )