brahma-trade-widget/src/components/morphoStrategy/store.ts

A React component for trade automation within the Brahma ecosystem.

import { create } from "zustand"
import { ConsoleKit, TaskIdStatusType } from "brahma-console-kit"
import { signTypedData } from "@wagmi/core"
import { Address, fromHex } from "viem"

import {
  AiResponseToDetect,
  MorphoTokensOptions,
  MorphoVault,
  MorphoYieldAutomationLog,
  MorphoYieldUserPosition,
} from "./types"
import { SupportedChainIds, TAsset } from "@/types"
import { fetchAndFilterAssetBalances, formatUnits } from "@/utils"
import { BASE_CHAIN_ID, URL_API_KEY_PROD } from "@/constants"
import { baseWagmiConfig as wagmiConfig } from "@/wagmi"
import { dispatchToast } from "../shared/components"
import {
  getBaseBackendUrl,
  getFormattedDate,
  getMorphoYieldExecutorId,
  getTokenData,
  waitFor,
} from "./utils"
import { MORPHO_YIELD_AUTOMATION_ACTIVE_STATE } from "./constants"
import localStorageService from "../shared/services/localStoage"
import { MORPHO_STRATEGY_TOKENS } from "../shared/constants"
import { base } from "viem/chains"
import {
  callIndexer,
  fetchMorphoYieldAutomation,
  fetchMorphoYieldLogs,
  fetchMorphoYieldVaultPosition,
  fetchMorphoYieldVaults,
} from "../shared/api"

const kit = new ConsoleKit(URL_API_KEY_PROD, getBaseBackendUrl())

type State = {
  loading: boolean
  isDepositModalOpen: boolean
  strategyDeployedOn: Address | null
  deploymentStatus: {
    status: TaskIdStatusType
    taskId: string
  } | null
  preComputedConsoleAddress: Address | null
  balances: {
    data: TAsset[]
    loading: boolean
  }
  eoaBalances: {
    data: TAsset[]
    loading: boolean
  }
  feeEstimateSignature: string | null
  feeEstimate: string | null
  automationConsoleAddress: Record<
    MorphoTokensOptions,
    { address: Address | null; loading: boolean }
  >
  userPositions: Record<
    MorphoTokensOptions,
    {
      loading: boolean
      data: MorphoYieldUserPosition[]
      consoleAddress: Address
    }
  >
  loopFetchingPositions: boolean
  morphoVaults: {
    loading: boolean
    vaultMapping: Map<`0x${string}`, MorphoVault>
    vaults: MorphoVault[]
  }
  hasNftData: {
    loading: boolean
    value: boolean
  }
  showPageSpinner: boolean
  paramsDetectedInAiRes: null | AiResponseToDetect
}

type Actions = {
  setParamsDetectedInAiRes: (value: null | AiResponseToDetect) => void

  fetchPreComputedConsoleAddress: (
    owner: Address,
    chainId: SupportedChainIds,
    feeToken: Address,
  ) => Promise<void>
  generateAndDeploySubAccount: (
    eoa: Address,
    chainId: SupportedChainIds,
    feeToken: Address,
    feeEstimate: string,
    tokens: Address[],
    amounts: string[],
    duration: number,
    tokenInputs: Record<Address, string>,
    tokenLimits: Record<Address, string>,
    preferredVaults: Address[],
  ) => Promise<void>
  fetchDeploymentStatus: (
    taskId: string,
    token: MorphoTokensOptions,
    eoa: Address,
  ) => Promise<void>
  fetchPreComputedConsoleBalances: (
    assets: TAsset[],
    addressToDetect?: Address,
  ) => Promise<boolean>
  fetchEoaAssets: (eoa: Address, assets: TAsset[]) => Promise<void>
  fetchMorphoUserPositions: (
    consoleAddress: Address,
    chainId: SupportedChainIds,
    token: MorphoTokensOptions,
    isPollingExecution?: boolean,
  ) => Promise<boolean>
  fetchMorphoVaults: () => Promise<void>

  setNftData: (value: boolean, loading?: boolean) => void
  setPageSpinnerState: (value: boolean) => void
  resetUserPostion: (token: MorphoTokensOptions) => void
  setAutomationConsoleAddress: (
    consoleAddress: Address | null,
    token: MorphoTokensOptions,
  ) => void
  resetPreComputedConsoleAddress: () => void
  resetAutomationConsoleAddress: (token: MorphoTokensOptions) => void
  openDepositModal: () => void
  closeDepositModal: () => void
}

