import streamlit as st
import pydicom
import zipfile
import os
import subprocess

subprocess_executed = False

# Function to install dependencies
@st.cache_resource
def install_dependencies():
    command = "chmod +x install.sh"
    subprocess.run(command, shell=True, check=True)

    command = "./install.sh"
    subprocess.run(command, shell=True, check=True)

# Function to run inference
@st.cache_resource
def run_inference():
    command = "chmod +x inference.sh"
    subprocess.run(command, shell=True, check=True)

    command = "./inference.sh"
    subprocess.run(command, shell=True, check=True)

# Function to zip and download files
def zip_and_download(directory, zip_filename):
    zipf = zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED)
    for root, _, files in os.walk(directory):
        for file in files:
            zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), directory))
    zipf.close()

    with open(zip_filename, 'rb') as f:
        bytes_data = f.read()
    st.download_button(label="Download Outputs", data=bytes_data, file_name=zip_filename, mime="application/zip")

# Main Streamlit code
def main():
    st.write("Upload a ZIP file containing DICOM slices")
    uploaded_zip_file = st.file_uploader("Upload a .zip file", type=["zip"])

    if uploaded_zip_file is not None:
        try:
            install_dependencies()

            temp_dir = "/home/user/app/C2C/temp_dicom_dir"
            os.makedirs(temp_dir, exist_ok=True)

            with zipfile.ZipFile(uploaded_zip_file, "r") as zip_ref:
                zip_ref.extractall(temp_dir)

            st.success("Zip file uploaded successfully")

            if st.button("Analyze"):
                st.success("Analysis started (expected time: 5 mins)")
                run_inference()

                outputs_directory = "/home/user/app/C2C/outputs"

                if os.path.exists(outputs_directory):
                    subdirectories = [subdir for subdir in os.listdir(outputs_directory)
                                      if os.path.isdir(os.path.join(outputs_directory, subdir))]

                    first_subdirectory = subdirectories[0] if subdirectories else None

                    if first_subdirectory:
                        subdirectory_path = os.path.join(outputs_directory, first_subdirectory)
                        temp_dicom_dir_path = os.path.join(subdirectory_path, "temp_dicom_dir")

                        video_path = os.path.join(temp_dicom_dir_path, "images/summary/aaa.mp4")
                        image_path = os.path.join(temp_dicom_dir_path, "images/summary/diameter_graph.png")
                        largest_slice = os.path.join(temp_dicom_dir_path, "images/summary/out.png")

                        if os.path.exists(video_path) and os.path.exists(image_path) and os.path.exists(largest_slice):
                            zip_filename = os.path.join(temp_dir, "outputs.zip")
                            zip_and_download(temp_dicom_dir_path, zip_filename)

                            st.title("Largest Slice")
                            st.image(largest_slice, use_column_width=True)

                            st.title("Video")
                            st.video(video_path, format="video/mp4")

                            st.title("Diameter Graph")
                            st.image(image_path, use_column_width=True)
                        else:
                            st.error("Output files not found.")
                    else:
                        st.warning("No subdirectories found inside 'outputs'.")
                else:
                    st.error("The 'outputs' directory does not exist.")
        except Exception as e:
            st.error(f"Error: {str(e)}")

if __name__ == "__main__":
    main()