//Constants we need
const APP_BASE_URL = '<APP_BASE_URL>';

const AWS_COGNITO_POOL_CLIENT_ID = '<COGNITO_POOL_CLIENT_ID>';
const AWS_COGNITO_POOL_URI = '<COGNITO_POOL_URI>';
//The following must be registered with the Cognito app client
const AWS_COGNITO_LOGIN_CALLBACK_URI = APP_BASE_URL + '/?action=login';
const AWS_COGNITO_LOGOUT_CALLBACK_URI = APP_BASE_URL + '/?action=logout';

//API endpoints
const AWS_COGNITO_AUTHORIZE_ENDPOINT = AWS_COGNITO_POOL_URI + '/oauth2/authorize';
const AWS_COGNITO_CODE_EXCHANGE_ENDPOINT = AWS_COGNITO_POOL_URI + '/oauth2/token';
const AWS_COGNITO_USER_INFO_ENDPOINT = AWS_COGNITO_POOL_URI + '/oauth2/userInfo';
const AWS_COGNITO_LOGOUT_ENDPOINT = AWS_COGNITO_POOL_URI + '/logout?client_id=' + AWS_COGNITO_POOL_CLIENT_ID + '&logout_uri=' + AWS_COGNITO_LOGOUT_CALLBACK_URI;

//Somewhere to store auth tokens, we don't use cookies or local storage because we have no way of knowing if the user has selected 'remember me' or not, so we only 'remember' for the length of a session
const sessionStorage = window.sessionStorage;

//Setup the UI
document.getElementById('hrefLogin').onclick = loginClickFunction;
document.getElementById('hrefLogout').onclick = logoutClickFunction;
document.getElementById('hrefLogout').style.display = 'none';

//Parse the URL params
const urlParams = new URL(window.location.href).searchParams;

const action = urlParams.get('action');

if (action === 'login') {
    if (sessionStorage.getItem('cognito_login_state') && sessionStorage.getItem('cognito_login_state') === urlParams.get('state') && sessionStorage.getItem('cognito_pkce_code_verifier')) {
        //Clear the state variable, it shouldn't be reused
        sessionStorage.removeItem('cognito_login_state');

        //Exchange the auth token for an access token, get the current user
        const authCode = urlParams.get('code');

        exchangeAuthToken(authCode, sessionStorage.getItem('cognito_pkce_code_verifier'), tokenSuccessCallback, errorCallback);
    }
}
else if (action === 'logout') {
    clearAuthData();

    //Redirect to base URL
    window.location.href = APP_BASE_URL;
}
else if (sessionStorage.getItem('cognito_expiration') && sessionStorage.getItem('cognito_access_token')) {
    if (Math.floor(Date.now() / 1000) < sessionStorage.getItem('cognito_expiration')) {
        //We have a current token
        getCurrentUser(sessionStorage.getItem('cognito_access_token'), getUserSuccessCallback, errorCallback);
    }
    else {
        //Token is expired, clean up now defunct auth data
        clearAuthData();
    }
}

//Setup async event callbacks
function tokenSuccessCallback(tokenData) {
    //Save the data
    sessionStorage.setItem('cognito_access_token', tokenData['access_token']);
    sessionStorage.setItem('cognito_id_token', tokenData['id_token']);
    sessionStorage.setItem('cognito_refresh_token', tokenData['refresh_token']);
    sessionStorage.setItem('cognito_expiration', Math.floor(Date.now() / 1000) + tokenData['expires_in']);

    //Don't reuse PKCE token
    sessionStorage.removeItem('cognito_pkce_code_verifier');

    //Redirect to base URL, the user info request will be made there
    window.location.href = APP_BASE_URL;
}

function getUserSuccessCallback(userData) {
    //Hide the login link
    document.getElementById('hrefLogin').style.display = 'none';

    //Show the logout link
    document.getElementById('hrefLogout').style.display = 'inline';

    console.log(userData);
}

function errorCallback(errorData) {
    console.log(errorData);
}

function loginClickFunction() {
    //Generate a 'state' variable we can use to make sure we're returning from an authentication flow we started, this adds no security, it's purely for flow validation
    const stateParameter = generateSecret();

    //Store the state so we can verify it later
    sessionStorage.setItem('cognito_login_state', stateParameter);

    //Generate a PKCE code verifier, again, since this is a public client, it provides no security, use a hash because it's easy, and satisifes length requirements
    //From RFC 7636, the code verifier is base64url-encoded ASCII
    const codeVerifier = base64UrlEncode(generateSecret());
    sessionStorage.setItem('cognito_pkce_code_verifier', codeVerifier);

    //The challenge, by definition, is base64 encoded SHA256 hash of the code verifier
    const codeChallenge = hexToBase64UrlEncode(sha256(codeVerifier));

    //Build the request query
    const loginRequestPairs = [];

    loginRequestPairs.push(encodeURIComponent('client_id') + '=' + encodeURIComponent(AWS_COGNITO_POOL_CLIENT_ID));
    loginRequestPairs.push(encodeURIComponent('redirect_uri') + '=' + encodeURIComponent(AWS_COGNITO_LOGIN_CALLBACK_URI));
    loginRequestPairs.push(encodeURIComponent('response_type') + '=' + encodeURIComponent('code'));
    loginRequestPairs.push(encodeURIComponent('code_challenge') + '=' + encodeURIComponent(codeChallenge));
    loginRequestPairs.push(encodeURIComponent('code_challenge_method') + '=' + encodeURIComponent('S256'));
    loginRequestPairs.push(encodeURIComponent('state') + '=' + encodeURIComponent(stateParameter));

    const loginRequestQuery = loginRequestPairs.join('&').replace(/%20/g, '+');

    //Redirect to login
    window.location.href = AWS_COGNITO_AUTHORIZE_ENDPOINT + '?' + loginRequestQuery;
    return false;
}

