import type { KernelSmartAccount } from "@zerodev/sdk"
import { getEntryPointVersion } from "permissionless"
import type { EstimateUserOperationGasReturnType } from "permissionless/actions"
import type {
    ENTRYPOINT_ADDRESS_V06_TYPE,
    ENTRYPOINT_ADDRESS_V07_TYPE,
    EntryPoint,
    GetEntryPointVersion,
    UserOperation
} from "permissionless/types"
import type { Address, Chain, Transport } from "viem"
import { getChainId } from "viem/actions"
import type {
    GetCabPaymasterDataReturnType,
    GetCabPaymasterStubDataReturnType
} from "../../cabPaymaster"
import type { KernelCABClient } from "../../clients/createKernelCABClient"
import type { RepayTokens } from "../../types"
import { createInvoiceCall, getGasParameters, withdrawCall } from "../../utils"
import { cabPaymasterAddress } from "../../utils"
import { decodeCalldata } from "../../utils/decodeCalldata"
import { getChain } from "../../utils/getPublicClient"

export type SponsorUserOperationParameters<TEntryPoint extends EntryPoint> = {
    userOperation: UserOperation<GetEntryPointVersion<TEntryPoint>>
    bundlerTransport: Transport
    repayTokens: RepayTokens
    cabPaymasterAddress?: Address
}

export type PartialPick<T, K extends keyof T> = Partial<Pick<T, K>>

export type SponsorUserOperationReturnType<entryPoint extends EntryPoint> =
    entryPoint extends ENTRYPOINT_ADDRESS_V06_TYPE
        ? Pick<
              UserOperation<"v0.6">,
              | "callGasLimit"
              | "preVerificationGas"
              | "verificationGasLimit"
              | "paymasterAndData"
          > &
              PartialPick<
                  UserOperation<"v0.6">,
                  "maxFeePerGas" | "maxPriorityFeePerGas"
              >
        : Pick<
              UserOperation<"v0.7">,
              | "callGasLimit"
              | "verificationGasLimit"
              | "preVerificationGas"
              | "paymaster"
              | "paymasterVerificationGasLimit"
              | "paymasterPostOpGasLimit"
              | "paymasterData"
          > &
              PartialPick<
                  UserOperation<"v0.7">,
                  "maxFeePerGas" | "maxPriorityFeePerGas"
              >

