import { getUserOperationGasPrice } from "@zerodev/sdk"
import { getEntryPointVersion } from "permissionless"
import type { SmartAccount } from "permissionless/accounts"
import {
    type SendTransactionsWithPaymasterParameters,
    prepareUserOperationRequest
} from "permissionless/actions/smartAccount"
import type {
    ENTRYPOINT_ADDRESS_V06_TYPE,
    ENTRYPOINT_ADDRESS_V07_TYPE,
    EntryPoint,
    GetEntryPointVersion,
    Prettify,
    UserOperation
} from "permissionless/types"
import type { Address, Chain, Client, Transport } from "viem"
import { getCabPaymasterSponsorTokens } from "../../cabPaymaster/actions/getCabPaymasterSponsorTokens"
import {
    type GetCabPaymasterStubDataReturnType,
    getCabPaymasterStubData
} from "../../cabPaymaster/actions/getCabPaymasterStubData"
import type {
    CabPaymasterRpcScheme,
    RepayTokenInfo,
    RepayTokens,
    SponsorTokenInfo,
    StateOverrides
} from "../../types"
import { cabPaymasterAddress } 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
        | SmartAccount<TEntryPoint, string, TTransport, TChain>
        | undefined
> = SendTransactionsWithPaymasterParameters<
    TEntryPoint,
    TTransport,
    TChain,
    TSmartAccount
> & {
    repayTokens: RepayTokens
    cabPaymasterAddress?: Address
}

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

export async function prepareUserOperationRequestCAB<
    TEntryPoint extends EntryPoint,
    TChain extends Chain | undefined = 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>
    >,
    args: Prettify<
        PrepareUserOperationRequestCABParameters<
            TEntryPoint,
            TTransport,
            TChain,
            TSmartAccount
        >
    >,
    transport: Transport,
    stateOverrides?: StateOverrides
): Promise<PrepareUserOperationRequestCABReturnType<TEntryPoint>> {
    const {
        transactions,
        account: account_,
        cabPaymasterAddress: cabPaymasterAddress_
    } = 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 gasPrice = () => getUserOperationGasPrice(client)
    const userOperation = (await prepareUserOperationRequest(
        cabClient,
        {
            account: account as unknown as SmartAccount<
                EntryPoint,
                string,
                TTransport,
                TChain
            >,
            userOperation: {
                callData: await account.encodeCallData(
                    transactions.map(({ to, value, data }) => ({
                        to,
                        value,
                        data
                    }))
                )
            },
            middleware: {
                sponsorUserOperation: async ({ userOperation }) => {
                    return {
                        ...userOperation,
                        callGasLimit: 0n,
                        preVerificationGas: 0n,
                        verificationGasLimit: 0n
                    }
                },
                gasPrice
            }
        },
        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
        }),
        ...transactions,
        createInvoiceCall({
            chainId: chain.id,
            account: account.address,
            nonce: userOperation.nonce,
            paymaster: sponsorTokenResponse.paymaster,
            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
        }),
        ...transactions,
        createInvoiceCall({
            chainId: chain.id,
            account: account.address,
            nonce: userOperation.nonce,
            paymaster: paymaster,
            repayTokensInfo: paymasterStubDataRes.repayTokensInfo
        })
    ])
    return {
        userOperation,
        repayTokensInfo: paymasterStubDataRes.repayTokensInfo,
        sponsorTokensInfo: sponsorTokenResponse.sponsorTokensInfo
    }
}