import os
from typing import Callable, Any
import streamlit as st
from descope.descope_client import DescopeClient
from descope.exceptions import AuthException
from helpers.activities import restoreUserActivity
import utils as U

from dotenv import load_dotenv
load_dotenv()

DESCOPE_PROJECT_ID = os.environ.get("DESCOPE_PROJECT_ID")
descopeClient = DescopeClient(project_id=DESCOPE_PROJECT_ID)


def runWithAuth(func: Callable[[], Any]):
    if "token" not in st.session_state:
        if "code" in st.query_params:
            code = st.query_params["code"]
            st.query_params.clear()
            try:
                with st.spinner("Authenticating ..."):
                    jwtResponse = descopeClient.sso.exchange_token(code)
                st.session_state["token"] = jwtResponse["sessionToken"].get("jwt")
                st.session_state["refreshToken"] = jwtResponse["refreshSessionToken"].get(
                    "jwt"
                )
                st.session_state["user"] = jwtResponse["user"]
                with st.spinner("Restoring your chats ..."):
                    restoreUserActivity()
                st.rerun()
            except AuthException:
                st.error("Login failed!")

        st.warning("Login to Unlock the Magic ✨", icon=":material/login:")
        with st.container(border=False):
            if st.button(
                "Sign in with Google",
                type="primary",
                use_container_width=True
            ):
                oauthResponse = descopeClient.oauth.start(
                    provider="google", return_url=st.context.headers["Origin"]
                )
                url = oauthResponse["url"]
                # Redirect to Google
                st.markdown(
                    f'<meta http-equiv="refresh" content="0; url={url}">',
                    unsafe_allow_html=True,
                )
    else:
        try:
            with st.spinner("Verifying your identity ..."):
                jwtResponse = descopeClient.validate_and_refresh_session(
                    st.session_state.token, st.session_state.refreshToken
                )
                st.session_state["token"] = jwtResponse["sessionToken"].get("jwt")
                userEmail = st.session_state["user"].get("email")
                U.pprint(f"{userEmail} successfully authenticated!")

            func()
        except AuthException:
            # Log out user
            del st.session_state.token
            st.rerun()