import { b64Decode } from "encoding/base64";

export type PEM = string;

export type ExportedKey = JsonWebKey | PEM;

export interface KeyEncoder {
    encode(key: CryptoKey): Promise<JsonWebKey>;
    decode(
        key: ExportedKey,
        algorithm: RsaHashedImportParams | AesKeyAlgorithm,
        usages: KeyUsage[],
        extractable: boolean
    ): Promise<CryptoKey>;
}

type KeyType = "public" | "private";
type KeyFormat = "jwk" | "pem";

export function KeyEncoder(): KeyEncoder {
    const defaultExportFormat = "jwk";

    const PEM_CONSTANTS = Object.freeze({
        public: Object.freeze({
            header: "-----BEGIN PUBLIC KEY-----",
            footer: "-----END PUBLIC KEY-----"
        }),
        // NOTE: Will not handle "BEGIN RSA PRIVATE KEY"
        private: Object.freeze({
            header: "-----BEGIN PRIVATE KEY-----",
            footer: "-----END PRIVATE KEY-----"
        })
    });

    function getKeyFormat(exportedKey: ExportedKey, type: KeyType): KeyFormat {
        if (typeof exportedKey === "string") {
            if (type && exportedKey.startsWith(PEM_CONSTANTS[type].header)) {
                return "pem";
            }
        }
        return defaultExportFormat;
    }

    function getKeyType(exportedKey: ExportedKey): KeyType {
        if (typeof exportedKey === "string") {
            if (exportedKey.startsWith(PEM_CONSTANTS.public.header)) {
                return "public";
            } else if (exportedKey.startsWith(PEM_CONSTANTS.private.header)) {
                return "private";
            }
        }
        return null;
    }

    function preprocessPEM(key: PEM, type: KeyType): Uint8Array {
        const header = PEM_CONSTANTS[type].header;
        const footer = PEM_CONSTANTS[type].footer;

        const pemContents = key.substring(
            header.length,
            key.length - footer.length
        );
        return b64Decode(pemContents);
    }

    async function importPEM(
        key: PEM,
        algorithm: RsaHashedImportParams,
        usages: KeyUsage[],
        extractable: boolean,
        type: KeyType
    ): Promise<CryptoKey> {
        // This doesn't actually work for private keys in PEM format.
        // TODO: Fix
        return await crypto.subtle.importKey(
            "spki",
            preprocessPEM(key, type),
            algorithm,
            extractable,
            usages
        );
    }

    async function importJWK(
        key: JsonWebKey,
        algorithm: AesKeyAlgorithm | RsaHashedImportParams,
        usages: KeyUsage[],
        extractable: boolean
    ): Promise<CryptoKey> {
        return await crypto.subtle.importKey(
            "jwk",
            key,
            algorithm,
            extractable,
            usages
        );
    }

    return {
        async encode(key: CryptoKey): Promise<JsonWebKey> {
            return await crypto.subtle.exportKey(defaultExportFormat, key);
        },
        async decode(
            exportedKey: ExportedKey,
            algorithm: RsaHashedImportParams | AesKeyAlgorithm,
            usages: KeyUsage[],
            extractable = false
        ): Promise<CryptoKey> {
            const keyType = getKeyType(exportedKey);
            const keyFormat = getKeyFormat(exportedKey, keyType);

            if (keyFormat === "pem") {
                return await importPEM(
                    exportedKey as PEM,
                    algorithm as RsaHashedImportParams,
                    usages,
                    extractable,
                    keyType
                );
            } else {
                return await importJWK(
                    exportedKey as JsonWebKey,
                    algorithm,
                    usages,
                    extractable
                );
            }
        }
    };
}