import streamlit as st
import requests
import logging
import jwt

from urllib.parse import urlencode

logging.basicConfig(
    level=logging.DEBUG,
    )
logger = logging.getLogger("python")



def get_discovery_keys(token_response):
    tenant = st.secrets["auth"]["microsoft"]["tenant"]
    url = f"https://login.microsoftonline.com/{tenant}/discovery/v2.0/keys"
    try:
        logger.info("Getting discovery keys...")
        response = requests.get(url=url)
        response.raise_for_status()
        disc_keys = response.json()
        logger.info(f"Discovery keys: {disc_keys}")
        if "id_token" in token_response:
            token = token_response["id_token"]
            decoded_payload = jwt.decode(token, options={"verify_signature": False})
            decoded_header = jwt.get_unverified_header(token)

            logger.info(f"ID Token Header:{str(decoded_header)}")
            logger.info(f"ID Token Payload:{str(decoded_payload)}")
        else:
            logger.error("No ID Token in the token_response.")

        st.write("Logged with manual set up.")
        logger.info("Manual Log in logs end.")
        st.query_params.clear()
    except requests.exceptions.RequestException as error:
        logger.error(f"Error on gettin the discovery keys: {str(error)}")

    st.stop()

def get_token_from_code(auth_code):
    tenant = st.secrets["auth"]["microsoft"]["tenant"]
    client_id = st.secrets["auth"]["microsoft"]["client_id"]
    client_secret = st.secrets["auth"]["microsoft"]["client_secret"] 
    redirect_uri = st.secrets["auth"]["microsoft"]["redirect_uri"]

    token_url = f"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
    headers = {
        'Content-Type': 'application/x-www-form-urlencoded'
    }
    body = {
        'grant_type': 'authorization_code',
        'response_type': 'id_token',
        'code': auth_code,
        'redirect_uri': redirect_uri,
        'client_id': client_id,
        'client_secret': client_secret,
        'scope': 'openid profile email'
    }

    try:
        logger.info("Manual Log in logs start...")
        response = requests.post(token_url, headers=headers, data=body)
        response.raise_for_status()
        token_response = response.json()
        logger.info(f"Response token: {token_response}")
        logger.info(f"Context Headers: {st.context.headers.to_dict()}")
        logger.info(f"Context Cookes: {st.context.cookies.to_dict()}")
        logger.info(f"Experimental user: {st.experimental_user.to_dict()}")
        logger.info(f"Manual parameters: {st.query_params.to_dict()}")
        get_discovery_keys(token_response)
    except requests.exceptions.RequestException as error:
        logger.error(f"Error on gettin the token from code: {str(error)}")



def manual_login():
    tenant = st.secrets["auth"]["microsoft"]["tenant"]
    params = {
        "client_id": st.secrets["auth"]["microsoft"]["client_id"],
        "response_type": "code",
        "scope": "openid profile email",
        "state": "12345",
        "nonce": "678910",
        'redirect_uri': st.secrets["auth"]["microsoft"]["redirect_uri"],
        "prompt": "select_account",
    }

    url = f"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/authorize?{urlencode(params)}"

    url_link = st.button("Manual Login",type="secondary")
    if url_link:
        st.markdown(
            f'<meta http-equiv="refresh" content="0; url={url}">',
            unsafe_allow_html=True
        )

    if "code" in st.query_params:
        code = st.query_params["code"]
        logger.info("Getting token from code...")
        get_token_from_code(code)
    elif "error" in st.query_params:
        logger.error(f'Failed to get authorization code: {st.query_params.get("error")}')



def streamlit_login():

    if st.button("Streamlit Login"):
        tenant = st.secrets["auth"]["microsoft"]["tenant"]
        url = f"https://login.microsoftonline.com/{tenant}/discovery/v2.0/keys"
        try:
            logger.info("Warming discovery keys...")
            response = requests.get(url=url)
            response.raise_for_status()
        except requests.exceptions.RequestException as error:
            logger.error(f"Error on warming the discovery keys: {str(error)}")

        st.login("microsoft")

        logger.info("Streamltit Log in logs without 'is_logged_in' start...")
        logger.info(f"Context Headers: {st.context.headers.to_dict()}")
        logger.info(f"Context Cookes: {st.context.cookies.to_dict()}")
        logger.info(f"Experimental user: {st.experimental_user.to_dict()}")
        logger.info(f"Streamlit parameters: {st.query_params.to_dict()}")
        logger.info("Streamltit Log in logs without 'is_logged_in' end.")
    
    if st.experimental_user.is_logged_in:
        st.write("Logged in with Streamlit Log In.")
        logger.info("Streamltit Log in logs with 'is_logged_in' start...")
        logger.info(f"Context Headers: {st.context.headers.to_dict()}")
        logger.info(f"Context Cookes: {st.context.cookies.to_dict()}")
        logger.info(f"Experimental user: {st.experimental_user.to_dict()}")
        logger.info(f"Streamlit parameters: {st.query_params.to_dict()}")
        logger.info("Streamltit Log in logs with 'is_logged_in' end.")

    if st.button("clear cookies"):
        st.logout()

    st.stop()


def main():
    pg = st.navigation(
        {"Home":[
                st.Page(
                streamlit_login,
                title="Streamlit Login",
                ),
                st.Page(
                    manual_login,
                    title="Manual Authorization Setup"
                )
            ]
        }
    )
    pg.run()