import {
    type KernelSmartAccount,
    type SponsorUserOperationReturnType,
    getUserOperationGasPrice
} from "@zerodev/sdk"
import { bundlerActions, getEntryPointVersion } from "permissionless"
import type { SmartAccount } from "permissionless/accounts"
import { pimlicoBundlerActions } from "permissionless/actions/pimlico"
import { prepareUserOperationRequest } from "permissionless/actions/smartAccount"
import type {
    ENTRYPOINT_ADDRESS_V06_TYPE,
    ENTRYPOINT_ADDRESS_V07_TYPE,
    EntryPoint,
    GetEntryPointVersion
} from "permissionless/types"
import type { UserOperation } from "permissionless/types"
import {
    http,
    type Address,
    type Chain,
    type Client,
    type Hex,
    type Transport,
    createClient
} from "viem"
import { getCabPaymasterSponsorTokens } from "../../cabPaymaster/actions/getCabPaymasterSponsorTokens"
import {
    type GetCabPaymasterStubDataReturnType,
    getCabPaymasterStubData
} from "../../cabPaymaster/actions/getCabPaymasterStubData"
import type { RepayTokens, StateOverrides, SwapTokenInfo } from "../../types"
import type {
    CabPaymasterRpcScheme,
    RepayTokenInfo,
    SponsorTokenInfo
} from "../../types"
import { cabPaymasterAddress, swapCall } from "../../utils"
import { createInvoiceCall } from "../../utils/createInvoiceCall"
import { withdrawCall } from "../../utils/withdrawCall"

export type PrepareUserOperationRequestCABParameters<
    TEntryPoint extends EntryPoint,
    TTransport extends Transport,
    TChain extends Chain | undefined,
    TSmartAccount extends
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined =
        | KernelSmartAccount<TEntryPoint, TTransport, TChain>
        | undefined
> = {
    calls: { to: Address; value: bigint; data: Hex }[]
    repayTokens: RepayTokens
    account?: TSmartAccount
    cabPaymasterAddress?: Address
} & {
    userOperation?: TEntryPoint extends ENTRYPOINT_ADDRESS_V06_TYPE
        ? Partial<UserOperation<"v0.6">>
        : Partial<UserOperation<"v0.7">>
}

export type PrepareUserOperationRequestCABReturnType<
    TEntryPoint extends EntryPoint
> = {
    userOperation: UserOperation<GetEntryPointVersion<TEntryPoint>>
    sponsorTokensInfo: SponsorTokenInfo[]
    swapTokensInfo: SwapTokenInfo[]
    repayTokensInfo: RepayTokenInfo[]
}

