import type { SmartAccount } from "permissionless/accounts"
import type { EntryPoint, Prettify } from "permissionless/types"
import type { Chain, Client, PublicClientConfig, Transport } from "viem"
import {
    type CABPaymasterActions,
    cabPaymasterActions
} from "../cabPaymaster/clients/decorators/cabPaymasterActions"
import type { CAB_VERSION, CabPaymasterRpcScheme } from "../types"
import { type CABActions, cabActions } from "./decorators/cabActions"

export type CABClient<
    TEntryPoint extends EntryPoint,
    TTransport extends Transport,
    TChain extends Chain,
    TSmartAccount extends
        | SmartAccount<TEntryPoint, string, TTransport, TChain>
        | undefined =
        | SmartAccount<TEntryPoint, string, TTransport, TChain>
        | undefined
> = Client<
    TTransport,
    TChain,
    TSmartAccount,
    CabPaymasterRpcScheme<TEntryPoint>,
    CABPaymasterActions<TEntryPoint> &
        CABActions<TEntryPoint, TTransport, TChain>
>

export const createCABClient = <
    TEntryPoint extends EntryPoint,
    TTransport extends Transport,
    TChain extends Chain,
    TSmartAccount extends SmartAccount<TEntryPoint, string, TTransport, TChain>
>(
    client: Prettify<Client<TTransport, TChain, TSmartAccount>>,
    parameters: PublicClientConfig<TTransport, TChain> & {
        entryPoint: TEntryPoint
        cabVersion?: CAB_VERSION
    }
): CABClient<TEntryPoint, TTransport, TChain, TSmartAccount> => {
    const { transport, entryPoint: entryPoint_, cabVersion } = parameters
    const entryPoint = entryPoint_ ?? client.account.entryPoint

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