import streamlit as st
from authlib.integrations.requests_client import OAuth2Session
from datetime import datetime, timedelta

# Set up OAuth2 client credentials
CLIENT_ID = '611244422056-9ls8ul7vveedbmvloc9jugogmhp8ogrm.apps.googleusercontent.com'
CLIENT_SECRET = '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'
OAUTH2_PROVIDER_URL = 'https://accounts.google.com'
OAUTH2_ACCESS_TOKEN_URL = 'https://oauth2.googleapis.com/token'
OAUTH2_AUTHORIZE_URL = 'https://accounts.google.com/o/oauth2/auth'

# Define the login function
def login():
    # Set up the OAuth2 session
    oauth2_client = OAuth2Session(CLIENT_ID, CLIENT_SECRET, scope='openid email profile', redirect_uri='urn:ietf:wg:oauth:2.0:oob')

    # Create the authorization URL
    authorization_url, state = oauth2_client.create_authorization_url(OAUTH2_AUTHORIZE_URL, hd='yourdomain.com', prompt='select_account')

    # Show the authorization URL and get the authorization code from the user
    st.write('Please authenticate via Google Authenticator.')
    st.write('Go to the following URL and enter the authorization code:')
    st.write(authorization_url)
    authorization_code = st.text_input('Authorization code')

    # Exchange the authorization code for an access token
    try:
        token = oauth2_client.fetch_token(
            OAUTH2_ACCESS_TOKEN_URL,
            authorization_response={'code': authorization_code},
            client_id=CLIENT_ID,
            client_secret=CLIENT_SECRET,
        )

        # Store the access token and expiration time in session state
        st.session_state.access_token = token.get('access_token')
        st.session_state.expires_at = datetime.now() + timedelta(seconds=token.get('expires_in'))

        st.success('Login successful!')
        return True
    except Exception as e:
        st.error('Error: {}'.format(str(e)))
        return False

# Define the logout function
def logout():
    # Remove the access token and expiration time from session state
    if 'access_token' in st.session_state:
        del st.session_state.access_token
    if 'expires_at' in st.session_state:
        del st.session_state.expires_at

    st.success('Logout successful!')
    return True

# Define the app
def app():
    # Set up the menu
    menu = ['Home', 'Page 1', 'Page 2']
    choice = st.sidebar.selectbox('Select a page', menu)

    # Show the appropriate page based on the user's menu choice
    if choice == 'Home':
        st.write('Welcome to the home page!')

        # Check if the user is logged in
        if 'access_token' not in st.session_state or datetime.now() > st.session_state.expires_at:
            st.write('Please log in to continue.')
            if login():
                st.write('You are now logged in.')
        else:
            st.write('You are currently logged in.')
            if st.button('Log out'):
                logout()
                st.write('You are now logged out.')
    elif choice == 'Page 1':
        st.write('Welcome to page 1!')
    elif choice == 'Page 2':
        st.write('Welcome to page 2!')

if __name__ == '__main__':
    app()