import 'dotenv/config';
import { AES_GCM_Decrypt, AES_GCM_Encrypt } from "./util/aes.js";
import { webcrypto } from "crypto";
import { Buffer } from "buffer";
import { sha512 } from "./util/sha2.js";
import { BSON } from "bson";

/**
 * Preshared Token Secret을 이용한 AES256-GCM 토큰
 * key, iv = SHA512(presharedSecret + nonce) 으로 생성
 */

export type SecureEncryptedToken = {
    nonce: string; // [key, iv] = SHA512(presharedSecret + nonce)
    encrypted: 1;
    type: string; // aad[0]
    validUntil: string; // aad[1]
    payload: string; //BASE64(AES(BSON(aad),BSON(payload)))
}

export type SecurePlaintextToken = {
    nonce: string; // [key, iv] = SHA512(presharedSecret + nonce)
    encrypted: 0;
    type: string; // aad[0]
    validUntil: string; // aad[1]
    payload: string; // JSON => aad[2]
    tag: string; //BASE64(AES(BSON(aad),null)))
}

export type SecureToken = SecureEncryptedToken | SecurePlaintextToken;

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

export async function generateSecureEncryptedToken<T extends object>(validUntil: Date, type: string, payload: T, presharedTokenSecret?: string): Promise<SecureEncryptedToken> {
    if (!presharedTokenSecret) {
        if (!staticPresharedTokenSecret) {
            throw new Error("PRESHARED_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');

    //secretBuffer에서 Buffer를 바로 준비해도 되지만, 혹시모를 안전상의 이유로 별도로 할당하고 복사
    const nonce = Buffer.from(webcrypto.getRandomValues(new Uint8Array(16)));
    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 validUntilText = validUntil.toISOString();

    const aad = [type, validUntilText];

    const aadRaw = BSON.serialize(aad);
    const payloadRaw = BSON.serialize(payload);

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

    return {
        nonce: nonce.toString('base64'),
        encrypted: 1,
        type,
        validUntil: validUntilText,
        payload: ciphertext.toString('base64'),
    }
}

export async function parseSecureEncryptedToken<T extends object>(secureToken: SecureEncryptedToken, presharedTokenSecret?: string): Promise<T> {
    const now = new Date();
    const validUntil = new Date(secureToken.validUntil);
    if (now > validUntil) {
        throw new Error("token expired");
    }

    if(!secureToken.encrypted) {
        throw new Error("token is not encrypted");
    }

    if (!presharedTokenSecret) {
        if (!staticPresharedTokenSecret) {
            throw new Error("PRESHARED_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(secureToken.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 aad = [secureToken.type, secureToken.validUntil];

    const aadRaw = BSON.serialize(aad);
    const payloadCiphertext = Buffer.from(secureToken.payload, 'base64');

    const payloadRaw = new Uint8Array(await AES_GCM_Decrypt(key, iv, payloadCiphertext, aadRaw));
    const payload = BSON.deserialize(payloadRaw);

    return payload as T;
}

export async function generateSecurePlaintextToken<T extends object>(validUntil: Date, type: string, payload: T, presharedTokenSecret?: string): Promise<SecurePlaintextToken> {
    if (!presharedTokenSecret) {
        if (!staticPresharedTokenSecret) {
            throw new Error("PRESHARED_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');

    //secretBuffer에서 Buffer를 바로 준비해도 되지만, 혹시모를 안전상의 이유로 별도로 할당하고 복사
    const nonce = Buffer.from(webcrypto.getRandomValues(new Uint8Array(16)));
    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 validUntilText = validUntil.toISOString();
    const jsonPayload = JSON.stringify(payload);

    const aad = [type, validUntilText, jsonPayload];
    const aadRaw = BSON.serialize(aad);

    const dummyPayload = new ArrayBuffer(0);

    const tag = Buffer.from(await AES_GCM_Encrypt(key, iv, dummyPayload, aadRaw));

    return {
        nonce: nonce.toString('base64'),
        encrypted: 0,
        type,
        validUntil: validUntilText,
        payload: jsonPayload,
        tag: tag.toString('base64'),
    }
}

export async function parseSecurePlaintextToken<T extends object>(secureToken: SecurePlaintextToken, presharedTokenSecret?: string): Promise<T> {
    if (!presharedTokenSecret) {
        if (!staticPresharedTokenSecret) {
            throw new Error("PRESHARED_TOKEN_SECRET is not set");
        }
        presharedTokenSecret = staticPresharedTokenSecret;
    }

    if (secureToken.encrypted) {
        throw new Error("token is encrypted");
    }

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

    const nonce = Buffer.from(secureToken.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 aad = [secureToken.type, secureToken.validUntil, secureToken.payload];
    const aadRaw = BSON.serialize(aad);

    const tag = Buffer.from(secureToken.tag, 'base64');

    await AES_GCM_Decrypt(key, iv, tag, aadRaw);

    return JSON.parse(secureToken.payload);
}

export async function generateSecureToken<T extends object>(encrypted: boolean, validUntil: Date, type: string, payload: T, presharedTokenSecret?: string): Promise<SecureToken> {
    if (encrypted) {
        return await generateSecureEncryptedToken(validUntil, type, payload, presharedTokenSecret);
    } else {
        return await generateSecurePlaintextToken(validUntil, type, payload, presharedTokenSecret);
    }
}

export async function parseSecureToken<T extends object>(secureToken: SecureToken, presharedTokenSecret?: string): Promise<T> {
    if (secureToken.encrypted) {
        return await parseSecureEncryptedToken(secureToken, presharedTokenSecret);
    } else {
        return await parseSecurePlaintextToken(secureToken, presharedTokenSecret);
    }
}