/*
Copyright 2023 The Matrix.org Foundation C.I.C.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

import { type IdTokenClaims, OidcClient, WebStorageStateStore, ErrorResponse } from "oidc-client-ts";

import { type AccessTokens, TokenRefreshLogoutError } from "../http-api/index.ts";
import { generateScope } from "./authorize.ts";
import { discoverAndValidateOIDCIssuerWellKnown } from "./discovery.ts";
import { logger } from "../logger.ts";

/**
 * @experimental
 * Class responsible for refreshing OIDC access tokens
 *
 * Client implementations will likely want to override {@link persistTokens} to persist tokens after successful refresh
 *
 */
export class OidcTokenRefresher {
    /**
     * This is now just a resolved promise and will be removed in a future version.
     * Initialisation is done lazily at token refresh time.
     * @deprecated Consumers no longer need to wait for this promise.
     */
    public readonly oidcClientReady!: Promise<void>;

    // If there is a initialisation attempt in progress, we keep track of it here.
    private initPromise?: Promise<void>;

    private oidcClient!: OidcClient;
    private inflightRefreshRequest?: Promise<AccessTokens>;

    public constructor(
        /**
         * The OIDC issuer as returned by the /auth_issuer API
         */
        private issuer: string,
        /**
         * id of this client as registered with the OP
         */
        private clientId: string,
        /**
         * redirectUri as registered with OP
         */
        private redirectUri: string,
        /**
         * Device ID of current session
         */
        protected deviceId: string,
        /**
         * idTokenClaims as returned from authorization grant
         * used to validate tokens
         */
        private readonly idTokenClaims: IdTokenClaims,
    ) {
        this.oidcClientReady = Promise.resolve();
    }

    /**
     * Ensures that the client is initialised.
     * @returns Promise that resolves when initialisation is complete
     * @throws if initialisation fails
     */
    private async ensureInit(): Promise<void> {
        if (!this.oidcClient) {
            if (this.initPromise) {
                return this.initPromise;
            }

            this.initPromise = this.initialiseOidcClient(this.issuer, this.clientId, this.deviceId, this.redirectUri);
            try {
                await this.initPromise;
            } finally {
                this.initPromise = undefined;
            }
        }
    }

    private async initialiseOidcClient(
        issuer: string,
        clientId: string,
        deviceId: string,
        redirectUri: string,
    ): Promise<void> {
        try {
            const config = await discoverAndValidateOIDCIssuerWellKnown(issuer);

            const scope = generateScope(deviceId);

            this.oidcClient = new OidcClient({
                metadata: config,
                signingKeys: config.signingKeys ?? undefined,
                client_id: clientId,
                scope,
                redirect_uri: redirectUri,
                authority: config.issuer,
                stateStore: new WebStorageStateStore({ prefix: "mx_oidc_", store: window.sessionStorage }),
            });
        } catch (error) {
            logger.error("Failed to initialise OIDC client.", error);
            throw new Error("Failed to initialise OIDC client.");
        }
    }

    /**
     * Attempt token refresh using given refresh token
     * @param refreshToken - refresh token to use in request with token issuer
     * @returns tokens - Promise that resolves with new access and refresh tokens
     * @throws when token refresh fails
     */
    public async doRefreshAccessToken(refreshToken: string): Promise<AccessTokens> {
        await this.ensureInit();

        if (!this.inflightRefreshRequest) {
            this.inflightRefreshRequest = this.getNewTokens(refreshToken);
        }
        try {
            const tokens = await this.inflightRefreshRequest;
            return tokens;
        } catch (e) {
            // If we encounter an OIDC error then signal that it should cause a logout by upgrading it to a TokenRefreshLogoutError
            if (e instanceof ErrorResponse) {
                throw new TokenRefreshLogoutError(e);
            }
            throw e;
        } finally {
            this.inflightRefreshRequest = undefined;
        }
    }

    /**
     * Persist the new tokens, called after tokens are successfully refreshed.
     *
     * This function is intended to be overriden by the consumer when persistence is necessary.
     *
     * @param tokens.accessToken - new access token
     * @param tokens.refreshToken - OPTIONAL new refresh token
     */
    protected async persistTokens(tokens: { accessToken: string; refreshToken?: string }): Promise<void> {
        // NOOP
    }

    private async getNewTokens(refreshToken: string): Promise<AccessTokens> {
        if (!this.oidcClient) {
            throw new Error("Cannot get new token before OIDC client is initialised.");
        }

        const refreshTokenState = {
            refresh_token: refreshToken,
            session_state: "test",
            data: undefined,
            profile: this.idTokenClaims,
        };

        const requestStart = Date.now();
        const response = await this.oidcClient.useRefreshToken({
            state: refreshTokenState,
            timeoutInSeconds: 300,
        });

        const tokens = {
            accessToken: response.access_token,
            refreshToken: response.refresh_token,
            // We use the request start time to calculate the expiry time as we don't know when the server received our request
            expiry: response.expires_in ? new Date(requestStart + response.expires_in * 1000) : undefined,
        } satisfies AccessTokens;

        await this.persistTokens(tokens);

        return tokens;
    }
}