import gradio as gr
import threading
import os
from train import main as train_main
from compile_video import compile_gif_to_video

def train_model():
    try:
        yield "Training Started...\n"

        train_main()

        gif_path = "models/frames/training.gif"
        video_path = "models/training.mp4"
        if os.path.exists(gif_path):
            compile_gif_to_video(gif_path, video_path)
            yield "Training Completed! Model and video have been saved.\n"
        else:
            yield "Training Completed! Model has been saved.\n"
    except Exception as e:
        yield f"An error occurred during training: {e}\n"

with gr.Blocks() as demo:
    gr.Markdown("## Tetris RL Training")
    status = gr.Textbox(label="Status", lines=10, interactive=False)

    # Add a Start Training button
    start_button = gr.Button("Start Training")

    start_button.click(
        fn=train_model,
        inputs=[],
        outputs=status,
        show_progress=True,
    )

demo.launch()



os.makedirs("models/frames", exist_ok=True)