import math
import streamlit as st
from PIL import Image
import cv2
import os
import threading

hide_streamlit_style = """
<style>
    #root > div:nth-child(1) > div > div > div > div > section > div {padding-top: 0rem;
    padding-left: 1%;
    }
</style>
"""

st.set_page_config(layout="wide")

def loadModel(n):
    super_res = cv2.dnn_superres.DnnSuperResImpl_create()
    super_res.readModel('models/ESPCN_x'+n+'.pb')
    return super_res

# on removing (show_spinner=False), it will show that function is running on web app
@st.experimental_memo(show_spinner=False)
def upscale(file, task):
    with open(file.name, "wb") as f:
        f.write(file.getbuffer())
        print('No file found, so added in list files')
    if isinstance(task, str):
        super_res = loadModel(task)
        super_res.setModel('espcn', int(task))
        if file.type.split('/')[0] == 'image':
            img = cv2.imread(file.name)
            upscaled_image = super_res.upsample(img)
            print('I upscaled up to', task, 'times')
            output_file_path = "processed_" + file.name.split('.')[0] + ".png"
            cv2.imwrite(output_file_path, upscaled_image)
            with st.sidebar:
                st.success('Done!', icon="✅")
            return output_file_path
        elif file.type.split('/')[0] == 'video':
            # 영상 업스케일링 코드는 생략 (필요시 추가 가능)
            pass
        return True
    else:
        # 커스텀 사이즈 코드 (필요시 추가)
        pass

    return True

if 'disable_opt2' not in st.session_state:
    st.session_state.disable_opt2 = True
if 'disable_opt1' not in st.session_state:
    st.session_state.disable_opt1 = False
if 'disable_download' not in st.session_state:
    st.session_state.disable_download = True
if 'disable_proceed' not in st.session_state:
    st.session_state.disable_proceed = False

st.markdown(hide_streamlit_style, unsafe_allow_html=True)

col1, _, col2 = st.columns([6, 1, 3], gap="small")

def toggle_state_opt1():
    if st.session_state.get("opt1") == True:
        st.session_state.opt2 = False
        st.session_state.disable_opt2 = True
    else:
        st.session_state.opt2 = True
        st.session_state.disable_opt2 = False

def toggle_state_opt2():
    if st.session_state.get("opt2") == True:
        st.session_state.opt1 = False
        st.session_state.disable_opt1 = True
    else:
        st.session_state.opt1 = True
        st.session_state.disable_opt1 = False

toggle_state_opt2()
toggle_state_opt1()
options = ["2", "3", "4"]

with col1:
    file = st.file_uploader(" ", type=['png', 'jpeg', 'jpg', 'pgm', 'jpe', 'mp4', 'mov'])
    if file is not None:
        bytes_data = file.getvalue()
        file_size = len(bytes_data)
        print("File size: ", file_size)

        if file.type.split('/')[0] == "image" and file_size > 1550000:
            st.session_state.disable_proceed = True
            with st.sidebar:
                st.info('Sorry, maximum size of image is 1.5MB', icon="ℹ️")
        elif file.type.split('/')[0] == "image":
            image = Image.open(file)
            st.session_state.disable_proceed = False
            st.image(image, caption="Upload Image", use_column_width=True)
            st.session_state.disable_proceed = False
        elif file.type.split('/')[0] == 'video' and file_size > 200000000:
            with st.sidebar:
                options = ["2", "3"]
                st.info('Sorry, maximum size of video is 200MB', icon="ℹ️")
                st.session_state.disable_proceed = True
        elif file.type.split('/')[0] == 'video':
            video = st.video(file)
            print(type(video))
            options = ["2", "3"]
            st.session_state.disable_proceed = False
            with st.sidebar:
                st.info('For custom size, currently I can process video without AI.', icon="ℹ️")

with col2:
    st.markdown("\n")
    st.markdown("\n")
    st.markdown("\n")

    st.subheader(" UPSCALE RESOLUTION UP TO")
    st.markdown("\n")
    st.markdown("\n")

    opt1 = st.checkbox("MULTIPLES OF", key="opt1", value=True, on_change=toggle_state_opt1)
    st.selectbox("SELECT", options, key="opt1_selBox", disabled=st.session_state.disable_opt1)

    st.markdown("\n")
    st.markdown("\n")
    opt2 = st.checkbox("CUSTOM SIZE", key="opt2", on_change=toggle_state_opt2)

    st.number_input("Width", step=1, min_value=150, max_value=3840, value=900, key="width", disabled=st.session_state.disable_opt2)
    st.number_input("Height", step=1, min_value=150, max_value=2160, value=900, key="height", disabled=st.session_state.disable_opt2)

    st.markdown("\n")
    st.markdown("\n")

    if st.button(15*"&nbsp;"+"PROCEED"+"&nbsp;"*15, disabled=st.session_state.disable_proceed) and file is not None:
        if st.session_state.get('opt1') == True:
            task = st.session_state.opt1_selBox
        else:
            task = [st.session_state.width, st.session_state.height]
        print(task)
        result_file_path = upscale(file, task)

        if result_file_path:
            st.session_state.disable_download = False
            st.session_state.result_file_path = result_file_path

    st.markdown("\n")
    st.markdown("\n")

    if file is None:
        st.session_state.disable_download = True

    if st.session_state.disable_download:
        st.button(13*"&nbsp;"+"DOWNLOAD"+"&nbsp;"*13, disabled=True)
    else:
        with open(st.session_state.result_file_path, "rb") as download_file:
            st.download_button(label=13*"&nbsp;"+"DOWNLOAD"+"&nbsp;"*13, data=download_file, file_name='processed_' + file.name.split('.')[0] + ".png", mime="image/png")

st.markdown("\n")
st.markdown("\n")