import streamlit as st
import numpy as np
from apps.data_preprocessing import preprocess_input
import os
import tempfile
import nibabel as nib

st.title("Dental medical imaging 3D")
st.write("This tool accepts NIfTI `.nii` or `.nii.gz` files smaller than 400MB.")
uploaded_file = st.file_uploader("Upload a NIfTI file (.nii or .nii.gz)", type=["nii", "nii.gz"])

if uploaded_file is not None:
    with tempfile.TemporaryDirectory() as temp_dir:
        try:
            # Save uploaded file to a temporary location
            nifti_path = os.path.join(temp_dir, "uploaded_file.nii.gz")
            with open(nifti_path, "wb") as f:
                f.write(uploaded_file.read())

            # Preprocess the uploaded NIfTI file
            output_tensor, dataset = preprocess_input(nifti_path)
            imgs = dataset[0]["CT"]["data"]
            pred = output_tensor.argmax(0)
            numpy_array = pred.numpy().astype(np.int16)
            affine = np.eye(4)

            # Save finalized mask
            mask_path = os.path.join(temp_dir, "finalized_mask.nii")
            nifti_image_mask = nib.Nifti1Image(numpy_array, affine)
            nib.save(nifti_image_mask, mask_path)

            # Save finalized image
            nifti_array = imgs.squeeze(0).numpy()
            image_path = os.path.join(temp_dir, "finalized_image.nii.gz")
            nifti_image_img = nib.Nifti1Image(nifti_array, affine)
            nib.save(nifti_image_img, image_path)

            # Provide download buttons for the finalized files
            with open(mask_path, "rb") as f:
                st.download_button(
                    label="Download Output Mask",
                    data=f,
                    file_name="finalized_mask.nii",
                    mime="application/octet-stream"
                )

            with open(image_path, "rb") as f:
                st.download_button(
                    label="Download The Image",
                    data=f,
                    file_name="finalized_image.nii.gz",
                    mime="application/octet-stream"
                )

        except Exception as e:
            st.error(f"An error occurred: {e}")
else:
    st.info("Please upload a NIfTI file to get preprocessed image and mask.")