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('idle') const [error, setError] = useState(null) const [pendingAmount, setPendingAmount] = useState('') const [requestId, setRequestId] = useState(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, } }