import { KMSClient, SignCommand } from "@aws-sdk/client-kms";
import { Integer, Sequence, verifySchema } from "asn1js";
import { tryit } from "radash";
import {
    Address,
    Hex,
    Signature,
    isAddressEqual,
    recoverAddress,
    signatureToHex,
    toBytes,
    toHex,
} from "viem";

/**
 * ASN1 schema to parse the signature from KMS
 */
const EcdsaSigAsnParse = new Sequence({
    name: "EcdsaSig",
    value: [new Integer({ name: "r" }), new Integer({ name: "s" })],
});

/**
 * Get a KMS signature for the given message
 * @param keyId
 * @param address
 * @param msg
 * @param client
 */
export async function getKmsSignature({
    keyId,
    address,
    msg,
    client,
}: {
    keyId: string;
    address: Address;
    msg: Hex;
    client: KMSClient;
}) {
    const msgHash = Buffer.from(toBytes(msg));
    const command = new SignCommand({
        KeyId: keyId,
        Message: msgHash,
        SigningAlgorithm: "ECDSA_SHA_256",
        MessageType: "DIGEST",
    });
    const res = await client.send(command);
    if (!res.Signature) {
        throw new Error("Missing signature");
    }
    const signature = Buffer.from(res.Signature);
    return extractSignature({ signature, msgHash, address });
}

/**
 * Extract the signature from the given buffer
 * @param signature
 * @param msgHash
 * @param address
 */
async function extractSignature({
    signature,
    msgHash,
    address,
}: {
    signature: Buffer;
    msgHash: Buffer;
    address: Address;
}): Promise<Signature> {
    const baseSignature = getSigRs(signature);
    const { v } = await getSigV({ msgHash, address, baseSignature });
    return { ...baseSignature, v };
}

/**
 * Extract 'r' and 's' from the signature
 * @param signature
 */
function getSigRs(signature: Buffer) {
    // Ensurethe signature match our expected format
    const decodedSignature = verifySchema(signature, EcdsaSigAsnParse);
    if (!decodedSignature.verified) {
        throw new Error("Invalid signature");
    }

    // Extract R & S from the signature
    // We are using hex here since the value overflow int
    const rawR = toHex(
        (decodedSignature.result.valueBlock as any).value[0].valueBlock
            .valueHexView
    );
    const rawS = toHex(
        (decodedSignature.result.valueBlock as any).value[1].valueBlock
            .valueHexView
    );
    const r = BigInt(rawR);
    let s = BigInt(rawS);

    // Reput the signature on a valid secp256 formt
    const secp256k1N = BigInt(
        "0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141"
    );

    const secp256k1halfN = secp256k1N / 2n;
    if (s > secp256k1halfN) {
        s = secp256k1N - s;
    }

    return { r: toHex(r), s: toHex(s) };
}

/**
 * Find the right 'v' value for the signature to match the given 'address'
 * @param msgHash
 * @param address
 * @param baseSignature
 */
async function getSigV({
    msgHash,
    address,
    baseSignature,
}: {
    msgHash: Buffer;
    address: Address;
    baseSignature: { r: Hex; s: Hex };
}) {
    let v = 27n;
    let signature = signatureToHex({
        r: baseSignature.r,
        s: baseSignature.s,
        v,
    });
    let [, recovered] = await tryit(recoverAddress)({
        hash: msgHash,
        signature,
    });
    if (!recovered || !isAddressEqual(recovered, address)) {
        v = 28n;
        signature = signatureToHex({
            r: baseSignature.r,
            s: baseSignature.s,
            v,
        });
        [, recovered] = await tryit(recoverAddress)({
            hash: msgHash,
            signature,
        });
    }
    if (!recovered || !isAddressEqual(recovered, address)) {
        throw new Error(
            "signature is invalid. recovered address does not match"
        );
    }
    return { v };
}