import streamlit as st
from utils.levels import complete_level, render_page, initialize_level
from utils.login import get_login
import os
import uuid
from utils.login import initialize_login

initialize_login()
initialize_level()

LEVEL = 2


def step2_page():
    st.header("Collecting Your Data")
    st.write(
        "Now it's time to collect the pictures we need to teach our application about emotions. But remember, we should always ask for permission before taking someone's picture. We can use a smartphone or a digital camera to capture different facial expressions. It's important to take pictures of different people and show different emotions like happiness, sadness, surprise, and more. This will help our application understand emotions better!"
    )

    img_dir = os.path.join(".sessions", get_login()["username"], "images")
    os.makedirs(img_dir, exist_ok=True)

    emotion = st.selectbox(
        "Select an emotion",
        ("Happy", "Sad", "Surprised", "Angry", "Disgusted", "Neutral", "Fearful"),
    )
    picture = st.camera_input("Take a picture")
    if picture:
        if st.button("Save"):
            img_name = str(uuid.uuid4()) + f"_{emotion}_" + ".jpg"
            img_path = os.path.join(img_dir, img_name)
            with open(img_path, "wb") as f:
                f.write(picture.getvalue())
            st.success("Image added successfully!")

    images = os.listdir(img_dir)
    if len(images) > 0:
        st.subheader("Your Images")
        cols = st.columns(4)
        for i, img in enumerate(images):
            img_emotion = img.split("_")[1]
            cols[i % 4].image(os.path.join(img_dir, img), use_column_width=True)
            cols[i % 4].write(img_emotion)

    if st.button("Clear All"):
        for img in images:
            os.remove(os.path.join(img_dir, img))
        st.success("All images cleared!")
        st.experimental_rerun()

    st.info("If you are satisfied with your images, click on the button below to complete this level.")
    if st.button("Complete"):
        complete_level(LEVEL)


render_page(step2_page, LEVEL)