From 6a11f4267ec54a5e224962d4a6687c73df6d114b Mon Sep 17 00:00:00 2001
From: aon <21188659+aon@users.noreply.github.com>
Date: Fri, 13 Dec 2024 12:53:38 -0300
Subject: [PATCH] feat: update cpamm

---
 contracts/src/CPAMM.sol   |  6 +++++-
 web/components/Header.tsx | 11 +++++++++--
 web/contracts/cpamm.ts    |  7 +++++++
 web/hooks/use-cpamm.ts    | 34 ++++++++++++++++++----------------
 4 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/contracts/src/CPAMM.sol b/contracts/src/CPAMM.sol
index 5a74491..bba4017 100644
--- a/contracts/src/CPAMM.sol
+++ b/contracts/src/CPAMM.sol
@@ -110,6 +110,10 @@ contract CPAMM {
         );
     }
 
+    function allowedToAddLiquidity() external view returns (bool) {
+      return userFeeTier[msg.sender] == 1 || userFeeTier[msg.sender] == 2;
+    }
+
     function addLiquidity(
         uint256 _amount0,
         uint256 _amount1
@@ -316,7 +320,7 @@ contract CPAMM {
 
     modifier onlyPremiumOrVip() {
         require(
-            userFeeTier[msg.sender] == 2 || userFeeTier[msg.sender] == 1,
+            userFeeTier[msg.sender] == 1 || userFeeTier[msg.sender] == 2,
             "Only VIP or premium users can call this function"
         );
         _;
diff --git a/web/components/Header.tsx b/web/components/Header.tsx
index 486fc9b..65b9aaf 100644
--- a/web/components/Header.tsx
+++ b/web/components/Header.tsx
@@ -1,5 +1,6 @@
 "use client";
 
+import { useEffect, useState } from "react";
 import Image from "next/image";
 import Link from "next/link";
 import { usePathname } from "next/navigation";
@@ -14,6 +15,12 @@ type HeaderMenuLink = {
 
 export const Header = () => {
   const { addLiquidityAllowed } = useCpamm();
+  const [showPool, setShowPool] = useState(false);
+
+  // Use it to update the state after hydration
+  useEffect(() => {
+    setShowPool(addLiquidityAllowed);
+  }, [addLiquidityAllowed]);
 
   return (
     <div className="flex h-[74px] justify-between z-20 w-full p-4">
@@ -23,9 +30,9 @@ export const Header = () => {
         </div>
         <span className="text-xl font-normal">Double Zero Swap</span>
       </Link>
-      <nav className="hidden md:flex lg:flex-nowrap px-1 gap-2">
+      <nav suppressHydrationWarning className="hidden md:flex lg:flex-nowrap px-1 gap-2">
         <HeaderMenuLink label="Swap" href="/" matches={["/"]} />
-        {addLiquidityAllowed && <HeaderMenuLink label="Pool" href="/pool" matches={["/pool", "/pool/add"]} />}
+        {showPool && <HeaderMenuLink label="Pool" href="/pool" matches={["/pool", "/pool/add"]} />}
       </nav>
       <div className="mr-4">
         <RainbowKitCustomConnectButton />
diff --git a/web/contracts/cpamm.ts b/web/contracts/cpamm.ts
index 9a06205..4e46ac1 100644
--- a/web/contracts/cpamm.ts
+++ b/web/contracts/cpamm.ts
@@ -19,6 +19,13 @@ export const CPAMM_ABI = [
     outputs: [{ name: "shares", type: "uint256", internalType: "uint256" }],
     stateMutability: "nonpayable",
   },
+  {
+    type: "function",
+    name: "allowedToAddLiquidity",
+    inputs: [],
+    outputs: [{ name: "", type: "bool", internalType: "bool" }],
+    stateMutability: "view",
+  },
   {
     type: "function",
     name: "balanceOf",
diff --git a/web/hooks/use-cpamm.ts b/web/hooks/use-cpamm.ts
index 61e27cf..bab8942 100644
--- a/web/hooks/use-cpamm.ts
+++ b/web/hooks/use-cpamm.ts
@@ -1,5 +1,5 @@
-import { useCallback, useMemo } from "react";
-import { useAccount, useReadContract, useSimulateContract, useWriteContract } from "wagmi";
+import { useCallback, useEffect, useMemo } from "react";
+import { useAccount, useReadContract, useWriteContract } from "wagmi";
 import { CPAMM_ABI, CPAMM_ADDRESS } from "~~/contracts/cpamm";
 
 const options = {
@@ -19,13 +19,10 @@ export default function useCpamm() {
     functionName: "balanceOf",
     args: [address ?? ""],
   });
-  const { error: simulateAddLiquidityError } = useSimulateContract({
+  const { data: allowedToAddLiquidity, refetch: refetchAllowedToAddLiquidity } = useReadContract({
     ...options,
-    functionName: "addLiquidity",
-    args: [1n, 1n],
-    query: {
-      retry: 1,
-    },
+    functionName: "allowedToAddLiquidity",
+    account: address ?? "",
   });
   const { data: fee, refetch: refetchFee } = useReadContract({
     ...options,
@@ -42,21 +39,26 @@ export default function useCpamm() {
     refetchLiquidity();
     refetchFee();
     refetchRemainingDailyAllowance();
-  }, [refetchLiquidity, refetchFee, refetchRemainingDailyAllowance]);
+    refetchAllowedToAddLiquidity();
+  }, [refetchLiquidity, refetchFee, refetchRemainingDailyAllowance, refetchAllowedToAddLiquidity]);
+
+  // Refetch all when the address changes
+  useEffect(() => {
+    refetchAll();
+  }, [address, refetchAll]);
 
   const addLiquidityAllowed = useMemo(() => {
     // Get the last known state from localStorage, default to true if not set
-    const lastKnownState = localStorage.getItem("addLiquidityAllowed") === "false" ? false : true;
+    const lastKnownState = globalThis?.localStorage?.getItem("addLiquidityAllowed") === "false" ? false : true;
 
-    if (simulateAddLiquidityError === null) {
+    // If the simulateAddLiquidityError is loading, then we are on the first render
+    if (allowedToAddLiquidity === undefined) {
       return lastKnownState;
     }
 
-    const currentState = !simulateAddLiquidityError.message.includes("Internal JSON-RPC error.");
-    // Store the new state
-    localStorage.setItem("addLiquidityAllowed", currentState.toString());
-    return currentState;
-  }, [simulateAddLiquidityError]);
+    globalThis?.localStorage?.setItem("addLiquidityAllowed", allowedToAddLiquidity.toString());
+    return allowedToAddLiquidity;
+  }, [allowedToAddLiquidity]);
 
   return {
     daiPoolLiquidity: liquidity?.[0],