import type { KernelAccountClient, KernelSmartAccount } from "@zerodev/sdk"
import type { EntryPoint } from "permissionless/types"
import type { Chain, Client, PublicClientConfig, Transport } from "viem"
import { cabPaymasterActions } from "../cabPaymaster/clients/decorators/cabPaymasterActions"
import type { CABPaymasterActions } from "../cabPaymaster/clients/decorators/cabPaymasterActions"
import type { CAB_VERSION, CabPaymasterRpcScheme } from "../types"
import {
    type KernelCABActions,
    kernelCABActions
} from "./decorators/kernelCABActions"

export type KernelCABClient<
    TEntryPoint extends EntryPoint,
    TTransport extends Transport = Transport,
    TChain extends Chain | undefined = Chain | undefined,
    TSmartAccount extends
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined =
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined
> = Client<
    TTransport,
    TChain,
    TSmartAccount,
    CabPaymasterRpcScheme<TEntryPoint>,
    CABPaymasterActions<TEntryPoint> &
        KernelCABActions<TEntryPoint, TTransport, TChain>
>

export const createKernelCABClient = <
    TEntryPoint extends EntryPoint,
    TTransport extends Transport,
    TChain extends Chain | undefined,
    TSmartAccount extends
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined =
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined
>(
    client: KernelAccountClient<TEntryPoint, TTransport, TChain, TSmartAccount>,
    parameters: PublicClientConfig<TTransport, TChain> & {
        entryPoint: TEntryPoint
        cabVersion?: CAB_VERSION
    }
): KernelCABClient<TEntryPoint, TTransport, TChain, TSmartAccount> => {
    const { transport, entryPoint, cabVersion } = parameters

    return client
        .extend(cabPaymasterActions(entryPoint, transport, cabVersion))
        .extend(kernelCABActions(transport, cabVersion))
}