import { KMSClient } from "@aws-sdk/client-kms";
import { logger } from "@frak-backend/core";
import {
    Hex,
    hashMessage,
    hashTypedData,
    keccak256,
    serializeTransaction,
    signatureToHex,
} from "viem";
import { toAccount } from "viem/accounts";
import { getKmsAddress } from "./address";
import { getKmsSignature } from "./signature";

/**
 * Our KMS client
 */
const kmsClient = new KMSClient({
    region: process.env.AWS_REGION,
});

/**
 * Build a kms viem account
 * details: https://ethereum.stackexchange.com/a/73371/5093
 * @param string The key id of the KMS to use
 */
export const getKmsAccount = async ({ keyId }: { keyId: string }) => {
    // Get base data & methods for our accounts
    const address = await getKmsAddress({ keyId, client: kmsClient });
    const signMsg = async (msg: Hex) =>
        getKmsSignature({ keyId, address, msg, client: kmsClient });

    // Build the account
    return toAccount({
        address,
        async signTransaction(
            transaction,
            { serializer = serializeTransaction } = {}
        ) {
            logger.debug(
                { serialisedTransaction: serializer(transaction) },
                "Signing a new transaction"
            );
            const signature = await signMsg(keccak256(serializer(transaction)));
            return serializer(transaction, signature);
        },
        async signTypedData(typedData) {
            logger.debug({ typedData }, "Signing a new typed data");
            return signatureToHex(await signMsg(hashTypedData(typedData)));
        },
        async signMessage({ message }) {
            logger.debug({ message }, "Signing a new message");
            return signatureToHex(await signMsg(hashMessage(message)));
        },
    });
};