import random
from PIL import Image
import os
import shutil
import streamlit as st
from about import about

st.set_page_config(
    page_title=None,
    page_icon=None,
    layout="wide",
    initial_sidebar_state="auto",
    menu_items=None,
)
st.title("Image Enhancer")

if os.path.exists("results"):
    shutil.rmtree(os.path.join("results"))

if os.path.exists("tempDir"):
    shutil.rmtree(os.path.join("tempDir"))


def create_dir(dirname: str):
    if not os.path.exists(dirname):
        os.makedirs(dirname, exist_ok=True)


create_dir("results/cmp")
create_dir("results/cropped_faces")
create_dir("results/restored_faces")
create_dir("results/restored_imgs")
create_dir("tempDir")


def save_uploadedfile(uploadedfile):
    file_extension = os.path.splitext(uploadedfile.name)[-1].lstrip(".")
    with open(os.path.join("tempDir", f"uploaded_image.{file_extension}"), "wb") as f:
        f.write(uploadedfile.getbuffer())

    name, path = f'uploaded_image.{file_extension}', os.path.join("tempDir", f"uploaded_image.{file_extension}")
    return name, path


def get_random_sample_image():
    sample_images = os.listdir(os.path.join('sample_images'))
    random_image = random.choice(sample_images)
    return random_image, f"{os.getcwd()}/sample_images/{random_image}"


def results_view(name, path):
    with st.spinner("Please wait while we process your image.."):
        os.system(f"python inference_gfpgan.py -i {path} -o results -v {version} -s 2")
    with st.expander("Results", expanded=True):
        col_1_1, col_2_2 = st.columns(2)
        with col_1_1:
            st.write("Sample Image")
            st.image(path)
        with col_2_2:
            st.write("Processed Image")
            st.image(Image.open(os.path.join("results", "restored_imgs", name)))
    with st.expander("Comparative Results", expanded=True):
        files = os.listdir(os.path.join("results", "cmp"))
        for f in files:
            st.write(f)
            st.image(Image.open(os.path.join("results", "cmp", f)))


options = ["Face Restore", "About"]
models = ['1.3', '1.4', "RestoreFormer"]

st.sidebar.image("assets/gfpgan_logo.png")
menu = st.sidebar.selectbox("Select an Option", options)

if menu == "Face Restore":
    col1, col2 = st.columns([1, 0.3])
    with st.sidebar:
        version = st.selectbox("Select Version", models)
    with col1:
        uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
    with col2:
        st.markdown("###")
        st.markdown("###")
        sample = st.button("Use Sample Image")

    if sample:
        name, path = get_random_sample_image()
        # with st.spinner("Please wait while we process your image.."):
        #     os.system(f"python inference_gfpgan.py -i {path} -o results -v {version} -s 2")
        #     with st.expander("Results", expanded=True):
        #         col_1_1, col_2_2 = st.columns(2)
        #         with col_1_1:
        #             st.write("Sample Image")
        #             st.image(path)
        #         with col_2_2:
        #             st.write("Processed Image")
        #             st.image(Image.open(os.path.join("results", "restored_imgs", name)))
        #     with st.expander("Comparative Results", expanded=True):
        #         files = os.listdir(os.path.join("results", "cmp"))
        #         for f in files:
        #             st.write(f)
        #             st.image(Image.open(os.path.join("results", "cmp", f)))
        results_view(name, path)
    if uploaded_file is not None:
        name, path = save_uploadedfile(uploaded_file)
        results_view(name, path)

if menu == 'About':
    about()