import type { KernelSmartAccount } from "@zerodev/sdk"
import { getEntryPointVersion, sendUserOperation } from "permissionless"
import type {
    ENTRYPOINT_ADDRESS_V06_TYPE,
    ENTRYPOINT_ADDRESS_V07_TYPE,
    EntryPoint,
    GetEntryPointVersion,
    Prettify,
    UserOperation
} from "permissionless/types"
import type { Address, Chain, Client, Hash, Transport } from "viem"
import { getCabPaymasterData } from "../../cabPaymaster"
import type { CabPaymasterRpcScheme } from "../../types"
import type { GetRpcCabPaymasterDataReturnType } from "../../types"
import { cabPaymasterAddress } from "../../utils"

export type SendUserOperationCABParameters<TEntryPoint extends EntryPoint> = {
    userOperation: UserOperation<GetEntryPointVersion<TEntryPoint>>
    cabPaymasterAddress?: Address
}

export async function sendUserOperationCAB<
    TEntryPoint extends EntryPoint,
    TChain extends Chain | undefined = Chain | undefined,
    TTransport extends Transport = Transport,
    TSmartAccount extends
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined =
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined
>(
    client: Client<
        TTransport,
        TChain,
        TSmartAccount,
        CabPaymasterRpcScheme<TEntryPoint>
    >,
    args: Prettify<SendUserOperationCABParameters<TEntryPoint>>,
    transport: Transport
): Promise<Hash> {
    const { userOperation, cabPaymasterAddress: cabPaymasterAddress_ } = args
    const paymasterAddress = cabPaymasterAddress_ ?? cabPaymasterAddress

    const account = client.account
    const chain = client.chain
    if (!account || !chain) {
        throw new Error("Invalid account or chain")
    }
    const entryPointAddress = account.entryPoint

    const { value, request, config } = transport({
        chain: client.chain
    })
    const cabClient = {
        ...client,
        transport: { ...value, ...config },
        request: request
    } as Client<
        TTransport,
        TChain,
        undefined,
        CabPaymasterRpcScheme<TEntryPoint>
    >
    const paymasterDataRes = await getCabPaymasterData(cabClient, {
        userOperation: userOperation,
        entryPoint: entryPointAddress,
        sponsorChain: chain,
        cabPaymasterAddress: paymasterAddress
    })
    if (getEntryPointVersion(entryPointAddress) === "v0.6") {
        const paymasterDataResV06 =
            paymasterDataRes as GetRpcCabPaymasterDataReturnType<ENTRYPOINT_ADDRESS_V06_TYPE>
        userOperation.paymasterAndData = paymasterDataResV06.paymasterAndData
    } else {
        const paymasterDataResV07 =
            paymasterDataRes as GetRpcCabPaymasterDataReturnType<ENTRYPOINT_ADDRESS_V07_TYPE>
        userOperation.paymaster = paymasterDataResV07.paymaster
        userOperation.paymasterData = paymasterDataResV07.paymasterData
    }
    const signature = await account.signUserOperation(userOperation)
    userOperation.signature = signature

    return await sendUserOperation(client, {
        userOperation: userOperation,
        entryPoint: entryPointAddress
    })
}