import type { SmartAccount } from "permissionless/accounts"
import type { EntryPoint } from "permissionless/types"
import type { Chain, Client, Transport } from "viem"
import type { CAB_VERSION, CabPaymasterRpcScheme } from "../../../types"
import { getCabPaymasterAddress } from "../../../utils"
import {
    type GetCabAllTokensParameters,
    type GetCabAllTokensReturnType,
    type GetCabAllVaultsParameters,
    type GetCabAllVaultsReturnType,
    type GetCabBalanceParameters,
    type GetCabBalanceReturnType,
    type GetCabPaymasterDataParameters,
    type GetCabPaymasterDataReturnType,
    type GetCabPaymasterSponsorTokensParameters,
    type GetCabPaymasterSponsorTokensReturnType,
    type GetCabPaymasterStubDataParameters,
    type GetCabPaymasterStubDataReturnType,
    getCabAllTokens,
    getCabAllVaults,
    getCabBalance,
    getCabPaymasterData,
    getCabPaymasterSponsorTokens,
    getCabPaymasterStubData
} from "../../actions"

export type CABPaymasterActions<entryPoint extends EntryPoint> = {
    getCabAllTokens: (
        args: GetCabAllTokensParameters | undefined
    ) => Promise<GetCabAllTokensReturnType>
    getCabAllVaults: (
        args: GetCabAllVaultsParameters | undefined
    ) => Promise<GetCabAllVaultsReturnType>
    getCabBalance: (
        args: GetCabBalanceParameters
    ) => Promise<GetCabBalanceReturnType>
    getCabPaymasterSponsorTokens: (
        args: GetCabPaymasterSponsorTokensParameters<entryPoint>
    ) => Promise<GetCabPaymasterSponsorTokensReturnType>
    getCabPaymasterStubData: (
        args: GetCabPaymasterStubDataParameters<entryPoint>
    ) => Promise<GetCabPaymasterStubDataReturnType<entryPoint>>
    getCabPaymasterData: (
        args: GetCabPaymasterDataParameters<entryPoint>
    ) => Promise<GetCabPaymasterDataReturnType<entryPoint>>
}

export const cabPaymasterActions =
    <
        TEntryPoint extends EntryPoint,
        TTransport extends Transport,
        TChain extends Chain | undefined,
        TSmartAccount extends
            | SmartAccount<TEntryPoint, string, TTransport, TChain>
            | undefined
    >(
        entryPointAddress: TEntryPoint,
        transport?: Transport,
        cabVersion?: CAB_VERSION
    ) =>
    (
        client: Client<TTransport, TChain, TSmartAccount>
    ): CABPaymasterActions<TEntryPoint> => {
        let cabClient = client
        const cabPaymasterAddress = getCabPaymasterAddress(cabVersion)
        if (transport) {
            const { value, request, config } = transport({
                chain: client.chain
            })
            cabClient = {
                ...client,
                transport: { ...value, ...config },
                request: request
            }
        }
        return {
            getCabAllTokens: async (args = {}) =>
                getCabAllTokens(
                    cabClient as Client<
                        TTransport,
                        TChain,
                        TSmartAccount,
                        CabPaymasterRpcScheme<TEntryPoint>
                    >,
                    {
                        chainId: args.chainId
                    }
                ),
            getCabAllVaults: async (args = {}) =>
                getCabAllVaults(
                    cabClient as Client<
                        TTransport,
                        TChain,
                        TSmartAccount,
                        CabPaymasterRpcScheme<TEntryPoint>
                    >,
                    {
                        chainId: args.chainId,
                        token: args.token
                    }
                ),
            getCabBalance: async (args) =>
                getCabBalance(
                    cabClient as Client<
                        TTransport,
                        TChain,
                        TSmartAccount,
                        CabPaymasterRpcScheme<TEntryPoint>
                    >,
                    {
                        address: args.address,
                        token: args.token,
                        repayTokens: args.repayTokens
                    }
                ),
            getCabPaymasterSponsorTokens: async (args) =>
                getCabPaymasterSponsorTokens(
                    cabClient as Client<
                        TTransport,
                        TChain,
                        TSmartAccount,
                        CabPaymasterRpcScheme<TEntryPoint>
                    >,
                    {
                        userOperation: args.userOperation,
                        entryPoint: entryPointAddress,
                        sponsorChain: args.sponsorChain,
                        cabPaymasterAddress
                    }
                ),
            getCabPaymasterStubData: async (args) =>
                getCabPaymasterStubData(
                    cabClient as Client<
                        TTransport,
                        TChain,
                        TSmartAccount,
                        CabPaymasterRpcScheme<TEntryPoint>
                    >,
                    {
                        userOperation: args.userOperation,
                        entryPoint: entryPointAddress,
                        sponsorChain: args.sponsorChain,
                        repayTokens: args.repayTokens,
                        cabPaymasterAddress
                    }
                ),
            getCabPaymasterData: async (args) =>
                getCabPaymasterData(
                    cabClient as Client<
                        TTransport,
                        TChain,
                        TSmartAccount,
                        CabPaymasterRpcScheme<TEntryPoint>
                    >,
                    {
                        userOperation: args.userOperation,
                        entryPoint: entryPointAddress,
                        sponsorChain: args.sponsorChain,
                        cabPaymasterAddress
                    }
                )
        }
    }