import streamlit as st
import hmac

# Standard imports
import numpy as np
import pandas as pd
from datetime import datetime
import os
import random

# Import torch
import torch

# Data visualization
import matplotlib.pyplot as plt
import seaborn as sns

# Path manipulation
from pathlib import Path
import sys

# Custom and other imports
import project_config
from utils import add_logo
from menu import menu

# Insert logo
# add_logo(project_config.MEDIA_DIR / 'gravity_logo.png')

# Initialize st.session_state.role to None
if "role" not in st.session_state:
    st.session_state.role = None

# # Retrieve the role from Session State to initialize the widget
# st.session_state._role = st.session_state.role

# def set_role():
#     # Callback function to save the role selection to Session State
#     st.session_state.role = st.session_state._role

def check_password():
    """Returns `True` if the user had a correct password."""

    def login_form():
        """Form with widgets to collect user information"""
        with st.form("Credentials"):
            st.text_input("Username", key="username")
            st.text_input("Password", type="password", key="password")
            st.form_submit_button("Log In", on_click=password_entered)

    def password_entered():
        """Checks whether a password entered by the user is correct."""
        if st.session_state["username"] in st.secrets[
            "passwords"
        ] and hmac.compare_digest(
            st.session_state["password"],
            st.secrets.passwords[st.session_state["username"]],
        ):
            st.session_state["password_correct"] = True

            # Retrieve user name
            if st.session_state["username"] in st.secrets["users"]:
                st.session_state["name"] = st.secrets["users"][st.session_state["username"]]
            else:
                st.session_state["name"] = st.session_state["username"]

            # Check if the username is an admin
            if st.session_state["username"] in st.secrets["admins"]['admins']:
                st.session_state["role"] = "admin"
            else:
                st.session_state["role"] = "user"

            # Check if username is in a team
            if st.session_state["username"] in st.secrets["teams"]:
                st.session_state["team"] = st.secrets["teams"][st.session_state["username"]]
            else:
                st.session_state["team"] = None

            # Don't store the username or password
            del st.session_state["password"]
            # del st.session_state["username"]
        else:
            st.session_state["password_correct"] = False

    # Return True if the username + password is validated
    if st.session_state.get("password_correct", False):
        return True

    # Show inputs for username + password
    login_form()
    if "password_correct" in st.session_state:
        st.error("😕 User not known or password incorrect.")
    return False

# # Selectbox to choose role
# st.selectbox(
#     "Select your role:",
#     [None, "user", "admin", "super-admin"],
#     key="_role",
#     on_change=set_role,
# )
menu() # Render the dynamic menu!

if not check_password():
    st.stop()