assetx/webapp/hooks/usePoolWithdraw.ts

import { useState, useEffect } from 'react'
import { useAccount, useWriteContract, useWaitForTransactionReceipt } from 'wagmi'
import { parseUnits } from 'viem'
import { abis, getContractAddress } from '@/lib/contracts'
import { handleContractCatchError, isValidAmount, parseContractError } from '@/lib/errors'
import { SWAP_GAS_LIMIT } from '@/lib/constants'

type WithdrawStatus = 'idle' | 'approving' | 'approved' | 'withdrawing' | 'success' | 'error'

/**
 * Hook for removing liquidity from YT Pool
 * User burns YT LP tokens and receives USDC back
 */
export function usePoolWithdraw() {
  const { address, chainId } = useAccount()
  const [status, setStatus] = useState<WithdrawStatus>('idle')
  const [error, setError] = useState<string | null>(null)
  const [approveHash, setApproveHash] = useState<`0x${string}` | undefined>()
  const [withdrawHash, setWithdrawHash] = useState<`0x${string}` | undefined>()

  const { writeContractAsync: approveWrite } = useWriteContract()
  const { writeContractAsync: withdrawWrite } = useWriteContract()

  const { isLoading: isApproving, isSuccess: isApproveSuccess, isError: isApproveError, error: approveReceiptError } = useWaitForTransactionReceipt({
    hash: approveHash,
  })

  const { isLoading: isWithdrawing, isSuccess: isWithdrawSuccess, isError: isWithdrawError, error: withdrawReceiptError } = useWaitForTransactionReceipt({
    hash: withdrawHash,
  })

  const executeApprove = async (amount: string) => {
    if (status !== 'idle') return false
    if (!chainId) {
      setError('Please connect your wallet')
      return false
    }
    if (!isValidAmount(amount)) {
      setError('Invalid amount')
      return false
    }

    try {
      setStatus('approving')
      setError(null)

      const ytLPAddress = getContractAddress('YTLPToken', chainId)
      const rewardRouterAddress = getContractAddress('YTRewardRouter', chainId)

      if (!ytLPAddress || !rewardRouterAddress) {
        throw new Error('Contract not deployed on this chain')
      }

      const amountInWei = parseUnits(amount, 18)

      const hash = await approveWrite({
        address: ytLPAddress,
        abi: abis.YTLPToken,
        functionName: 'approve',
        args: [rewardRouterAddress, amountInWei],
      })

      setApproveHash(hash)
      return true
    } catch (err: unknown) {
      handleContractCatchError(err, 'Approval failed', setError, setStatus as (s: string) => void)
      return false
    }
  }

  const executeWithdraw = async (
    tokenOutAddress: string,
    tokenOutDecimals: number,
    amount: string,
    minOut: string = '0'
  ) => {
    if (status === 'withdrawing' || status === 'success' || status === 'error') return false
    if (!chainId || !address) {
      setError('Please connect your wallet')
      return false
    }
    if (!isValidAmount(amount)) {
      setError('Invalid amount')
      return false
    }

    try {
      setStatus('withdrawing')
      setError(null)

      const rewardRouterAddress = getContractAddress('YTRewardRouter', chainId)

      if (!tokenOutAddress || !rewardRouterAddress) {
        throw new Error('Contract not deployed on this chain')
      }

      const amountInWei = parseUnits(amount, 18)
      const minOutWei = parseUnits(minOut, tokenOutDecimals)

      const hash = await withdrawWrite({
        address: rewardRouterAddress,
        abi: abis.YTRewardRouter,
        functionName: 'removeLiquidity',
        args: [tokenOutAddress as `0x${string}`, amountInWei, minOutWei, address],
        gas: SWAP_GAS_LIMIT,
      })

      setWithdrawHash(hash)
      return true
    } catch (err: unknown) {
      handleContractCatchError(err, 'Transaction failed', setError, setStatus as (s: string) => void)
      return false
    }
  }

  const executeApproveAndWithdraw = async (
    tokenOutAddress: string,
    tokenOutDecimals: number,
    amount: string,
    minOut: string = '0'
  ) => {
    if (status !== 'idle') return
    const approved = await executeApprove(amount)
    if (!approved) return

    await new Promise(resolve => setTimeout(resolve, 2000))

    await executeWithdraw(tokenOutAddress, tokenOutDecimals, amount, minOut)
  }

  const reset = () => {
    setStatus('idle')
    setError(null)
    setApproveHash(undefined)
    setWithdrawHash(undefined)
  }

  useEffect(() => {
    if (isApproveSuccess && status === 'approving') {
      setStatus('approved')
    }
  }, [isApproveSuccess, status])

  useEffect(() => {
    if (isApproveError && status === 'approving') {
      setError(parseContractError(approveReceiptError))
      setStatus('error')
    }
  }, [isApproveError, status, approveReceiptError])

  useEffect(() => {
    if (isWithdrawSuccess && status === 'withdrawing') {
      setStatus('success')
    }
  }, [isWithdrawSuccess, status])

  useEffect(() => {
    if (isWithdrawError && status === 'withdrawing') {
      setError(parseContractError(withdrawReceiptError))
      setStatus('error')
    }
  }, [isWithdrawError, status, withdrawReceiptError])

  return {
    status,
    error,
    isLoading: isApproving || isWithdrawing,
    approveHash,
    withdrawHash,
    executeApprove,
    executeWithdraw,
    executeApproveAndWithdraw,
    reset,
  }
}