import type { KernelSmartAccount } from "@zerodev/sdk"
import type { EntryPoint } from "permissionless/types"
import type { Chain, Client, Hash, Transport } from "viem"
import {
    type EnableCABParameters,
    type GetEnabledChainsParameters,
    type GetEnabledChainsResult,
    type PrepareUserOperationRequestCABParameters,
    type PrepareUserOperationRequestCABReturnType,
    type SendUserOperationCABParameters,
    type SponsorUserOperationParameters,
    type SponsorUserOperationReturnType,
    enableCAB,
    getEnabledChains,
    prepareUserOperationRequestCAB,
    sendUserOperationCAB,
    sponsorUserOperation
} from "../../actions/kernel"
import type { CAB_VERSION } from "../../types"
import { getCabPaymasterAddress } from "../../utils"
import { getInvoiceManagerAddress } from "../../utils/getInvoiceManagerAddress"
import type { KernelCABClient } from "../createKernelCABClient"

export type KernelCABActions<
    TEntryPoint extends EntryPoint,
    TTransport extends Transport,
    TChain extends Chain | undefined,
    TSmartAccount extends
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined =
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined
> = {
    enableCAB: (args: EnableCABParameters) => Promise<void>
    sponsorUserOperation: (
        args: SponsorUserOperationParameters<TEntryPoint>
    ) => Promise<SponsorUserOperationReturnType<TEntryPoint>>
    prepareUserOperationRequestCAB: (
        args: PrepareUserOperationRequestCABParameters<
            TEntryPoint,
            TTransport,
            TChain,
            TSmartAccount
        >
    ) => Promise<PrepareUserOperationRequestCABReturnType<TEntryPoint>>
    sendUserOperationCAB: (
        args: SendUserOperationCABParameters<TEntryPoint>
    ) => Promise<Hash>
    getEnabledChains: (
        args?: GetEnabledChainsParameters
    ) => Promise<GetEnabledChainsResult>
}

export const kernelCABActions =
    <
        TEntryPoint extends EntryPoint,
        TTransport extends Transport,
        TChain extends Chain | undefined,
        TSmartAccount extends
            | KernelSmartAccount<TEntryPoint, TTransport, TChain>
            | undefined =
            | KernelSmartAccount<TEntryPoint, TTransport, TChain>
            | undefined
    >(
        transport: Transport,
        cabVersion?: CAB_VERSION
    ) =>
    (
        client: Client<TTransport, TChain, TSmartAccount>
    ): KernelCABActions<TEntryPoint, TTransport, TChain> => ({
        enableCAB: async (args) =>
            enableCAB<TEntryPoint, TChain, TTransport, TSmartAccount>(
                client,
                {
                    ...args,
                    cabPaymasterAddress: getCabPaymasterAddress(cabVersion)
                },
                transport
            ),
        sponsorUserOperation: async (args) =>
            sponsorUserOperation<
                TEntryPoint,
                TTransport,
                TChain,
                TSmartAccount
            >(
                client as KernelCABClient<
                    TEntryPoint,
                    TTransport,
                    TChain,
                    TSmartAccount
                >,
                {
                    ...args,
                    cabPaymasterAddress: getCabPaymasterAddress(cabVersion)
                }
            ),
        prepareUserOperationRequestCAB: async (args) =>
            prepareUserOperationRequestCAB<
                TEntryPoint,
                TChain,
                TTransport,
                TSmartAccount
            >(
                client,
                {
                    ...args,
                    cabPaymasterAddress: getCabPaymasterAddress(cabVersion)
                } as PrepareUserOperationRequestCABParameters<
                    TEntryPoint,
                    TTransport,
                    TChain,
                    TSmartAccount
                >,
                transport
            ),
        sendUserOperationCAB: async (args) =>
            sendUserOperationCAB<
                TEntryPoint,
                TChain,
                TTransport,
                TSmartAccount
            >(
                client,
                {
                    ...args,
                    cabPaymasterAddress: getCabPaymasterAddress(cabVersion)
                },
                transport
            ),
        getEnabledChains: async (args = {}) =>
            getEnabledChains<TEntryPoint, TChain, TTransport, TSmartAccount>(
                client,
                {
                    ...args,
                    invoiceManagerAddress: getInvoiceManagerAddress(cabVersion)
                },
                transport
            )
    })