import gradio as gr
import subprocess
import os
import shutil
import zipfile


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_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


def zip_files(file_paths):
    if not file_paths:
        raise gr.Error("No files to download. Please upload and convert AVIF files first.")
    zip_path = "/tmp/converted_images.zip"
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for file_path in file_paths:
            if os.path.isfile(file_path):
                zipf.write(file_path, os.path.basename(file_path))
    return zip_path


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")
        download_button = gr.Button("Download All as ZIP")

    converted_files = []

    def handle_conversion(avif_files):
        global converted_files
        converted_files = convert_avif_to_png(avif_files)
        return converted_files

    def handle_download():
        return zip_files(converted_files)

    run_button.click(
        fn=handle_conversion,
        inputs=[avif_file],
        outputs=[result]
    )
    download_button.click(
        fn=handle_download,
        inputs=[],
        outputs=[gr.File()]
    )
demo.queue().launch()