from face_deid_ct import drown_volume

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


def process_file(input_file):
    cache_dir = "cache"
    cache_out_dir = "cache_out"
    output_zip_file = "output.zip"

    # Check if input file is a zip file
    if zipfile.is_zipfile(input_file.name):
        with zipfile.ZipFile(input_file.name, 'r') as zip_ref:
            # Unzip the file in 'cache' directory
            zip_ref.extractall(cache_dir)
        
        # Run deid function
        drown_volume(cache_dir, cache_out_dir, replacer='face')
        
        # Create a Zip file for 'cache_out' directory
        with zipfile.ZipFile(output_zip_file, 'w') as zipf:
            for root, dirs, files in os.walk(cache_out_dir):
                for file in files:
                    zipf.write(os.path.join(root, file), 
                               os.path.relpath(os.path.join(root, file), 
                                               os.path.join(cache_out_dir, '..')))
        
        # Cleanup cache directories
        shutil.rmtree(cache_dir)
        shutil.rmtree(cache_out_dir)

        return output_zip_file

    else:
        raise ValueError("The provided file is not a zip file.")

description = "Upload a ZIP file containing a folder with a head CT's DICOM files. The ZIP file might also contain subfolders, each one containing a head CT."

inputs = gr.components.File(label="Input File")
outputs = gr.components.File(label="Output File")
demo = gr.Interface(fn=process_file, description=description, inputs=inputs, outputs=outputs)
demo.launch()