import os
import streamlit as st
import requests
import msal
import urllib.parse
import base64
import hashlib
import secrets

# Configuration
APPLICATION_ID_KEY = os.getenv('APPLICATION_ID_KEY')
CLIENT_SECRET_KEY = os.getenv('CLIENT_SECRET_KEY')
AUTHORITY_URL = 'https://login.microsoftonline.com/common'
REDIRECT_URI = 'https://huggingface.co/spaces/awacke1/MSGraphAPI'

# Define product to scope mapping
PRODUCT_SCOPES = {
    "📧 Outlook": ['Mail.Read', 'Mail.Send', 'Calendars.ReadWrite'],
    "📒 OneNote": ['Notes.Read', 'Notes.Create'],
    "📊 Excel": ['Files.ReadWrite.All'],
    "📄 Word": ['Files.ReadWrite.All'],
    "🗃️ SharePoint": ['Sites.Read.All', 'Sites.ReadWrite.All'],
    "📅 Teams": ['Team.ReadBasic.All', 'Channel.ReadBasic.All'],
    "💬 Viva": ['Analytics.Read'],
    "🚀 Power Platform": ['Flow.Read.All'],
    "🧠 Copilot": ['Cognitive.Read'],
    "🗂️ OneDrive": ['Files.ReadWrite.All'],
    "💡 PowerPoint": ['Files.ReadWrite.All'],
    "📚 Microsoft Bookings": ['Bookings.Read.All', 'Bookings.ReadWrite.All'],
    "📓 Loop": ['Files.ReadWrite.All'],
    "🗣️ Translator": ['Translation.Read'],
    "📋 To Do & Planner": ['Tasks.ReadWrite'],
    "🔗 Azure OpenAI Service": ['AzureAIServices.ReadWrite.All']
}

# Base scopes (non-reserved)
BASE_SCOPES = ['User.Read']

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

def get_msal_app(code_verifier=None):
    return msal.PublicClientApplication(
        client_id=APPLICATION_ID_KEY,
        authority=AUTHORITY_URL,
        token_cache=msal.SerializableTokenCache()
    )

def get_access_token(code):
    client_instance = get_msal_app()
    
    try:
        result = client_instance.acquire_token_by_authorization_code(
            code=code,
            scopes=st.session_state.get('request_scopes', BASE_SCOPES),
            redirect_uri=REDIRECT_URI
        )
        
        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

def main():
    st.title("🦄 MS Graph API with AI & Cloud Integration for M365")

    st.sidebar.title("📝 M365 Products")
    st.sidebar.write("Select products to integrate:")
    
    selected_products = {}
    for product in PRODUCT_SCOPES.keys():
        selected = st.sidebar.checkbox(product)
        if selected:
            selected_products[product] = True

    # Dynamically build scopes based on selected products
    request_scopes = BASE_SCOPES.copy()
    for product in selected_products:
        request_scopes.extend(PRODUCT_SCOPES[product])
    request_scopes = list(set(request_scopes))  # Remove duplicates
    
    # Store request scopes in session state
    st.session_state['request_scopes'] = request_scopes

    if 'access_token' not in st.session_state:
        if 'code_verifier' not in st.session_state:
            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('=')
        
        client_instance = get_msal_app()
        auth_url = client_instance.get_authorization_request_url(
            scopes=request_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))
        
        query_params = st.query_params
        if 'code' in query_params:
            code = query_params.get('code')
            st.write('🔑 Authorization Code Obtained:', code[:10] + '...')
            
            try:
                access_token = get_access_token(code)
                st.session_state['access_token'] = access_token
                st.success("Access token acquired successfully!")
                st.rerun()
            except Exception as e:
                st.error(f"Error acquiring access token: {str(e)}")
                st.stop()
    else:
        access_token = st.session_state['access_token']
        
        user_info = get_user_info(access_token)
        if user_info:
            st.sidebar.write(f"👋 Hello, {user_info.get('displayName', 'User')}!")
        
        if selected_products:
            st.header("🧩 M365 Product Integrations")
            for product in selected_products:
                st.subheader(f"{product}")
                handle_product_integration(access_token, product)
        else:
            st.write("No products selected. Please select products from the sidebar.")

def get_user_info(access_token):
    headers = {'Authorization': f'Bearer {access_token}'}
    response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers)
    if response.status_code == 200:
        return response.json()
    else:
        st.error('Failed to fetch user info.')
        return None

def handle_product_integration(access_token, product):
    st.write(f"Integrating {product}...")
    # Implement product-specific API calls here

if __name__ == "__main__":
    main()