export async function prepareUserOperationRequestCAB<
    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: PrepareUserOperationRequestCABParameters<
        TEntryPoint,
        TTransport,
        TChain,
        TSmartAccount
    >,
    transport: Transport,
    stateOverrides?: StateOverrides
): Promise<PrepareUserOperationRequestCABReturnType<TEntryPoint>> {
    const {
        calls,
        account,
        cabPaymasterAddress: cabPaymasterAddress_,
        userOperation: partialUserOperation
    } = args
    const paymasterAddress = cabPaymasterAddress_ ?? cabPaymasterAddress

    const account_ = account ?? client.account
    const chain = client.chain
    const entryPointAddress = account_?.entryPoint
    if (!entryPointAddress || !chain) {
        throw new Error("Invalid entryPoint or chain")
    }
    const { value, request, config } = transport({
        chain: client.chain
    })
    const cabClient = {
        ...client,
        transport: { ...value, ...config },
        request: request
    } as Client<TTransport, TChain, undefined>
    const pimlicoBundlerClient = createClient({
        chain,
        transport: http(client.transport.url)
    })
        .extend(bundlerActions(entryPointAddress))
        .extend(pimlicoBundlerActions(entryPointAddress))

    const [nonce, factory, factoryData] = await Promise.all([
        partialUserOperation?.nonce ?? account_.getNonce(),
        account_.getFactory(),
        account_.getFactoryData()
    ])
    let gasEstimation = partialUserOperation?.maxFeePerGas &&
        partialUserOperation.maxPriorityFeePerGas && {
            maxFeePerGas: partialUserOperation.maxFeePerGas,
            maxPriorityFeePerGas: partialUserOperation.maxPriorityFeePerGas
        }
    try {
        gasEstimation = await getUserOperationGasPrice(client)
    } catch (e) {
        gasEstimation = (await pimlicoBundlerClient.getUserOperationGasPrice())
            .standard
    }

    const userOperation = (await prepareUserOperationRequest(
        cabClient,
        {
            account: account_ as unknown as SmartAccount<
                EntryPoint,
                string,
                TTransport,
                TChain
            >,
            userOperation: {
                callData: await account_.encodeCallData(
                    calls.map(({ to, value, data }) => ({
                        to,
                        value,
                        data
                    }))
                ),
                nonce,
                factory,
                factoryData,
                maxFeePerGas: gasEstimation.maxFeePerGas,
                maxPriorityFeePerGas: gasEstimation.maxPriorityFeePerGas
            },
            middleware: {
                sponsorUserOperation: async ({ userOperation }) => {
                    return {
                        ...userOperation,
                        callGasLimit: 0n,
                        preVerificationGas: 0n,
                        verificationGasLimit: 0n
                    } as SponsorUserOperationReturnType<TEntryPoint>
                }
            }
        },
        stateOverrides
    )) as UserOperation<GetEntryPointVersion<TEntryPoint>>

    const sponsorTokenResponse = await getCabPaymasterSponsorTokens(cabClient, {
        userOperation,
        entryPoint: entryPointAddress,
        sponsorChain: chain,
        cabPaymasterAddress: paymasterAddress
    })

    userOperation.callData = await account_.encodeCallData([
        ...withdrawCall({
            paymaster: sponsorTokenResponse.paymaster,
            accountAddress: account_.address,
            sponsorTokensInfo: sponsorTokenResponse.sponsorTokensInfo
        }),
        ...swapCall({
            swapTokensInfo: sponsorTokenResponse.swapTokensInfo
        }),
        ...calls,
        createInvoiceCall({
            chainId: chain.id,
            account: account_.address,
            nonce: userOperation.nonce,
            paymaster: sponsorTokenResponse.paymaster,
            repayTokensInfo: sponsorTokenResponse.repayTokensInfo
        })
    ])
    const paymasterStubDataRes = await getCabPaymasterStubData(cabClient, {
        userOperation,
        entryPoint: account_.entryPoint,
        sponsorChain: chain,
        repayTokens: args.repayTokens,
        cabPaymasterAddress: paymasterAddress
    })

    let paymaster: Address
    userOperation.callGasLimit = paymasterStubDataRes.callGasLimit ?? 0n
    userOperation.verificationGasLimit =
        paymasterStubDataRes.verificationGasLimit ?? 0n
    userOperation.preVerificationGas =
        paymasterStubDataRes.preVerificationGas ?? 0n

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

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

        paymaster = paymasterStubData07.paymaster
        userOperation.paymaster = paymasterStubData07.paymaster
        userOperation.paymasterData = paymasterStubData07.paymasterData
        userOperation.paymasterVerificationGasLimit =
            paymasterStubData07.paymasterVerificationGasLimit ?? 0n
        userOperation.paymasterPostOpGasLimit =
            paymasterStubData07.paymasterPostOpGasLimit ?? 0n
    }

    userOperation.callData = await account_.encodeCallData([
        ...withdrawCall({
            paymaster: paymaster,
            accountAddress: account_.address,
            sponsorTokensInfo: sponsorTokenResponse.sponsorTokensInfo
        }),
        ...swapCall({
            swapTokensInfo: sponsorTokenResponse.swapTokensInfo
        }),
        ...calls,
        createInvoiceCall({
            chainId: chain.id,
            account: account_.address,
            nonce: userOperation.nonce,
            paymaster: paymaster,
            repayTokensInfo: paymasterStubDataRes.repayTokensInfo
        })
    ])
    return {
        userOperation,
        sponsorTokensInfo: sponsorTokenResponse.sponsorTokensInfo,
        swapTokensInfo: sponsorTokenResponse.swapTokensInfo,
        repayTokensInfo: paymasterStubDataRes.repayTokensInfo
    }
}