assetx/webapp/hooks/useWithdraw.ts

import { useState, useCallback, useEffect } from 'react'
import { useAccount, useWriteContract, useWaitForTransactionReceipt } from 'wagmi'
import { parseUnits } from 'viem'
import { abis } from '@/lib/contracts'
import { Token } from '@/lib/api/tokens'
import { handleContractCatchError, isValidAmount, parseContractError } from '@/lib/errors'
import { DEFAULT_GAS_LIMIT } from '@/lib/constants'

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

export function useWithdraw(ytToken: Token | undefined) {
  const { address } = useAccount()
  const [status, setStatus] = useState<WithdrawStatus>('idle')
  const [error, setError] = useState<string | null>(null)
  const [pendingAmount, setPendingAmount] = useState<string>('')
  const [requestId, setRequestId] = useState<string | null>(null)

  const {
    writeContractAsync: approveWrite,
    data: approveHash,
    isPending: isApprovePending,
    reset: resetApprove,
  } = useWriteContract()

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

  const {
    writeContractAsync: withdrawWrite,
    data: withdrawHash,
    isPending: isWithdrawPending,
    reset: resetWithdraw,
  } = useWriteContract()

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

  const executeApprove = useCallback(async (amount: string) => {
    if (status !== 'idle') return false
    if (!address || !ytToken?.contractAddress || !amount) {
      setError('Missing required parameters')
      return false
    }
    if (!isValidAmount(amount)) {
      setError('Invalid amount')
      return false
    }

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

      const decimals = ytToken.onChainDecimals ?? ytToken.decimals
      const amountInWei = parseUnits(amount, decimals)

      await approveWrite({
        address: ytToken.contractAddress as `0x${string}`,
        abi: abis.YTToken,
        functionName: 'approve',
        args: [ytToken.contractAddress as `0x${string}`, amountInWei],
      })

      return true
    } catch (err: unknown) {
      handleContractCatchError(err, 'Approve failed', setError, setStatus as (s: string) => void)
      return false
    }
  }, [address, approveWrite, ytToken, status])

  const executeWithdraw = useCallback(async (amount: string) => {
    if (status === 'withdrawing' || status === 'success' || status === 'error') return false
    if (!address || !ytToken?.contractAddress || !amount) {
      setError('Missing required parameters')
      return false
    }
    if (!isValidAmount(amount)) {
      setError('Invalid amount')
      return false
    }

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

      const decimals = ytToken.onChainDecimals ?? ytToken.decimals
      const amountInWei = parseUnits(amount, decimals)

      await withdrawWrite({
        address: ytToken.contractAddress as `0x${string}`,
        abi: abis.YTToken,
        functionName: 'withdrawYT',
        args: [amountInWei],
        gas: DEFAULT_GAS_LIMIT,
      })

      return true
    } catch (err: unknown) {
      handleContractCatchError(err, 'Withdraw failed', setError, setStatus as (s: string) => void)
      return false
    }
  }, [address, withdrawWrite, ytToken, status])

  const executeApproveAndWithdraw = useCallback(async (amount: string) => {
    if (status !== 'idle') return
    setPendingAmount(amount)
    const approveSuccess = await executeApprove(amount)
    if (!approveSuccess) return
  }, [executeApprove, status])

  // Auto-execute withdraw when approve is successful
  useEffect(() => {
    if (isApproveSuccess && status === 'approving' && pendingAmount) {
      setStatus('approved')
      executeWithdraw(pendingAmount)
    }
  }, [isApproveSuccess, status, pendingAmount, executeWithdraw])

  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])

  const reset = useCallback(() => {
    setStatus('idle')
    setError(null)
    setPendingAmount('')
    setRequestId(null)
    resetApprove()
    resetWithdraw()
  }, [resetApprove, resetWithdraw])

  return {
    status,
    error,
    isLoading: isApprovePending || isApproveConfirming || isWithdrawPending || isWithdrawConfirming,
    approveHash,
    withdrawHash,
    requestId,
    executeApprove,
    executeWithdraw,
    executeApproveAndWithdraw,
    reset,
  }
}