const useStore = create<State & Actions>((set, get) => ({
  isDepositModalOpen: false,
  paramsDetectedInAiRes: null,
  showPageSpinner: false,
  loading: false,
  hasNftData: {
    loading: true,
    value: false,
  },
  automationConsoleAddress: {
    usdc: {
      loading: true,
      address: null,
    },
    weth: {
      loading: true,
      address: null,
    },
  },
  deploymentStatus: null,
  balances: {
    data: [],
    loading: false,
  },
  morphoVaults: {
    loading: true,
    vaults: [],
    vaultMapping: new Map(),
  },
  userPositions: {
    weth: { data: [], loading: true, consoleAddress: "0x" as Address },
    usdc: { data: [], loading: true, consoleAddress: "0x" as Address },
  },
  strategyDeployedOn: null,
  eoaBalances: {
    data: [],
    loading: false,
  },
  preComputedConsoleAddress: null,
  feeEstimate: null,
  feeEstimateSignature: null,
  loopFetchingPositions: false,

  openDepositModal: () => set({ isDepositModalOpen: true }),
  closeDepositModal: () => set({ isDepositModalOpen: false }),
  setParamsDetectedInAiRes: (value: null | AiResponseToDetect) =>
    set({ paramsDetectedInAiRes: value }),

  resetAutomationConsoleAddress: (token: MorphoTokensOptions) => {
    set((prev) => ({
      automationConsoleAddress: {
        ...prev.automationConsoleAddress,
        [token]: {
          address: null,
          loading: true,
        },
      },
    }))
  },
  setAutomationConsoleAddress: (
    consoleAddress: Address | null,
    token: MorphoTokensOptions,
  ) => {
    set((prev) => ({
      automationConsoleAddress: {
        ...prev.automationConsoleAddress,
        [token]: {
          address: consoleAddress,
          loading: false,
        },
      },
    }))
  },
  resetUserPostion: (token: MorphoTokensOptions) => {
    set((prev) => ({
      userPositions: {
        ...prev.userPositions,
        [token]: { data: [], loading: false, consoleAddress: "0x" as Address },
        strategyDeployedOn: null,
      },
    }))
  },
  setPageSpinnerState: (value: boolean) => set({ showPageSpinner: value }),
  setNftData: (value: boolean, loading?: boolean) =>
    set({
      hasNftData: {
        value,
        loading: typeof loading === "boolean" ? loading : false,
      },
    }),

  fetchMorphoVaults: async () => {
    try {
      const response = (await fetchMorphoYieldVaults()) || []
      const baseVaults = response.filter(
        (vault) => vault.asset.chain.id === base.id,
      )

      // Get the best performing vault with with min TVL threshold as 100K USD
      const getFirstVault = (tokenAddress: string) =>
        baseVaults
          .filter(
            (vault) =>
              vault.asset.address.toLowerCase() ===
                tokenAddress.toLowerCase() && vault.supply.usd > 100_000,
          )
          .sort((a, b) => b.state.netApy - a.state.netApy)[0]

      const firstUSDCVault = getFirstVault(MORPHO_STRATEGY_TOKENS.usdc.address)
      const firstWETHVault = getFirstVault(MORPHO_STRATEGY_TOKENS.weth.address)

      const filteredOrderedVaults = [firstUSDCVault, firstWETHVault].filter(
        Boolean,
      ) as MorphoVault[]

      const vaultMapping = filteredOrderedVaults.reduce(
        (acc, { asset, ...data }) => {
          acc.set(asset.address.toLowerCase() as Address, { asset, ...data })
          return acc
        },
        new Map<Address, MorphoVault>(),
      )

      set({
        morphoVaults: { vaultMapping, vaults: baseVaults, loading: false },
      })
    } catch (err) {
      set({
        morphoVaults: { vaultMapping: new Map(), loading: false, vaults: [] },
      })
    }
  },

  resetPreComputedConsoleAddress: () => {
    set({
      preComputedConsoleAddress: null,
    })
  },

  fetchPreComputedConsoleAddress: async (owner, chainId, feeToken) => {
    set((state) => ({ ...state, loading: true }))
    try {
      const data = await kit.publicDeployer.fetchPreComputeData(
        owner,
        chainId,
        feeToken,
      )

      console.log({ data, owner, chainId, feeToken })
      if (
        !data ||
        !data.feeEstimate ||
        !data.feeEstimateSignature ||
        !data.precomputedAddress
      ) {
        throw new Error("Invalid data received from fetchPreComputeAddress")
      }

      set((state) => ({
        ...state,
        loading: false,
        feeEstimate: data.feeEstimate,
        feeEstimateSignature: data.feeEstimateSignature,
        preComputedConsoleAddress: data.precomputedAddress,
      }))
    } catch (err: any) {
      console.log(`Error fetching precompute address: ${err}`)
      dispatchToast({
        id: "fetch-precompute-address-error",
        title: "Error fetching precompute address",
        description: {
          value:
            err?.message ||
            "An error occurred while fetching precompute address",
        },
        type: "error",
      })
      set((state) => ({ ...state, loading: false })) // Ensure loading is set to false on error
    }
  },
  generateAndDeploySubAccount: async (
    eoa,
    chainId,
    feeToken,
    feeEstimate,
    tokens,
    amounts,
    duration,
    tokenInputs,
    tokenLimits,
    preferredVaults,
  ) => {
    const { preComputedConsoleAddress, feeEstimateSignature } = get()

    const HARDCODED_REGISTRY_ID =
      getMorphoYieldExecutorId(chainId) ||
      "8fa8bacb-7be0-4232-8dce-d080e1e56540"

    if (!preComputedConsoleAddress) {
      dispatchToast({
        id: "fetch-precomputed-account-address-error",
        title: "Error fetching precomputed account address",
        description: {
          value: "Precomputed account address not found",
        },
        type: "error",
      })
      return
    }

    set((state) => ({ ...state, loading: true }))
    try {
      // Generate Automation SubAccount
      const durationForPayload = duration > 3600 ? duration - 3600 : duration

      const generateData =
        await kit.publicDeployer.generateAutomationSubAccount(
          eoa,
          preComputedConsoleAddress,
          chainId,
          HARDCODED_REGISTRY_ID,
          feeToken,
          feeEstimate,
          tokens,
          amounts,
          {
            duration: durationForPayload,
            tokenInputs: tokenInputs,
            tokenLimits: tokenLimits,
          },
          {
            baseToken: feeToken,
            every: duration.toString(),
            type: "EARN",
            preferredVaults,
          },
        )

      if (!generateData || !feeEstimateSignature) {
        dispatchToast({
          id: "fetch-signature-error",
          title: "Error fetching signature",
          description: {
            value: "An error occurred while fetching signature",
          },
          type: "error",
        })
        return
      }

      const {
        signaturePayload: { domain, message, types, primaryType },
        subAccountPolicyCommit,
        subscriptionDraftID,
      } = generateData

      const signature = await signTypedData(wagmiConfig, {
        domain: {
          verifyingContract: domain.verifyingContract,
          chainId: fromHex(domain.chainId as Address, "number"),
        },
        types,
        primaryType,
        message,
      })

      if (!signature) {
        dispatchToast({
          id: "upgrade-console-error",
          type: "error",
          title: "Error",
          description: {
            value: "User rejected the transaction",
          },
        })
        return
      }

      // Deploy Brahma Account
      const deployData = await kit.publicDeployer.deployBrahmaAccount(
        eoa,
        chainId,
        HARDCODED_REGISTRY_ID,
        subscriptionDraftID,
        subAccountPolicyCommit,
        feeToken,
        tokens,
        amounts,
        signature, // Use the signature obtained from the previous step
        feeEstimateSignature,
        feeEstimate,
        {},
      )

      if (!deployData) {
        dispatchToast({
          id: "deploy-account-error",
          title: "Error deploying account and sub-account",
          description: {
            value: "An error occurred during deployment",
          },
          type: "error",
        })
        return
      }

      dispatchToast({
        id: "deployment-status-pending",
        title: "Deployment Status",
        description: {
          value: "The deployment is currently pending.",
        },
        type: "loading",
      })

      // Update state with taskId and set loading to false
      set((state) => ({
        ...state,
        deploymentStatus: {
          status: "pending",
          taskId: deployData.taskId,
        },
        signature, // Update the state with the signature
        loading: false,
        strategyDeployedOn: preComputedConsoleAddress,
      }))
    } catch (err: any) {
      console.error("Error in generate and deploy sub-account:", err)
      dispatchToast({
        id: "generate-deploy-error",
        title: "Error in generate and deploy sub-account",
        description: {
          value: err?.message || "An error occurred during the process",
        },
        type: "error",
      })
    }
  },
  fetchDeploymentStatus: async (
    taskId,
    token: MorphoTokensOptions,
    eoa: Address,
  ) => {
    const {
      strategyDeployedOn,
      userPositions,
      fetchMorphoUserPositions,
      setAutomationConsoleAddress,
    } = get()
    try {
      const data = await kit.publicDeployer.fetchDeploymentStatus(taskId)
      set({ showPageSpinner: true })
      if (!data) {
        dispatchToast({
          id: "fetch-deployment-status-error",
          title: "Error fetching deployment status",
          description: {
            value: "An error occurred while fetching deployment status",
          },
          type: "error",
        })
        set({ showPageSpinner: false })

        return
      }

      console.log({ deploymentData: data })
      switch (data.status) {
        case "pending":
          dispatchToast({
            id: "deployment-status-pending",
            title: "Deployment Status",
            description: {
              value: "The deployment is currently pending.",
            },
            type: "loading",
          })
          break
        case "executing":
          dispatchToast({
            id: "deployment-status-executing",
            title: "Deployment Status",
            description: {
              value: "The deployment is currently executing.",
            },
            type: "loading",
          })
          break
        case "cancelled":
          dispatchToast({
            id: "deployment-status-cancelled",
            title: "Deployment Status",
            description: {
              value: "The deployment has been cancelled.",
            },
            type: "error",
          })
          break
        case "successful":
          dispatchToast({
            id: "deployment-status-successful",
            title: "Deployment Status",
            description: {
              value:
                "Your account and sub-account have been deployed successfully",
            },
            type: "success",
          })
          strategyDeployedOn &&
            localStorageService.setDeployedMorphoStrategyConsole(
              strategyDeployedOn as Address,
              token,
              eoa,
            )

          break
        case "failed":
          {
            dispatchToast({
              id: "deployment-status-failed",
              title: "Deployment Status",
              description: {
                value: "The deployment has failed.",
              },
              type: "error",
            })
            set({ showPageSpinner: false })
          }
          break
        default:
          {
            dispatchToast({
              id: "deployment-status-unknown",
              title: "Deployment Status",
              description: {
                value: "The deployment status is unknown.",
              },
              type: "error",
            })
            set({ showPageSpinner: false })
          }
          break
      }

      set((state) => ({
        ...state,
        deploymentStatus: {
          status: data.status,
          taskId: taskId,
        },
      }))

      // add a loop to continue fetch positions
      if (data.status === "successful" && strategyDeployedOn) {
        await callIndexer(data?.outputTransactionHash || "")

        set({ loopFetchingPositions: true })
        let isPositionDeployed = false

        while (!isPositionDeployed) {
          const isPositionDeployedResponse = await fetchMorphoUserPositions(
            strategyDeployedOn,
            BASE_CHAIN_ID,
            token,
          )
          isPositionDeployed = isPositionDeployedResponse
          await waitFor(3000)
          console.log("LOOP FETCHING USER POSITIONS")
        }
        setAutomationConsoleAddress(strategyDeployedOn, token)
        set({
          loopFetchingPositions: false,
          showPageSpinner: false,
          balances: {
            data: [],
            loading: false,
          },
        })
      }
    } catch (err: any) {
      set({ loopFetchingPositions: false })
      set({ showPageSpinner: false })
      console.error("Error fetching deployment status:", err)
      dispatchToast({
        id: "fetch-deployment-status-error",
        title: "Error fetching deployment status",
        description: {
          value:
            err?.message ||
            "An error occurred while fetching deployment status",
        },
        type: "error",
      })
    }
  },
  fetchPreComputedConsoleBalances: async (assets, addressToDetect) => {
    const { preComputedConsoleAddress } = get()

    if (!preComputedConsoleAddress) {
      dispatchToast({
        id: "fetch-precomputed-account-balances-error",
        title: "Error fetching precomputed account balances",
        description: {
          value: "Precomputed account address not found",
        },
        type: "error",
      })
      return false
    }

    set((state) => ({ balances: { ...state.balances, loading: true } }))

    try {
      const filteredUserAssets = await fetchAndFilterAssetBalances(
        preComputedConsoleAddress,
        assets,
      )

      set({
        balances: {
          data: filteredUserAssets,
          loading: false,
        },
      })

      if (!addressToDetect) return false

      const targetToken = filteredUserAssets.find(
        (asset) =>
          asset.address.toLowerCase() === addressToDetect.toLowerCase(),
      )

      if (!targetToken) return false

      const targetTokenHasAmount = targetToken?.balanceOf?.value !== BigInt(0)
      console.log(targetTokenHasAmount, targetToken)
      return targetTokenHasAmount
    } catch (err) {
      console.error("Error fetching precomputed console balances:", err)
      set((state) => ({ balances: { ...state.balances, loading: false } }))
      return false
    }
  },

  fetchEoaAssets: async (eoa, assets) => {
    set((state) => ({ balances: { ...state.balances, loading: true } }))

    try {
      const filteredUserAssets = await fetchAndFilterAssetBalances(eoa, assets)

      set({
        eoaBalances: {
          data: filteredUserAssets,
          loading: false,
        },
      })
    } catch (err) {
      console.error("error on fetching eoa assets", err)
      set({ balances: { data: [], loading: false } })
    }
  },

  fetchMorphoUserPositions: async (
    consoleAddress: Address,
    chainId: SupportedChainIds,
    token: MorphoTokensOptions,
    isPollingExecution?: boolean,
  ) => {
    const { morphoVaults } = get()
    const automations = await fetchMorphoYieldAutomation(
      consoleAddress,
      chainId,
    )
    const morphoExecutorId = getMorphoYieldExecutorId(chainId)

    // filter by registery ID and staus
    const activeMorphoAutomations = automations.filter(
      (automation) =>
        automation.registryId.toLowerCase() ===
          (morphoExecutorId || "").toLowerCase() &&
        automation.status === MORPHO_YIELD_AUTOMATION_ACTIVE_STATE,
    )
    // for every automation, fetch vault info and append in the vault fields
    const automationWithVaultInfo: MorphoYieldUserPosition[] =
      await Promise.all(
        activeMorphoAutomations.map(async (_automation) => {
          // vault data
          const vaultResponse =
            (await fetchMorphoYieldVaultPosition(
              _automation.subAccountAddress,
            )) || []

          let vaultPositionData = null

          const baseToken = Object.entries(_automation.tokenInputs).map(
            ([address, value]) => ({
              address,
              value: Number(value), // Convert the value to a number if needed
            }),
          )[0]

          // we pick only first object, asked by BE.
          if (vaultResponse.length > 0) {
            vaultPositionData = vaultResponse[0]
          }

          const nonFormattedAmount = vaultPositionData
            ? vaultPositionData.assets
            : baseToken.address
              ? _automation.tokenInputs[baseToken.address as Address]
              : "0"

          const asset = getTokenData(baseToken.address as Address, chainId)

          const amount = formatUnits(nonFormattedAmount || "0", asset.decimals)

          // automation logs
          const logs = (await fetchMorphoYieldLogs(_automation.id)) || []

          const history: MorphoYieldAutomationLog[] = logs.map((log) => {
            const prevVault = log.metadata.transitionState.prev
            const currentVault = log.metadata.transitionState.current

            const name = prevVault ? "Rebalance" : "Deposit"

            const fromVault = prevVault
              ? prevVault.targetVault
              : currentVault.targetVault

            const toVault = prevVault ? currentVault.targetVault : null

            const time = getFormattedDate(new Date(log.createdAt), {
              year: "numeric",
              month: "short",
              day: "numeric",
              hour: "numeric",
              minute: "2-digit",
            })

            return {
              id: log.id,
              name,
              toVault,
              fromVault,
              amount: formatUnits(currentVault.inputAmount, asset.decimals),
              asset,
              time,
              outputTxHash: log.outputTxHash,
              underlyingVault: currentVault.targetVault,
            }
          })

          const targetVault = vaultPositionData
            ? morphoVaults.vaults.find(
                (vault) =>
                  vault.address.toLowerCase() ===
                  vaultPositionData.vault.address.toLowerCase(),
              )
            : null

          const vaultName = targetVault?.name || ""
          const vaultCurators = targetVault?.metadata?.curators || []

          const targetCurator =
            vaultCurators.length > 0 ? vaultCurators[0] : null

          return {
            ..._automation,
            // vaultPositionData,
            vaultPositionData: vaultPositionData
              ? { ...vaultPositionData, vaultName, curator: targetCurator }
              : null,
            history,
            tokenData: {
              asset,
              amount,
            },
          }
        }),
      )
    console.log({ automationWithVaultInfo })

    const { userPositions } = get()
    if (userPositions[token].data.length === 0 && isPollingExecution) {
      console.log("STOPPING OVERRIDING EMPTY STATE")
      return true // stop overriding empty state during polling}
    }

    set((prev) => ({
      userPositions: {
        ...prev.userPositions,
        [token]: {
          loading: false,
          data: automationWithVaultInfo,
          consoleAddress: consoleAddress,
        },
      },
    }))

    return automationWithVaultInfo.length > 0
  },
}))

export default useStore