import { sha512 } from "@sammo/crypto/SHA2";
import { BSON } from "bson";
import type { PlainBson } from '@sammo/util';
import { calcBase64Len } from "@sammo/util";
import { AES_GCM_Decrypt, AES_GCM_Encrypt } from "@sammo/crypto/AES";
import { z } from "zod";

const nonceLen = 16;
const tagLen = 16;

export type EncryptedSimpleItem<T> = {
    nonce: string;
    ct: string;
}

export const zodEncryptedSimpleItem = z.object({
    nonce: z.string().length(calcBase64Len(nonceLen)),
    ct: z.string().min(calcBase64Len(tagLen)),
});

export type DecryptedSimpleItem<T extends EncryptedSimpleItem<unknown>> = T extends EncryptedSimpleItem<infer U> ? U : never;

const staticPresharedTokenSecret: string | undefined = process.env.PRESHARED_SECURE_TOKEN_SECRET;

export async function encryptSimpleItem<T extends PlainBson>(data: T, presharedTokenSecret?: string): Promise<EncryptedSimpleItem<T>> {
    if (!presharedTokenSecret) {
        if (!staticPresharedTokenSecret) {
            throw new Error("PRESHARED_SECURE_TOKEN_SECRET is not set");
        }
        presharedTokenSecret = staticPresharedTokenSecret;
    }

    const secretLength = Buffer.byteLength(presharedTokenSecret, 'utf8');
    const secretBuffer = Buffer.alloc(secretLength + nonceLen);
    secretBuffer.write(presharedTokenSecret, 0, secretLength, 'utf8');

    //secretBuffer에서 Buffer를 바로 준비해도 되지만, 혹시모를 안전상의 이유로 별도로 할당하고 복사
    const nonce = Buffer.from(crypto.getRandomValues(new Uint8Array(nonceLen)));
    secretBuffer.set(nonce, secretLength);

    const keyBuffer = Buffer.from(await sha512(secretBuffer));

    const key = new Uint8Array(keyBuffer.buffer, 0, 32);
    const iv = new Uint8Array(keyBuffer.buffer, 32, 12);

    const payload = BSON.serialize(data);

    const ciphertext = Buffer.from(await AES_GCM_Encrypt(key, iv, payload));

    return {
        nonce: nonce.toString('base64'),
        ct: ciphertext.toString('base64'),
    };
}

export async function decryptSimpleItem<T extends PlainBson>(data: EncryptedSimpleItem<T>, presharedTokenSecret?: string): Promise<T> {
    if (!presharedTokenSecret) {
        if (!staticPresharedTokenSecret) {
            throw new Error("PRESHARED_SECURE_TOKEN_SECRET is not set");
        }
        presharedTokenSecret = staticPresharedTokenSecret;
    }

    const secretLength = Buffer.byteLength(presharedTokenSecret, 'utf8');
    const secretBuffer = Buffer.alloc(secretLength + 16);
    secretBuffer.write(presharedTokenSecret, 0, secretLength, 'utf8');

    const nonce = Buffer.from(data.nonce, 'base64');
    secretBuffer.set(nonce, secretLength);

    const keyBuffer = Buffer.from(await sha512(secretBuffer));
    const key = new Uint8Array(keyBuffer.buffer, 0, 32);
    const iv = new Uint8Array(keyBuffer.buffer, 32, 12);

    const ct = Buffer.from(data.ct, 'base64');

    const payload = BSON.deserialize(Buffer.from(await AES_GCM_Decrypt(key, iv, ct))) as T satisfies DecryptedSimpleItem<EncryptedSimpleItem<T>>;
    return payload;
}