# pages/oauth_callback.py - COMPLETE FIXED VERSION
import streamlit as st
import requests
import os
from datetime import datetime

st.set_page_config(page_title="Authenticating...", page_icon="⏳", layout="centered")

# Get credentials from environment
GOOGLE_CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
GOOGLE_CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")

# Show loading screen
st.markdown("""
<style>
    .spinner {
        border: 4px solid #f3f3f3;
        border-top: 4px solid #4285F4;
        border-radius: 50%;
        width: 40px;
        height: 40px;
        animation: spin 2s linear infinite;
        margin: 20px auto;
    }
    @keyframes spin {
        0% { transform: rotate(0deg); }
        100% { transform: rotate(360deg); }
    }
    .success-box {
        background: linear-gradient(135deg, #10b981, #34d399);
        color: white;
        padding: 2rem;
        border-radius: 10px;
        text-align: center;
        margin: 2rem 0;
    }
</style>
""", unsafe_allow_html=True)

# Initial loading message
st.markdown("""
<div style="text-align: center; padding: 50px;">
    <h1>🔐 Authenticating...</h1>
    <p>Please wait while we log you in with Google.</p>
    <div class="spinner"></div>
</div>
""", unsafe_allow_html=True)

# Process the OAuth callback
try:
    # Get the authorization code from URL
    # Try new method first, fallback to old
    try:
        query_params = st.query_params
        is_new_api = True
    except AttributeError:
        query_params = st.experimental_get_query_params()
        is_new_api = False
    
    if "code" in query_params and GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET:
        # Extract code (different format in old vs new)
        if is_new_api:
            code = query_params.get("code", "")
        else:
            code_list = query_params.get("code", [])
            code = code_list[0] if code_list else ""
        
        if not code:
            st.error("❌ No authorization code received")
            st.stop()
        
        # Exchange code for token
        with st.spinner("Verifying credentials..."):
            token_url = "https://oauth2.googleapis.com/token"
            data = {
                "client_id": GOOGLE_CLIENT_ID,
                "client_secret": GOOGLE_CLIENT_SECRET,
                "code": code,
                "redirect_uri": "https://refat81-social-media-data-extractor-chatbot.hf.space/oauth_callback",
                "grant_type": "authorization_code"
            }
            
            response = requests.post(token_url, data=data)
            tokens = response.json()
        
        if "access_token" in tokens:
            # Get user info
            with st.spinner("Loading your profile..."):
                userinfo_url = "https://www.googleapis.com/oauth2/v3/userinfo"
                headers = {"Authorization": f"Bearer {tokens['access_token']}"}
                user_response = requests.get(userinfo_url, headers=headers)
                user_info = user_response.json()
            
            # Store in session
            st.session_state.user_info = user_info
            st.session_state.authenticated = True
            st.session_state.login_time = datetime.now()
            
            # Clear the loading message and show success
            st.empty()  # Clear loading message
            
            st.markdown(f"""
            <div class="success-box">
                <h2 style="margin:0;">✅ Login Successful!</h2>
                <p style="font-size: 1.2rem; margin: 1rem 0;">Welcome, <strong>{user_info.get('name', 'User')}</strong>!</p>
                <p>Redirecting to dashboard...</p>
            </div>
            """, unsafe_allow_html=True)
            
            # Auto-redirect to main page
            st.markdown("""
            <script>
                setTimeout(function() {
                    window.location.href = "/";
                }, 2000);
            </script>
            
            <div style="text-align: center; margin-top: 2rem;">
                <p>If you are not redirected automatically, <a href="/">click here</a>.</p>
            </div>
            """, unsafe_allow_html=True)
            
        else:
            st.error("❌ Authentication failed. Please try again.")
            if "error" in tokens:
                st.write(f"Error: {tokens['error']}")
            
    elif not GOOGLE_CLIENT_ID or not GOOGLE_CLIENT_SECRET:
        st.error("""
        ❌ Server configuration error.
        
        Please contact administrator or check HuggingFace Space secrets.
        """)
        
    else:
        st.error("❌ No authorization code received. Please try logging in again.")
        
except Exception as e:
    st.error(f"❌ Error: {str(e)}")
    st.write("Please try logging in again from the main page.")

# Footer
st.markdown("---")
st.caption("Social Media Data Extractor • Secure OAuth Authentication")