import {
    type CreateDeriveFn,
    ecdsaCreateDerive,
    ed25519CreateDerive,
    sr25519CreateDerive,
} from "@polkadot-labs/hdkd"
import {
    sr25519,
    ed25519,
    ecdsa,
    KeyPair,
    Curve,
} from "@polkadot-labs/hdkd-helpers"
import { KeystoreMeta, keystoreV4, type KeystoreV4WithMeta } from "./keystore"
import { assert } from "./utils"
import * as storage from "./storage"
import {
    InsertCryptoKeyArgs,
    UpdateCryptoKeyArgs,
    RemoveCryptoKeyArgs,
    KeystoreAccount,
    RevealCryptoKeyArgs,
} from "./types"
import { toHex, fromHex } from "@polkadot-api/utils"
import { wellKnownChainIdByGenesisHash } from "../constants"

const createDeriveFnMap: Record<string, CreateDeriveFn> = {
    Sr25519: sr25519CreateDerive,
    Ed25519: ed25519CreateDerive,
    Ecdsa: ecdsaCreateDerive,
}

const curveFnMap: Record<string, Curve> = {
    Sr25519: sr25519,
    Ed25519: ed25519,
    Ecdsa: ecdsa,
}

const createKeyPair = (privateKey: string, scheme: string): KeyPair => {
    const curve = curveFnMap[scheme]
    if (!curve) throw new Error("unsupported signature scheme")

    return {
        publicKey: curve.getPublicKey(privateKey),
        sign(message) {
            return curve.sign(message, privateKey)
        },
    }
}

