import type { SmartAccount } from "permissionless/accounts"
import type { EntryPoint } from "permissionless/types"
import type { Address } from "viem"
import type { Chain, Client, Transport } from "viem"
import type { RepayToken } from "../../types"
import type {
    AllTokens,
    CabPaymasterRpcScheme,
    GetRpcCabPaymasterBalanceParameters
} from "../../types"

export type GetCabBalanceParameters = {
    address: Address
    token?: AllTokens | AllTokens[]
    repayTokens?: RepayToken[]
}

export type GetCabBalanceReturnType = bigint

export async function getCabBalance<
    TEntryPoint extends EntryPoint,
    TChain extends Chain | undefined,
    TTransport extends Transport = Transport,
    TSmartAccount extends
        | SmartAccount<TEntryPoint, string, TTransport, TChain>
        | undefined =
        | SmartAccount<TEntryPoint, string, TTransport, TChain>
        | undefined
>(
    client: Client<
        TTransport,
        TChain,
        TSmartAccount,
        CabPaymasterRpcScheme<TEntryPoint>
    >,
    args: GetCabBalanceParameters
): Promise<GetCabBalanceReturnType> {
    const { address, token, repayTokens } = args
    const params: GetRpcCabPaymasterBalanceParameters = [
        address,
        token ?? null,
        repayTokens ?? null
    ]

    const response = await client.request({
        method: "pm_getCabBalance",
        params: params
    })
    return BigInt(response)
}