import os
os.environ["MUJOCO_GL"] = "osmesa"  # or "egl" if you switch to EGL support

import gradio as gr
from tray_sim import run_tray_simulation

import time  # For unique file naming

# Global counter for unique file naming
run_counter = 0

def simulate_and_render(
    seed: int,
    azimuth: float = 35,
    elevation: float = -20,
    distance: float = 1.0
):
    """
    Runs a TraySim simulation with the specified camera parameters and renders the result.

    Parameters:
        seed (int): Random seed for simulation initialization, affects object placement and dynamics.
        azimuth (float, optional): Horizontal camera angle in degrees. Default is 35.
        elevation (float, optional): Vertical camera angle in degrees. Default is -20.
        distance (float, optional): Distance of the camera from the scene. Default is 1.0.

    Returns:
        str: Path to the rendered GIF file showing the simulation from the specified viewpoint.
    """
    global run_counter
    run_counter += 1
    unique_id = f"{seed}_{run_counter}_{int(time.time())}"  # Unique identifier for this run
    
    gif_path, stability_flags, physics_state, llm_friendly_output, json_path  = run_tray_simulation(
        seed=seed,
        azimuth=azimuth,
        elevation=elevation,
        distance=distance,
        unique_id=unique_id  # Pass unique_id to avoid file conflicts
    )

    # Optional debug
    print("Stability: ", stability_flags)
    print("Physics State: ", physics_state)
    print("LLM Output: ", llm_friendly_output)
    print("JSON Path: ", json_path)
    
    return gif_path

#iface = gr.Interface(
#    fn=simulate_and_render,
#    inputs=[
#        gr.Number(label="Random Seed", value=0),
#        gr.Slider(minimum=0, maximum=360, label="Azimuth", value=35),
#        gr.Slider(minimum=-90, maximum=90, label="Elevation", value=-20),
#        gr.Slider(minimum=0.2, maximum=1.0, label="Zoom (Distance)", value=1.0)
#    ],
#    outputs=gr.Image(type="filepath", label="Simulation"),
#    title="TraySim Simulation Viewer",
#    description="Simulates object drop on a tray using MuJoCo and shows the result as a GIF."
#)

# Define the custom layout using Blocks
with gr.Blocks(
    css="""#simulation_image {min-height: 600px; min-width: 700px;}""",  # Set minimum size for the image
    theme=gr.themes.Default()
) as iface:
    gr.Markdown("# TraySim Simulation Viewer")
    gr.Markdown("Simulates object drop on a tray using MuJoCo and shows the result as a GIF.")
    
    with gr.Row():
        # Input column (30% width)
        with gr.Column(scale=3):  # scale=3 for 30% width
            seed_input = gr.Number(label="Random Seed", value=0)
            azimuth_input = gr.Slider(minimum=0, maximum=360, label="Azimuth", value=35)
            elevation_input = gr.Slider(minimum=-90, maximum=90, label="Elevation", value=-20)
            distance_input = gr.Slider(minimum=0.2, maximum=1.0, label="Zoom (Distance)", value=1.0)
            submit_button = gr.Button("Run Simulation")
        
        # Output column (70% width)
        with gr.Column(scale=7):  # scale=7 for 70% width
            output_image = gr.Image(
                type="filepath",
                label="Simulation",
                elem_id="simulation_image",  # Assign an ID for CSS targeting
                width=700,  # Set default width
                height=600  # Set default height
            )

    # Connect the button to the simulation function
    submit_button.click(
        fn=simulate_and_render,
        inputs=[seed_input, azimuth_input, elevation_input, distance_input],
        outputs=output_image
    )

if __name__ == "__main__":
    iface.launch()