import streamlit as st
import os
import zipfile
import subprocess

logo_image_path = '1.png'

st.image(logo_image_path, width=150)
st.title("Automated Abdominal Aortic Aneurysm Detection")

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

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

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

def search_file(start_path, target_file):
    for root, _, files in os.walk(start_path):
        if target_file in files:
            return os.path.join(root, target_file)
    return None

def resolve_path(base_path):
    return os.path.abspath(os.path.join(base_path, "C2C/comp2comp/utils/../../outputs"))

def fetch_outputs():
    outputs_base_path = resolve_path("/home/user/app")
    if os.path.exists(outputs_base_path):
        subdirectories = [subdir for subdir in os.listdir(outputs_base_path)
                          if os.path.isdir(os.path.join(outputs_base_path, subdir))]

        first_subdirectory = subdirectories[0] if subdirectories else None

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

            for root, dirs, _ in os.walk(temp_dicom_dir_path):
                for subdir in dirs:
                    dicom_folder_path = os.path.join(root, subdir)
                    video_path = search_file(dicom_folder_path, "aaa.mp4")
                    if video_path:  # Use this folder if it contains the mp4 file
                        image_path = search_file(dicom_folder_path, "diameter_graph.png")
                        largest_slice = search_file(dicom_folder_path, "out.png")

                        if video_path and image_path and largest_slice:
                            zip_filename = os.path.join("/home/user/app/C2C/temp_dicom_dir", "outputs.zip")
                            zip_and_download(dicom_folder_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)
                            return
            st.error("Output files not found in any subdirectory")
        else:
            st.warning("Output files not found")
    else:
        st.error("Output directory not found")

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

            if st.button("Fetch Outputs"):
                fetch_outputs()
        except Exception as e:
            st.error(f"Error: {str(e)}")

if __name__ == "__main__":
    main()