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

def run_segmentation(uploaded_file, modality):
    job_id = str(uuid.uuid4())
    input_filename = f"input_{job_id}.nii.gz"
    output_folder = f"segmentations_{job_id}"
    
    with open(input_filename, "wb") as f:
        f.write(uploaded_file.read())

    command = ["TotalSegmentator", "-i", input_filename, "-o", output_folder]
    if modality == "MR":
        command.extend(["--task", "total_mr"])
    
    try:
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        return f"Error during segmentation: {e}"
    
    zip_filename = f"segmentations_{job_id}.zip"
    with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(output_folder):
            for file in files:
                file_path = os.path.join(root, file)
                arcname = os.path.relpath(file_path, output_folder)
                zipf.write(file_path, arcname)
    
    os.remove(input_filename)
    shutil.rmtree(output_folder)
    
    return zip_filename

with gr.Blocks() as demo:
    gr.Markdown("# TotalSegmentator Gradio App")
    gr.Markdown(
        "Upload a CT or MR image (in NIfTI format) and run segmentation using TotalSegmentator. "
        "For MR images, the task flag is set accordingly."
    )
    
    with gr.Row():
        uploaded_file = gr.File(label="Upload NIfTI Image (.nii.gz)")
        modality = gr.Radio(choices=["CT", "MR"], label="Select Image Modality", value="CT")
    
    output_file = gr.File(label="Download Segmentation Output (zip)")
    
    run_btn = gr.Button("Run Segmentation")
    run_btn.click(fn=run_segmentation, inputs=[uploaded_file, modality], outputs=output_file)

demo.launch()