import os
import streamlit as st
import requests
import msal
import base64
import hashlib
import secrets

# 🤓 Load environment variables (Ensure these are set!)
APPLICATION_ID_KEY = os.getenv('APPLICATION_ID_KEY')
CLIENT_SECRET_KEY = os.getenv('CLIENT_SECRET_KEY')
AUTHORITY_URL = 'https://login.microsoftonline.com/common'  # Use 'common' for multi-tenant apps
REDIRECT_URI = 'https://huggingface.co/spaces/awacke1/MSGraphAPI'  # Update this to match your app's redirect URI

# 🎯 Define the scopes your app will need
SCOPES = ['User.Read']

# New function to generate PKCE code verifier and challenge
def generate_pkce_codes():
    code_verifier = secrets.token_urlsafe(128)[:128]
    code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).decode().rstrip('=')
    return code_verifier, code_challenge

# 🛠️ Initialize the MSAL client for Public Client App (PKCE support)
def get_msal_app():
    return msal.PublicClientApplication(
        client_id=APPLICATION_ID_KEY,
        authority=AUTHORITY_URL
    )

# 🔐 Acquire access token using authorization code
def get_access_token(code, code_verifier=None):
    client_instance = get_msal_app()
    
    st.write("Debug: MSAL App Configuration:")
    st.write(f"Client ID: {APPLICATION_ID_KEY[:5]}...")
    st.write(f"Authority: {AUTHORITY_URL}")
    st.write(f"Redirect URI: {REDIRECT_URI}")
    
    try:
        # Attempt to acquire token, use PKCE code_verifier if provided
        result = client_instance.acquire_token_by_authorization_code(
            code=code,
            scopes=SCOPES,
            redirect_uri=REDIRECT_URI,
            code_verifier=code_verifier  # Include only if PKCE is enabled
        )
        
        if 'access_token' in result:
            return result['access_token']
        else:
            error_description = result.get('error_description', 'No error description provided')
            raise Exception(f"Error acquiring token: {error_description}")
    except Exception as e:
        st.error(f"Exception in get_access_token: {str(e)}")
        raise

# 🏃‍♂️ Main function to process the query parameters and handle the token exchange
def process_query_params():
    try:
        query_params = st.experimental_get_query_params()
        st.write("Debug: All query parameters:", query_params)

        if 'error' in query_params:
            error = query_params.get('error')
            error_description = query_params.get('error_description', 'No description provided')
            st.error(f"Authentication Error: {error}")
            st.error(f"Error Description: {error_description}")
            st.stop()

        if 'code' in query_params:
            code = query_params.get('code')[0]  # MS Graph returns the code as a list
            st.write('🔑 Authorization Code Obtained:', code[:10] + '...')
            
            try:
                # Retrieve code_verifier from session state
                code_verifier = st.session_state.get('code_verifier')
                if not code_verifier:
                    st.error("Code verifier not found in session state.")
                    st.stop()
                
                # Acquire the access token
                access_token = get_access_token(code, code_verifier)
                st.session_state['access_token'] = access_token
                st.success("Access token acquired successfully!")
                
                # Clear the query parameters from the URL
                st.experimental_set_query_params()
                st.experimental_rerun()
            except Exception as e:
                st.error(f"Error acquiring access token: {str(e)}")
                st.stop()
        else:
            st.warning("No authorization code found in the query parameters.")
    except Exception as e:
        st.error(f"Error processing query parameters: {str(e)}")
        st.stop()

# Main application function
def main():
    st.title("🦄 MS Graph API with AI & Cloud Integration with M365")

    if 'access_token' not in st.session_state:
        if 'code_verifier' not in st.session_state:
            # Generate PKCE codes and store code_verifier in session
            code_verifier, code_challenge = generate_pkce_codes()
            st.session_state['code_verifier'] = code_verifier
        else:
            code_verifier = st.session_state['code_verifier']
            code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).decode().rstrip('=')

        # Get MSAL app and construct the authorization URL
        client_instance = get_msal_app()
        auth_url = client_instance.get_authorization_request_url(
            scopes=SCOPES,
            redirect_uri=REDIRECT_URI,
            code_challenge=code_challenge,
            code_challenge_method="S256"
        )
        st.write('👋 Please [click here]({}) to log in and authorize the app.'.format(auth_url))
        st.stop()

    # Process query parameters to acquire token
    process_query_params()

    # If access token is present, greet the user
    if 'access_token' in st.session_state:
        access_token = st.session_state['access_token']
        headers = {'Authorization': 'Bearer ' + access_token}

        response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers)
        if response.status_code == 200:
            user_info = response.json()
            st.write(f"👋 Hello, {user_info['displayName']}!")
        else:
            st.error("Failed to fetch user info.")
            st.write(response.text)

# 🚀 Run the main function
if __name__ == "__main__":
    main()