import gradio as gr
import subprocess
import os
import shutil

# Function to install Rust and Cargo, clone and build avif-decode
def setup_avif_decode():
    # Install Rust and Cargo
    if not os.path.exists(os.path.expanduser("~/.cargo/bin/cargo")):
        subprocess.run("curl https://sh.rustup.rs -sSf | sh -s -- -y", shell=True, check=True)
        os.environ['PATH'] += os.pathsep + os.path.expanduser("~/.cargo/bin")

    # Clone avif-decode if it doesn't exist
    if not os.path.exists("avif-decode"):
        subprocess.run("git clone https://github.com/kornelski/avif-decode.git", shell=True, check=True)

    # Build avif-decode
    subprocess.run("cd avif-decode && cargo build --release", shell=True, check=True)

# Call setup function to ensure everything is installed and built
setup_avif_decode()

# Define the function to convert AVIF to PNG
def convert_avif_to_png(avif_file):
    avif_path = avif_file.name
    png_path = avif_path.rsplit('.', 1)[0] + '.png'
    
    # Run the avif-decode command
    result = subprocess.run(["avif-decode/target/release/avif_decode", "-f", avif_path, png_path], capture_output=True, text=True)
    
    if result.returncode == 0:
        return png_path
    else:
        return f"Error converting {avif_file.name}: {result.stderr}"

css = """
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
        # AVIF to PNG Converter
        Upload your AVIF files and get them converted to PNG.
        """)
        
        with gr.Row():
            
            avif_file = gr.File(
                label="Upload AVIF File",
                file_types=[".avif"],
                file_count="multiple"
            )
            
            run_button = gr.Button("Convert", scale=0)
        
        result = gr.Gallery(label="Result")

    run_button.click(
        fn=convert_avif_to_png,
        inputs=[avif_file],
        outputs=[result]
    )

demo.queue().launch()