export const sponsorUserOperation = async <
    TEntryPoint extends EntryPoint,
    TTransport extends Transport,
    TChain extends Chain | undefined,
    TSmartAccount extends
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined =
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined
>(
    client: KernelCABClient<TEntryPoint, TTransport, TChain, TSmartAccount>,
    args: SponsorUserOperationParameters<TEntryPoint>
): Promise<SponsorUserOperationReturnType<TEntryPoint>> => {
    const {
        userOperation,
        repayTokens,
        bundlerTransport,
        cabPaymasterAddress: cabPaymasterAddress_
    } = args
    const paymasterAddress = cabPaymasterAddress_ ?? cabPaymasterAddress

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

    const chainId = client.chain?.id || (await getChainId(client))
    const chain = client.chain ?? getChain(chainId)
    const calls = decodeCalldata({
        accountAddress: userOperation.sender,
        callData: userOperation.callData,
        entryPoint: entryPoint
    })

    // 1. get the sponsorTokens
    const sponsorTokenResponse = await client.getCabPaymasterSponsorTokens({
        userOperation,
        entryPoint,
        sponsorChain: chain,
        cabPaymasterAddress: paymasterAddress
    })
    userOperation.callData = await client.account.encodeCallData([
        ...withdrawCall({
            paymaster: sponsorTokenResponse.paymaster,
            accountAddress: client.account.address,
            sponsorTokensInfo: sponsorTokenResponse.sponsorTokensInfo
        }),
        ...calls,
        createInvoiceCall({
            chainId,
            account: client.account.address,
            nonce: userOperation.nonce,
            paymaster: sponsorTokenResponse.paymaster,
            repayTokensInfo: []
        })
    ])

    // 2. get paymaster stubData
    const paymasterStubDataRes = await client.getCabPaymasterStubData({
        userOperation,
        entryPoint,
        sponsorChain: chain,
        repayTokens,
        cabPaymasterAddress: paymasterAddress
    })
    let paymaster: Address

    if (getEntryPointVersion(entryPoint) === "v0.6") {
        const paymasterStubData06 =
            paymasterStubDataRes as GetCabPaymasterStubDataReturnType<ENTRYPOINT_ADDRESS_V06_TYPE>
        userOperation.paymasterAndData = paymasterStubData06.paymasterAndData

        const gasParameters =
            paymasterStubData06.preVerificationGas &&
            paymasterStubData06.verificationGasLimit &&
            paymasterStubData06.callGasLimit
                ? {
                      preVerificationGas:
                          paymasterStubData06.preVerificationGas,
                      verificationGasLimit:
                          paymasterStubData06.verificationGasLimit,
                      callGasLimit: paymasterStubData06.callGasLimit
                  }
                : await getGasParameters({
                      userOperation,
                      entryPoint,
                      bundlerTransport
                  })
        userOperation.preVerificationGas = gasParameters.preVerificationGas
        userOperation.verificationGasLimit =
            gasParameters.verificationGasLimit || 0n
        userOperation.callGasLimit = gasParameters.callGasLimit || 0n

        paymaster = paymasterStubData06.paymasterAndData.slice(0, 42) as Address
    } else {
        const paymasterStubData07 =
            paymasterStubDataRes as GetCabPaymasterStubDataReturnType<ENTRYPOINT_ADDRESS_V07_TYPE>
        userOperation.paymaster = paymasterStubData07.paymaster
        userOperation.paymasterData = paymasterStubData07.paymasterData

        const gasParameters =
            paymasterStubData07.preVerificationGas &&
            paymasterStubData07.verificationGasLimit &&
            paymasterStubData07.callGasLimit &&
            paymasterStubData07.paymasterVerificationGasLimit &&
            paymasterStubData07.paymasterPostOpGasLimit
                ? {
                      preVerificationGas:
                          paymasterStubData07.preVerificationGas,
                      verificationGasLimit:
                          paymasterStubData07.verificationGasLimit,
                      callGasLimit: paymasterStubData07.callGasLimit,
                      paymasterVerificationGasLimit:
                          paymasterStubData07.paymasterVerificationGasLimit,
                      paymasterPostOpGasLimit:
                          paymasterStubData07.paymasterPostOpGasLimit
                  }
                : ((await getGasParameters({
                      userOperation,
                      entryPoint,
                      bundlerTransport
                  })) as EstimateUserOperationGasReturnType<ENTRYPOINT_ADDRESS_V07_TYPE>)
        userOperation.preVerificationGas = gasParameters.preVerificationGas
        userOperation.verificationGasLimit = gasParameters.verificationGasLimit
        userOperation.callGasLimit = gasParameters.callGasLimit
        userOperation.paymasterVerificationGasLimit =
            gasParameters.paymasterVerificationGasLimit
        userOperation.paymasterPostOpGasLimit =
            gasParameters.paymasterPostOpGasLimit

        paymaster = paymasterStubData07.paymaster
    }

    userOperation.callData = await client.account.encodeCallData([
        ...withdrawCall({
            paymaster,
            accountAddress: client.account.address,
            sponsorTokensInfo: sponsorTokenResponse.sponsorTokensInfo
        }),
        ...calls,
        createInvoiceCall({
            chainId,
            account: client.account.address,
            nonce: userOperation.nonce,
            paymaster,
            repayTokensInfo: paymasterStubDataRes.repayTokensInfo
        })
    ])

    // 3. get paymaster data
    const paymasterDataRes = await client.getCabPaymasterData({
        userOperation,
        entryPoint,
        sponsorChain: chain
    })

    if (getEntryPointVersion(entryPoint) === "v0.6") {
        const paymasterData06 =
            paymasterDataRes as GetCabPaymasterDataReturnType<ENTRYPOINT_ADDRESS_V06_TYPE>

        userOperation.paymasterAndData = paymasterData06.paymasterAndData
    } else {
        const paymasterData07 =
            paymasterDataRes as GetCabPaymasterDataReturnType<ENTRYPOINT_ADDRESS_V07_TYPE>
        userOperation.paymaster = paymasterData07.paymaster
        userOperation.paymasterData = paymasterData07.paymasterData
    }

    return userOperation as SponsorUserOperationReturnType<TEntryPoint>
}