export const createKeyring = () => {
    const getKeystore = () => storage.get("keystore")
    const setKeystore = (keystore: KeystoreV4WithMeta) =>
        storage.set("keystore", keystore)
    const removeKeystore = () => storage.remove("keystore")

    const getKeystoreAccounts = (
        keystoreMeta: KeystoreMeta,
    ): KeystoreAccount[] => {
        switch (keystoreMeta.type) {
            case "KeysetKeystore":
                return keystoreMeta.derivationPaths.map((d) => ({
                    ...d,
                    type: "Keyset",
                }))
            case "KeypairKeystore":
                return [
                    {
                        type: "Keypair",
                        publicKey: keystoreMeta.publicKey,
                    },
                ]
        }
    }

    const getCryptoKeys = async () => {
        const keys = (await getKeystore())?.meta ?? []

        return keys.map((meta) => ({
            name: meta.name,
            scheme: meta.scheme,
            accounts: getKeystoreAccounts(meta),
            createdAt: meta.createdAt,
        }))
    }

    const decodeSecrets = (secrets: Uint8Array) =>
        JSON.parse(new TextDecoder().decode(secrets)) as string[]
    const encodeSecrets = (secrets: string[]) =>
        new TextEncoder().encode(JSON.stringify(secrets))

    const getAccounts = async (chainId: string) => {
        const keystore = await getKeystore()
        if (!keystore) return []

        return keystore.meta
            .flatMap(getKeystoreAccounts)
            .filter(
                (account) =>
                (account.type === "Keyset" && account.chainId === chainId) ||
                account.type !== "Keyset",
            )
    }

    const insertCryptoKey = async (args: InsertCryptoKeyArgs) => {
        const keystore = await getKeystore()
        assert(keystore, "keyring must be setup")
        assert(currentPassword, "keyring must be unlocked")
        assert(
            ["Sr25519", "Ed25519", "Ecdsa"].includes(args.scheme),
            "invalid signature scheme",
        )

        const secrets = decodeSecrets(keystoreV4.decrypt(keystore, currentPassword))
        const secret = args.type === "Keyset" ? args.miniSecret : args.privatekey
        const newKeystore = keystoreV4.create(
            currentPassword,
            encodeSecrets([...secrets, secret]),
        )

        const newCryptoKey =
            args.type === "Keyset"
            ? {
                type: "KeysetKeystore" as const,
                derivationPaths: args.derivationPaths,
            }
            : args.type === "Keypair"
            ? {
                type: "KeypairKeystore" as const,
                publicKey: toHex(
                    createKeyPair(args.privatekey, args.scheme).publicKey,
                ),
            }
            : undefined
        if (!newCryptoKey) throw new Error("invalid keystore type")

        setKeystore({
            ...newKeystore,
            meta: [
                ...keystore.meta,
                {
                    name: args.name,
                    scheme: args.scheme,
                    createdAt: args.createdAt,
                    ...newCryptoKey,
                },
            ],
        })
    }

    const updateCryptoKey = async (args: UpdateCryptoKeyArgs) => {
        const keystore = await getKeystore()
        assert(keystore, "keyring must be setup")
        assert(currentPassword, "keyring must be unlocked")
        const cryptoKeyIndex = keystore.meta.findIndex((k) => k.name === args.name)
        assert(cryptoKeyIndex != -1, "keyring name not found")

        const uniqueNetworks = new Map()
        keystore.meta[cryptoKeyIndex].derivationPaths.forEach((obj: any) =>
            uniqueNetworks.set(wellKnownChainIdByGenesisHash[obj.chainId], {
                network: wellKnownChainIdByGenesisHash[obj.chainId],
                chainId: obj.chainId,
            })
        )

        const secrets = decodeSecrets(keystoreV4.decrypt(keystore, currentPassword))
        const derive = createDeriveFnMap[keystore.meta[cryptoKeyIndex].scheme](secrets[cryptoKeyIndex])
        const newDerivationPaths = Array.from(uniqueNetworks.values())
            .map(({ network, chainId }) => {
                return {
                    chainId,
                    path: `//${network}/${args.deviation}`,
                    publicKey: toHex(derive(`//${network}/${args.deviation}`).publicKey),
                }
            })

        keystore.meta[cryptoKeyIndex].derivationPaths.push(...newDerivationPaths)
        setKeystore(keystore)
    }

    const clearCryptoKey = async (name: string, index: number) => {
        const keystore = await getKeystore()
        assert(keystore, "keyring must be setup")
        assert(currentPassword, "keyring must be unlocked")
        const cryptoKeyIndex = keystore.meta.findIndex(obj => obj.name === name)
        assert(cryptoKeyIndex != -1, "keyring name not found")
        const secrets = decodeSecrets(keystoreV4.decrypt(keystore, currentPassword))
        switch (keystore.meta.at(cryptoKeyIndex)?.type) {
            case "KeysetKeystore":
                keystore.meta.at(cryptoKeyIndex)?.derivationPaths.splice(index, 1)
                setKeystore(keystore)
                break
            case "KeypairKeystore":
                secrets.splice(cryptoKeyIndex, 1)
                keystore.meta.splice(cryptoKeyIndex, 1)
                const newKeystore = keystoreV4.create(
                    currentPassword,
                    encodeSecrets(secrets),
                )
                setKeystore({
                    ...newKeystore,
                    meta: keystore.meta,
                })
                break
            default:
                throw new Error("invalid keystore type")
        }
    }

    const clearCryptoKeys = async (args: RemoveCryptoKeyArgs) => {
        const keystore = await getKeystore()
        assert(keystore, "keyring must be setup")
        assert(currentPassword, "keyring must be unlocked")
        const cryptoKeyIndex = keystore.meta.findIndex(obj => obj.name === args.name)
        assert(cryptoKeyIndex != -1, "keyring name not found")
        const secrets = decodeSecrets(keystoreV4.decrypt(keystore, currentPassword))
        secrets.splice(cryptoKeyIndex, 1)
        const newKeystore = keystoreV4.create(
            currentPassword,
            encodeSecrets(secrets),
        )
        keystore.meta.splice(cryptoKeyIndex, 1)
        setKeystore({
            ...newKeystore,
            meta: keystore.meta,
        })
    }

    const revealCryptoKey = async (name: string, index: number) => {
        const keystore = await getKeystore()
        assert(keystore, "keyring must be setup")
        assert(currentPassword, "keyring must be unlocked")
        const cryptoKeyIndex = keystore.meta.findIndex((k) => k.name === name)
        assert(cryptoKeyIndex != -1, "keyring name not found")
        const secret = decodeSecrets(
            keystoreV4.decrypt(keystore, currentPassword)
        ).at(cryptoKeyIndex)
        assert(secret, "secret key should exists")
        // TODO: need to get secret from secret and derivation path
        return secret
    }

    let currentPassword: string | undefined

    return {
        async unlock(password: string) {
            const keystore = await getKeystore()
            assert(keystore, "keyring must be setup")
            if (!keystoreV4.verifyPassword(keystore, password))
            throw new Error("invalid password")
            currentPassword = password
        },
        async lock() {
            assert(await getKeystore(), "keyring must be setup")
            assert(currentPassword, "keyring must be unlocked")
            currentPassword = undefined
        },
        async isLocked() {
            return !currentPassword || !(await getKeystore())
        },
        async reset(password: string) {
            const keystore = await getKeystore()
            assert(keystore, "keyring must be setup")
            assert(currentPassword, "keyring must be unlocked")
            if (!keystoreV4.verifyPassword(keystore, password))
            throw new Error("invalid password")
            currentPassword = undefined
            await removeKeystore()
        },
        async changePassword(password: string, newPassword: string) {
            const keystore = await getKeystore()
            assert(keystore, "keyring must be setup")
            assert(currentPassword, "keyring must be unlocked")
            if (!keystoreV4.verifyPassword(keystore, password))
            throw new Error("invalid password")
            currentPassword = newPassword
            await setKeystore({
                ...keystoreV4.create(
                    newPassword,
                    keystoreV4.decrypt(keystore, password),
                ),
                meta: keystore.meta,
            })
        },
        async setup(password: string) {
            assert(!(await getKeystore()), "keyring is already setup")
            await setKeystore({
                ...keystoreV4.create(
                    password,
                    new TextEncoder().encode(JSON.stringify([])),
                ),
                meta: [],
            })
            currentPassword = password
        },
        async hasPassword() {
            return !!(await getKeystore())
        },
        getAccounts,
        insertCryptoKey,
        updateCryptoKey,
        getCryptoKeys,
        async getCryptoKey(name: string) {
            return (await getCryptoKeys())?.find((m) => m.name === name)
        },
        revealCryptoKey,
        clearCryptoKey,
        clearCryptoKeys,
        async getKeypair(chainId: string, publicKey: string) {
            assert(currentPassword, "keyring must be unlocked")
            const keystore = await getKeystore()
            assert(keystore, "keyring must be setup")

            const keysetIndex = keystore.meta.findIndex((keyset) => {
                switch (keyset.type) {
                    case "KeysetKeystore":
                    return keyset.derivationPaths.some(
                        (d) => d.chainId === chainId && d.publicKey === publicKey,
                    )
                    case "KeypairKeystore":
                    return keyset.publicKey === publicKey
                    default:
                    throw new Error("invalid keystore type")
                }
            })

            if (keysetIndex === -1) {
                throw new Error("unknown account")
            }

            const secret = decodeSecrets(
                keystoreV4.decrypt(keystore, currentPassword),
            )[keysetIndex]

            const cryptoKey = keystore.meta[keysetIndex]
            switch (cryptoKey.type) {
                case "KeysetKeystore": {
                    const { derivationPaths, scheme } = cryptoKey
                    const derivationPath = derivationPaths.find(
                        (d) => d.publicKey === publicKey,
                    )!
                    const createDeriveFn = createDeriveFnMap[scheme]
                    if (!createDeriveFn) throw new Error("invalid signature scheme")
                    return [
                        createDeriveFn(secret)(derivationPath.path),
                        scheme as "Sr25519" | "Ed25519" | "Ecdsa",
                    ] as const
                }
                case "KeypairKeystore": {
                    let keypair = createKeyPair(secret, cryptoKey.scheme)

                    return [keypair, cryptoKey.scheme] as const
                }
            }
        },
    }
}
