import gradio as gr
import subprocess
import os
import shutil

def get_env_info():
    files = os.listdir('.')
    return f"Current Directory: {os.getcwd()}\nFiles: {files}"

def rig_model(input_file):
    # The official script from the GitHub
    script = "./demo_rigging.sh"
    
    if not os.path.exists(script):
        return None, f"Error: {script} not found. Directory contains: {os.listdir('.')}"

    output_dir = "results/final_rigging"
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir, exist_ok=True)
            
    try:
        # Run the official bash command exactly as instructed
        # We use -x to see the bash progress in the logs
        result = subprocess.run(
            ["bash", "-x", script, input_file.name], 
            capture_output=True, 
            text=True
        )
        
        # Check for generated output
        generated = [f for f in os.listdir(output_dir) if f.endswith('.glb')]
        
        if generated:
            return os.path.join(output_dir, generated[0]), "Rigging Complete!"
        else:
            return None, f"Script finished but no GLB was found.\n\nLOGS:\n{result.stderr}\n{result.stdout}"
                
    except Exception as e:
        return None, f"System Error: {str(e)}"

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎭 Puppeteer Official Setup (Fixed Paths)")
    gr.Markdown(get_env_info())
    
    with gr.Row():
        input_file = gr.File(label="Upload Static GLB")
        run_btn = gr.Button("Start Rigging", variant="primary")
    
    out_file = gr.File(label="Rigged Result")
    out_logs = gr.Textbox(label="Process Logs", lines=12)

    run_btn.click(fn=rig_model, inputs=[input_file], outputs=[out_file, out_logs])

demo.launch(server_name="0.0.0.0", server_port=7860)