import {
    serializeMultiChainPermissionAccounts,
    serializePermissionAccount
} from "@zerodev/permissions"
import { toEmptyECDSASigner } from "@zerodev/permissions/signers"
import type { KernelSmartAccount } from "@zerodev/sdk"
import type { EntryPoint, Prettify } from "permissionless/types"
import type { Address, Chain, Client, Transport } from "viem"
import type { CabPaymasterRpcScheme } from "../../types"
import {
    type EnableCABSupportedToken,
    cabPaymasterAddress,
    getInvoiceManagerAddress,
    getSupportedTokens,
    getVaultManagerAddress
} from "../../utils/constants"
import { getPermissionAccount } from "../utils/getPermissionAccount"

export type EnableCABParameters = {
    tokens: EnableCABSupportedToken[]
    cabPaymasterAddress?: Address
}

export async function enableCAB<
    TEntryPoint extends EntryPoint,
    TChain extends Chain | undefined = Chain | undefined,
    TTransport extends Transport = Transport,
    TSmartAccount extends
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined =
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined
>(
    client: Client<
        TTransport,
        TChain,
        TSmartAccount,
        CabPaymasterRpcScheme<TEntryPoint>
    >,
    args: Prettify<EnableCABParameters>,
    transport: Transport
) {
    const { tokens, cabPaymasterAddress: cabPaymasterAddress_ } = args
    console.log("enable", cabPaymasterAddress_)
    const paymasterAddress = cabPaymasterAddress_ ?? cabPaymasterAddress
    const invoiceManagerAddress = getInvoiceManagerAddress({
        cabPaymasterAddress: paymasterAddress
    })
    const vaultManagerAddress = getVaultManagerAddress({
        cabPaymasterAddress: paymasterAddress
    })

    const account = client.account
    const chain = client.chain
    const entryPointAddress = account?.entryPoint

    // get validator source
    const validatorSource = account?.kernelPluginManager.sudoValidator?.source
    const supportedMultiNetworks =
        validatorSource === "MultiChainECDSAValidator" ||
        validatorSource === "MultiChainWebAuthnValidator"

    if (!entryPointAddress || !chain) {
        throw new Error("Invalid entryPoint or chain")
    }
    const { value } = transport({})
    if (!value?.url) {
        throw new Error("Invalid transport")
    }
    const paymasterUrl = value.url.replace(/\/$/, "")

    if (
        !supportedMultiNetworks &&
        tokens.some(({ networks }) => !networks || networks.length > 1)
    ) {
        throw new Error(
            "Multi network can only be used with multichain validator"
        )
    }

    // 1. get session signer address
    let getSessionAddressResponse: Response
    try {
        getSessionAddressResponse = await fetch(`${paymasterUrl}/session`, {
            method: "POST",
            body: JSON.stringify({
                accountAddress: account.address
            }),
            headers: {
                "Content-Type": "application/json"
            }
        })
        if (!getSessionAddressResponse.ok) {
            const errorBody = await getSessionAddressResponse.json()
            const errorMessage =
                errorBody.message || getSessionAddressResponse.statusText
            throw new Error(
                `Failed to get session signer address: ${errorMessage}`
            )
        }
        const sessionAddress: Address = (await getSessionAddressResponse.json())
            .sessionSignerAddress

        const sessionSigner = toEmptyECDSASigner(sessionAddress)

        // 2. construct the tokens by chainId
        const tokensByChainIdArray = Object.values(
            tokens.reduce(
                (acc, { name, networks }) => {
                    // enable on all network if chainId is not provided
                    const supportedTokens = getSupportedTokens({
                        cabPaymasterAddress: paymasterAddress
                    })
                    const chainIdsToUse = networks
                        ? networks
                        : Object.keys(supportedTokens[name]).map(Number)
                    for (const id of chainIdsToUse) {
                        if (supportedTokens[name][id]) {
                            if (!acc[id]) {
                                acc[id] = { chainId: id, vaults: [] }
                            }
                            acc[id].vaults.push(supportedTokens[name][id])
                        }
                    }
                    return acc
                },
                {} as Record<
                    number,
                    {
                        chainId: number
                        vaults: { vault: Address; token: Address }[]
                    }
                >
            )
        )

        const accounts = await Promise.all(
            tokensByChainIdArray.map(async ({ chainId, vaults }) => {
                return getPermissionAccount({
                    account: account as KernelSmartAccount<
                        TEntryPoint,
                        TTransport,
                        Chain
                    >,
                    entryPoint: account.entryPoint,
                    sessionSigner,
                    chainId,
                    vaults,
                    cabPaymasterAddress: paymasterAddress
                })
            })
        )

        const approvals = supportedMultiNetworks
            ? await serializeMultiChainPermissionAccounts(
                  accounts.map((account) => ({ account }))
              )
            : await Promise.all(
                  accounts.map((account) => serializePermissionAccount(account))
              )

        const serializedSessionAccounts = tokensByChainIdArray.map((c, i) => {
            return {
                chainId: c.chainId,
                vaultManager: vaultManagerAddress,
                invoiceManager: invoiceManagerAddress,
                vaults: c.vaults,
                approval: approvals[i]
            }
        })
        try {
            // Loop through each session account and enable CAB because if we send
            // as one reqeust the request body size will be too large
            const enableRequests = serializedSessionAccounts.map(
                (sessionAccount) =>
                    fetch(`${paymasterUrl}/cab`, {
                        method: "POST",
                        body: JSON.stringify({
                            address: account.address,
                            serializedSessionAccounts: [sessionAccount],
                            paymaster: paymasterAddress
                        }),
                        headers: {
                            "Content-Type": "application/json"
                        }
                    }).then(async (response) => {
                        if (!response.ok) {
                            const errorBody = await response.json()
                            throw new Error(
                                `Failed to enable CAB for chain ${sessionAccount.chainId}: ${
                                    errorBody.message || response.statusText
                                }`
                            )
                        }
                        return response
                    })
            )

            await Promise.all(enableRequests)
        } catch (err) {
            throw new Error(`Failed to enable cab: ${(err as Error).message}`)
        }
    } catch (err) {
        throw new Error(
            `Failed to get session address: ${(err as Error).message}`
        )
    }
}