import { crypto } from "jsr:@std/crypto";
import { encodeBase64Url } from "jsr:@std/encoding/base64url";
import { decodeHex } from "jsr:@std/encoding/hex";

const {
    createECDH,
    getRandomValues,
} = await import("node:crypto");

// Reference:
// - https://docs.deno.com/api/node/crypto/~/ECDH
// - https://developer.mozilla.org/en-US/docs/Web/API/SubtleCrypto/encrypt

class EcdhEphemeral {
    ePublicKey: Uint8Array;
    sharedSecret: ArrayBufferLike;

    constructor(ePublicKey: Uint8Array, sharedSecret: ArrayBufferLike) {
        this.ePublicKey = ePublicKey;
        this.sharedSecret = sharedSecret;
    }
}

class Ecdh {
    publicKeyPoint: Uint8Array;

    constructor(publicKeyPoint: Uint8Array) {
        this.publicKeyPoint = publicKeyPoint;
    }

    static fromHex(publicKeyPointHex: string): Ecdh {
        return new Ecdh(decodeHex(publicKeyPointHex));
    }

    ecdh(): EcdhEphemeral {
        const bitLength = ((this.publicKeyPoint.length - 1) / 2) * 8;
        if (bitLength !== 256) {
            // NOTICE current only supports P256
            throw `Invalid EC public key point, bit length: ${bitLength}`;
        }
        const temp = createECDH(`secp${bitLength}r1`);
        const ePublicKey = temp.generateKeys();
        const sharedSecret = temp.computeSecret(this.publicKeyPoint);
        return new EcdhEphemeral(ePublicKey, sharedSecret);
    }
}

async function simpleKdf256(
    input: ArrayBufferLike,
): Promise<ArrayBuffer> {
    const SHA256 = "SHA-256";
    for (let i = 0; i < 8; i++) {
        input = await crypto.subtle.digest(SHA256, input);
    }
    return input;
}

function randomNonce(): ArrayBufferLike {
    const buffer = new Uint8Array(12);
    getRandomValues(buffer);
    return buffer;
}

class WrapKeyHeader {
    kid: string;
    enc: string;
    ePubKey: string;

    constructor(kid: string, enc: string, ePubKey: string) {
        this.kid = kid;
        this.enc = enc;
        this.ePubKey = ePubKey;
    }
}

class WrapKey {
    header: WrapKeyHeader;
    nonce: ArrayBuffer;
    encryptedData: ArrayBuffer;

    constructor(
        header: WrapKeyHeader,
        nonce: ArrayBuffer,
        encryptedData: ArrayBuffer,
    ) {
        this.header = header;
        this.nonce = nonce;
        this.encryptedData = encryptedData;
    }

    toString(): string {
        const headerJson = JSON.stringify(this.header);
        return `WK:${encodeBase64Url(new TextEncoder().encode(headerJson))}.${
            encodeBase64Url(this.nonce)
        }.${encodeBase64Url(this.encryptedData)}`;
    }
}

function concatUint8Array(a: Uint8Array, b: Uint8Array): Uint8Array {
    const merge = new Uint8Array(a.byteLength + b.byteLength);
    merge.set(a, 0);
    merge.set(b, a.byteLength);
    return merge;
}

export async function encryptEcdhP256(
    kid: string,
    publicKeyPointHex: string,
    data: ArrayBuffer,
) {
    const ENC_AES256_GCM_P256 = "aes256-gcm-p256";
    const ecdh = Ecdh.fromHex(publicKeyPointHex);
    const ecdhEphemeral = ecdh.ecdh();
    const key = await simpleKdf256(ecdhEphemeral.sharedSecret);
    const nonce = randomNonce();

    const derPrefix = decodeHex(
        "3059301306072a8648ce3d020106082a8648ce3d030107034200",
    );
    const der = concatUint8Array(derPrefix, ecdhEphemeral.ePublicKey);
    const wrapkeyHeader = new WrapKeyHeader(
        kid,
        ENC_AES256_GCM_P256,
        encodeBase64Url(der),
    );
    const cryptoKey = await crypto.subtle.importKey(
        "raw",
        key,
        "AES-GCM",
        false,
        ["encrypt"],
    );
    const encryptedData = await crypto.subtle.encrypt(
        { name: "AES-GCM", iv: nonce },
        cryptoKey,
        data,
    );
    return new WrapKey(wrapkeyHeader, nonce, encryptedData);
}