import { randomBytes } from "node:crypto"
import {
  EntrypointAbi,
  NexusBootstrapAbi,
  createSmartAccountClient,
  toNexusAccount
} from "@biconomy/abstractjs"
import {
  http,
  type Address,
  type Hex,
  concatHex,
  encodeAbiParameters,
  encodeFunctionData,
  encodePacked,
  erc20Abi,
  keccak256,
  parseEther
} from "viem"
import { toPackedUserOperation } from "viem/account-abstraction"
import { recoverAuthorizationAddress } from "viem/utils"
import {
  ENTRYPOINT_V07_ADDRESS,
  FREE_MINT_ERC20,
  NEXUS_BOOTSTRAP_ADDRESS,
  NEXUS_IMPLEMENTATION_ADDRESS,
  getBalance,
  toClients
} from "../src"
import type { Infra } from "../src/toEcosystem"

export const benchmarkPrep = async ({
  bundler,
  network: { rpcUrl, chain, privateKey }
}: Infra) => {
  const { publicClient, walletClients, testClient, accounts } = await toClients(
    {
      rpcUrl,
      chain
    }
  )

  const nexusAccount = await toNexusAccount({
    signer: accounts[0],
    chain,
    transport: http(rpcUrl)
  })

  const nexusAccountClient = createSmartAccountClient({
    account: nexusAccount,
    chain,
    transport: http(bundler.url),
    mock: true
  })

  await testClient.setBalance({
    address: nexusAccount.address,
    value: parseEther("10")
  })

  const eoaSigner = accounts[8]
  const someWallet = accounts[9]

  const eoaSignerClient = walletClients[8]
  const someWalletClient = walletClients[9]

  const userOp = await nexusAccountClient.prepareUserOperation({
    calls: [
      {
        to: FREE_MINT_ERC20,
        value: 0n,
        data: encodeFunctionData({
          abi: erc20Abi,
          functionName: "transfer",
          args: [someWallet.address, parseEther("1")]
        })
      }
    ]
  })

  // sign eip7702 authorization with viem
  // We are going to modify the signature in it to fit our PREP algorithm
  const signedAuthorization = await eoaSignerClient.signAuthorization({
    contractAddress: NEXUS_IMPLEMENTATION_ADDRESS,
    chainId: 0,
    nonce: 0
  })

  // ========================================================
  // ===== PREPare data to build a proper PREP signature =====
  // ========================================================

  // ========== Prepare Nexus initdata ============
  const bootstrapData = encodeFunctionData({
    abi: NexusBootstrapAbi,
    functionName: "initNexusWithDefaultValidator",
    args: [encodePacked(["bytes"], [eoaSigner.address])]
  })

  // abi.encode(address, bytes)
  const initData = encodeAbiParameters(
    [{ type: "address" }, { type: "bytes" }],
    [NEXUS_BOOTSTRAP_ADDRESS, bootstrapData]
  )
  const initDataHash = keccak256(initData)

  // ========== Not let's build a proper r and s to make a one-time signature ============
  let found = false
  let saltAndDelegation: Hex = "0x"
  let recoveredAddress: Address = "0x"

  while (!found) {
    // generate random 12 bytes
    // this is the salt that we can vary to get different r and s
    // to find the valid combination that results in a recovered address
    const salt = randomBytes(12).toString("hex")

    // concat the salt and the delegation address (Nexus implementation address)
    // we are going to use that to build the userOp.signature
    saltAndDelegation = concatHex([salt as Hex, NEXUS_IMPLEMENTATION_ADDRESS])

    // convert bytes12 salt to bytes32 salt
    const salt32 = concatHex([
      "0x0000000000000000000000000000000000000000",
      salt as Hex
    ])

    // concat the initDataHash and the salt32
    const initDataHashAndSalt = encodePacked(
      ["bytes32", "bytes32"],
      [initDataHash, salt32]
    )

    // Finally, make an r value
    // This algorithm will be repeated in the smart contract to build r out of the `saltAndDelegation`
    const r = concatHex([
      "0x000000000000000000000000", // zero out first 12 bytes
      keccak256(initDataHashAndSalt).slice(26) as `0x${string}` // hash the `initDataHashAndSalt` and take the last 20 bytes of it
    ])

    // s is just the hash of r
    const s = keccak256(r)

    // cryptography related sanity check
    if (
      BigInt(s) >=
      BigInt(
        "0x7fffffffffffffffffffffffffffffff5d576e7357a4501ddfe92f46681b20a1"
      )
    ) {
      // ECDSA lib on-chain inverts the s value which is bigger than _HALF_N_PLUS_1
      // we just skip this case and look for a different s
      continue
    }

    // record this to the authorization
    // to make a PREP authorization
    signedAuthorization.r = r
    signedAuthorization.s = s
    signedAuthorization.yParity = 0
    signedAuthorization.v = 27n

    try {
      // recover the address from the signed authorization
      recoveredAddress = await recoverAuthorizationAddress({
        authorization: signedAuthorization
      })
      found = true
    } catch (e) {
      // if we ended up with a non recoverable authorization, try other salt => r,s
    }
  }

  // ========================================================
  // ===== PREP algorithm ends =====
  // ========================================================

  // fund the recovered address
  await testClient.setBalance({
    address: recoveredAddress,
    value: parseEther("100")
  })

  //send erc20 token to the recovered address
  const mintHash = await someWalletClient.sendTransaction({
    to: FREE_MINT_ERC20,
    data: encodeFunctionData({
      abi: erc20Abi,
      functionName: "transfer",
      args: [recoveredAddress, parseEther("2")]
    })
  })
  await publicClient.waitForTransactionReceipt({ hash: mintHash })

  // update sender to be the recovered address
  userOp.sender = recoveredAddress
  userOp.factory = "0x"
  userOp.factoryData = "0x"

  // ========================================================
  // ===== Build a proper Nonce with PREP mode =====
  // ========================================================
  // What I do this, I am just using 0x02 as a mode (4th MSB)
  // You have better way to do that in the SDK
  // so the nonce is [3 bytes key][1 byte mode = 02][20 bytes validator address, address 0 in this case][8 bytes nonce sequence]

  // Convert to hex string, pad to 64 chars (32 bytes)
  let hexNonce = userOp.nonce.toString(16).padStart(64, "0")
  // Replace the 4th byte (7th and 8th characters) with '02'
  hexNonce = `${hexNonce.substring(0, 6)}02${hexNonce.substring(8)}`
  //userOp.nonce = BigInt("0x" + hexNonce) // set the PREP mode nonce
  userOp.nonce = BigInt(`0x${hexNonce}`)

  // calculate the updated userOp hash
  const updatedHash = nexusAccount.getUserOpHash(userOp)
  // sign it
  // eoa signer will be the owner of the Nexus Prep at the recovered address
  const sig = await eoaSigner.signMessage({ message: { raw: updatedHash } })

  // encode the signature for the PREP mode
  // [bytes32 saltAndDelegation, bytes initData, bytes og signature]
  const prepSig = encodeAbiParameters(
    [{ type: "bytes32" }, { type: "bytes" }, { type: "bytes" }],
    [saltAndDelegation, initData, sig]
  )
  userOp.signature = prepSig

  // ========================================================
  // ===== Send the UserOp to the Entrypoint =====
  // ========================================================
  const packedUserOp = toPackedUserOperation(userOp)
  const balanceBefore = await getBalance(
    publicClient,
    someWallet.address,
    FREE_MINT_ERC20
  )

  const handleOpsHash = await someWalletClient.sendTransaction({
    authorizationList: [signedAuthorization],
    data: encodeFunctionData({
      abi: EntrypointAbi,
      functionName: "handleOps",
      args: [[packedUserOp], someWallet.address]
    }),
    to: ENTRYPOINT_V07_ADDRESS
  })

  const receipt = await publicClient.waitForTransactionReceipt({
    hash: handleOpsHash
  })

  const balanceAfter = await getBalance(
    publicClient,
    someWallet.address,
    FREE_MINT_ERC20
  )

  if (balanceAfter - balanceBefore !== parseEther("1")) {
    throw new Error("Balance has not changed properly")
  }
  console.log("Spin up a Nexus PREP + send ERC20", receipt.gasUsed)
}