import { getEntryPointVersion } from "permissionless"
import type { SmartAccount } from "permissionless/accounts"
import type {
    ENTRYPOINT_ADDRESS_V06_TYPE,
    ENTRYPOINT_ADDRESS_V07_TYPE,
    EntryPoint,
    GetEntryPointVersion,
    UserOperation
} from "permissionless/types"
import { deepHexlify } from "permissionless/utils"
import type { Address, Chain, Client, Hex, Transport } from "viem"
import type {
    CabPaymasterRpcScheme,
    GetRpcCabPaymasterDataReturnType,
    UserOperationWithBigIntAsHex
} from "../../types"
import { cabPaymasterAddress } from "../../utils"

export type GetCabPaymasterDataParameters<TEntryPoint extends EntryPoint> = {
    userOperation: UserOperation<GetEntryPointVersion<TEntryPoint>>
    entryPoint: TEntryPoint
    sponsorChain: Chain
    cabPaymasterAddress?: Address
}

export type GetCabPaymasterDataReturnType<TEntryPoint extends EntryPoint> =
    GetEntryPointVersion<TEntryPoint> extends "v0.6"
        ? {
              paymasterAndData: Hex
          }
        : {
              paymaster: Hex
              paymasterData: Hex
          }

export async function getCabPaymasterData<
    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>
    >,
    {
        userOperation,
        entryPoint,
        sponsorChain,
        cabPaymasterAddress: cabPaymasterAddress_
    }: GetCabPaymasterDataParameters<TEntryPoint>
): Promise<GetCabPaymasterDataReturnType<TEntryPoint>> {
    const paymasterAddress = cabPaymasterAddress_ ?? cabPaymasterAddress
    const params: [
        UserOperationWithBigIntAsHex<GetEntryPointVersion<TEntryPoint>>,
        TEntryPoint,
        number,
        unknown,
        Address
    ] = [
        deepHexlify(userOperation) as UserOperationWithBigIntAsHex<
            GetEntryPointVersion<TEntryPoint>
        >,
        entryPoint,
        sponsorChain.id,
        [],
        paymasterAddress
    ]

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

    const entryPointVersion = getEntryPointVersion(entryPoint)
    if (entryPointVersion === "v0.6") {
        const responseV06 =
            response as GetRpcCabPaymasterDataReturnType<ENTRYPOINT_ADDRESS_V06_TYPE>

        return {
            paymasterAndData: responseV06.paymasterAndData
        } as GetCabPaymasterDataReturnType<TEntryPoint>
    }
    const responseV07 =
        response as GetRpcCabPaymasterDataReturnType<ENTRYPOINT_ADDRESS_V07_TYPE>

    return {
        paymaster: responseV07.paymaster,
        paymasterData: responseV07.paymasterData
    } as GetCabPaymasterDataReturnType<TEntryPoint>
}