import gradio as gr
import subprocess
import os
import shutil


def setup_avif_decode():
    # Install Rust and Cargo
    if not os.path.exists(os.path.expanduser("~/.cargo/bin/cargo")):
        subprocess.run("apt-get update && apt-get install -y yasm", shell=True, check=True)
        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")
    if not os.path.exists("avif-decode"):
        subprocess.run("git clone https://github.com/kornelski/avif-decode.git", shell=True, check=True)
    subprocess.run("cd avif-decode && cargo build --release", shell=True, check=True)
setup_avif_decode()


def convert_avif_to_png(avif_file):
    def convert_avif_to_png(avif_files):
    output_paths = []
    for avif_file in avif_files:
        avif_path = avif_file.name
        png_path = avif_path.rsplit('.', 1)[0] + '.png'   
        result = subprocess.run(["avif-decode/target/release/avif_decode", "-f", avif_path, png_path], capture_output=True, text=True)
        if result.returncode == 0:
            output_paths.append(png_path)
        else:
            output_paths.append(f"Error converting {avif_file.name}: {result.stderr}")
    return output_paths


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()