function logoutClickFunction() {
    //Redirect to logout
    window.location.href = AWS_COGNITO_LOGOUT_ENDPOINT;
    return false;
}

function exchangeAuthToken(authorizationCode, codeVerifier, successCallback, errorCallback) {
    //Exchanges an authorization code for a set of user tokens
    //https://docs.aws.amazon.com/cognito/latest/developerguide/token-endpoint.html
    const tokenRequest = new XMLHttpRequest();

    //Build the form data
    const tokenRequestPairs = [];

    tokenRequestPairs.push(encodeURIComponent('grant_type') + '=' + encodeURIComponent('authorization_code'));
    tokenRequestPairs.push(encodeURIComponent('client_id') + '=' + encodeURIComponent(AWS_COGNITO_POOL_CLIENT_ID));
    tokenRequestPairs.push(encodeURIComponent('redirect_uri') + '=' + encodeURIComponent(AWS_COGNITO_LOGIN_CALLBACK_URI));
    tokenRequestPairs.push(encodeURIComponent('code_verifier') + '=' + encodeURIComponent(codeVerifier));
    tokenRequestPairs.push(encodeURIComponent('code') + '=' + encodeURIComponent(authorizationCode));

    const tokenRequestData = tokenRequestPairs.join('&').replace(/%20/g, '+');

    //Set up event handler
    tokenRequest.onreadystatechange = function() {
        if (tokenRequest.readyState == 4) {
            if (tokenRequest.status == 200) {
                successCallback(JSON.parse(tokenRequest.responseText));
            }
            else {
                errorCallback(JSON.parse(tokenRequest.responseText));
            }
        }
    };

    //Submit the form
    tokenRequest.open('POST', AWS_COGNITO_CODE_EXCHANGE_ENDPOINT, true);
    tokenRequest.setRequestHeader('Content-Type', 'application/x-www-form-urlencoded');
    tokenRequest.send(tokenRequestData);
}

function getCurrentUser(accessToken, successCallback, errorCallback) {
    //Gets the user info
    //https://docs.aws.amazon.com/cognito/latest/developerguide/userinfo-endpoint.html
    const userInfoRequest = new XMLHttpRequest();

    //Setup handler
    userInfoRequest.onreadystatechange = function() {
        if (userInfoRequest.readyState == 4) {
            if (userInfoRequest.status == 200) {
                successCallback(JSON.parse(userInfoRequest.responseText));
            }
            else {
                errorCallback(JSON.parse(tokenRequest.responseText));
            }
        }
    };

    userInfoRequest.open('GET', AWS_COGNITO_USER_INFO_ENDPOINT, true);
    userInfoRequest.setRequestHeader('Authorization', 'Bearer ' + accessToken);
    userInfoRequest.send(null);
}

function clearAuthData() {
    sessionStorage.removeItem('cognito_access_token');
    sessionStorage.removeItem('cognito_id_token');
    sessionStorage.removeItem('cognito_refresh_token');
    sessionStorage.removeItem('cognito_expiration');
    sessionStorage.removeItem('cognito_login_state');
    sessionStorage.removeItem('cognito_pkce_code_verifier');
}

function base64UrlEncode(toEncode) {
    var encodedString = btoa(toEncode);

    encodedString = encodedString.replace(/=+$/, '');
    encodedString = encodedString.replace(/\+/g, '-');
    encodedString = encodedString.replace(/\//g, '_');

    return encodedString;
}

function hexToBase64UrlEncode(hexString) {
    //https://stackoverflow.com/a/30614192
    hexString = hexString.replace(/\r|\n/g, '');
    hexString = hexString.replace(/([\da-fA-F]{2}) ?/g, '0x$1 ');
    hexString = hexString.replace(/ +$/, '');

    const hexArray = hexString.split(' ');
    const byteString = String.fromCharCode.apply(null, hexArray);

    return base64UrlEncode(byteString);
}

function generateSecret() {
    //Generates a cryptographically secure "secret" with 256 bits of entropy, returned as a hex string
    return sha256(window.crypto.getRandomValues(new Uint32Array(8)));
}

<!DOCTYPE html>
<html>
	<head>
		<meta charset="UTF-8">

		<title>Auth Test</title>
	</head>

	<body>
		<p>This is the auth test. <a id="hrefLogin" href="#">Login</a><a id="hrefLogout" href="#">Logout</a></p>
	</body>

	<script src="https://cdnjs.cloudflare.com/ajax/libs/js-sha256/0.9.0/sha256.js"></script>
	<script src="/cognito.js"></script>
</html>

<!doctype html>
<html>
	<head>
		<meta charset="UTF-8">

		<title>Auth Test - Privacy Policy</title>
	</head>

	<body>
		<p>This is purely for testing purposes. No privacy or security should be expected. By signing up, you acknowledge that your data may become public.</p>
	</body>
</html>