import streamlit as st
import requests
import os
import json
import urllib.parse
from dotenv import load_dotenv

# Load environment variables (if using a .env file)
load_dotenv()

# Hugging Face OAuth Credentials
CLIENT_ID = os.getenv("HUGGINGFACE_CLIENT_ID")
CLIENT_SECRET = os.getenv("HUGGINGFACE_CLIENT_SECRET")
REDIRECT_URI = os.getenv("HUGGINGFACE_REDIRECT_URI")

# Hugging Face OAuth URLs
AUTHORIZATION_URL = "https://huggingface.co/oauth/authorize"
TOKEN_URL = "https://huggingface.co/oauth/token"
USER_INFO_URL = "https://huggingface.co/api/whoami-v2"

# Streamlit Session State to Track User
if "user" not in st.session_state:
    st.session_state.user = None

def get_authorization_url():
    """Generate the Hugging Face authorization URL."""
    params = {
        "response_type": "code",
        "client_id": CLIENT_ID,
        "redirect_uri": REDIRECT_URI,
    }
    return f"{AUTHORIZATION_URL}?{urllib.parse.urlencode(params)}"

def exchange_code_for_token(auth_code):
    """Exchange authorization code for an access token."""
    data = {
        "grant_type": "authorization_code",
        "code": auth_code,
        "client_id": CLIENT_ID,
        "client_secret": CLIENT_SECRET,
        "redirect_uri": REDIRECT_URI,
    }
    response = requests.post(TOKEN_URL, data=data)
    return response.json()

def get_user_info(access_token):
    """Fetch user details from Hugging Face API."""
    headers = {"Authorization": f"Bearer {access_token}"}
    response = requests.get(USER_INFO_URL, headers=headers)
    return response.json()

def main():
    # Display App Title
    st.title("Sign In with Hugging Face")

    # Check if user is logged in
    if st.session_state.user:
        st.success(f"Logged in as **{st.session_state.user['name']}**")
        st.image(st.session_state.user['avatar'], width=100)
        if st.button("Logout"):
            st.session_state.user = None
            st.experimental_rerun()
    else:
        auth_url = get_authorization_url()
        st.markdown(f"[![Sign in with Hugging Face](https://huggingface.co/front/assets/huggingface_logo-noborder.svg)]({auth_url})")

    # Capture OAuth Redirect URL with Auth Code
    query_params = st.query_params
    if "code" in query_params:
        auth_code = query_params["code"][0]
        token_data = exchange_code_for_token(auth_code)

        if "access_token" in token_data:
            user_info = get_user_info(token_data["access_token"])
            st.session_state.user = {
                "name": user_info.get("name", "Unknown User"),
                "avatar": user_info.get("avatar", ""),
            }
            st.rerun()
        else:
            st.error("Failed to authenticate with Hugging Face.")

if __name__ == "__main__":
    main()