trading_bot_v4/lib/drift/withdraw.ts

/**
 * Withdrawal Execution Service
 * 
 * Handles actual withdrawal from Drift Protocol to wallet
 */

import BN from 'bn.js'
import { PublicKey, Keypair } from '@solana/web3.js'
import { getAssociatedTokenAddress, TOKEN_PROGRAM_ID } from '@solana/spl-token'
import { initializeDriftService } from './client'
import bs58 from 'bs58'

export interface WithdrawalResult {
  success: boolean
  amount?: number
  signature?: string
  error?: string
}

export async function withdrawFromDrift(
  amountUSD: number,
  destinationWallet?: string
): Promise<WithdrawalResult> {
  try {
    console.log(`💸 Initiating withdrawal: $${amountUSD.toFixed(2)}`)
    
    const driftService = await initializeDriftService()
    
    // Get USDC spot market index (typically 0 for USDC)
    const usdcMarketIndex = 0
    
    // Convert USD amount to token amount (USDC has 6 decimals)
    // $50.25 → 50,250,000 (raw token amount with 6 decimals)
    const rawAmount = Math.floor(amountUSD * 1_000_000)
    console.log(`📊 Raw amount before BN: ${rawAmount} (type: ${typeof rawAmount})`)
    const tokenAmount = new BN(rawAmount.toString())
    
    console.log(`📊 Withdrawal details:`)
    console.log(`   Amount: $${amountUSD.toFixed(2)} USDC`)
    console.log(`   Token amount: ${tokenAmount.toString()} (raw with 6 decimals)`)
    console.log(`   Market index: ${usdcMarketIndex}`)
    
    // Get destination address
    // Need to get the Associated Token Account (ATA) for USDC, not just the wallet address
    let walletPublicKey: PublicKey
    let destination: PublicKey
    try {
      if (destinationWallet) {
        console.log(`🔑 Using provided destination wallet`)
        walletPublicKey = new PublicKey(destinationWallet)
      } else {
        console.log(`🔑 Deriving destination from private key`)
        // Derive public key from private key (same wallet as trader)
        const privateKeyStr = process.env.DRIFT_WALLET_PRIVATE_KEY
        if (!privateKeyStr) {
          throw new Error('DRIFT_WALLET_PRIVATE_KEY environment variable not set')
        }
        console.log(`🔑 Private key length: ${privateKeyStr.length}`)
        const privateKeyBytes = privateKeyStr.startsWith('[') 
          ? Uint8Array.from(JSON.parse(privateKeyStr))
          : bs58.decode(privateKeyStr)
        console.log(`🔑 Private key bytes length: ${privateKeyBytes.length}`)
        const keypair = Keypair.fromSecretKey(privateKeyBytes)
        walletPublicKey = keypair.publicKey
        console.log(`🔑 Derived wallet public key: ${walletPublicKey.toString()}`)
      }
      
      // USDC mint address on Solana mainnet
      const usdcMint = new PublicKey('EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v')
      console.log(`🔑 USDC Mint: ${usdcMint.toString()}`)
      console.log(`🔑 Wallet public key: ${walletPublicKey.toString()}`)
      console.log(`🔑 About to call getAssociatedTokenAddress...`)
      
      // Get the Associated Token Account for USDC
      destination = await getAssociatedTokenAddress(
        usdcMint,
        walletPublicKey,
        false, // allowOwnerOffCurve
        TOKEN_PROGRAM_ID
      )
      
      console.log(`✅ Got ATA successfully!`)
      console.log(`   Wallet: ${walletPublicKey.toString()}`)
      console.log(`   USDC ATA: ${destination.toString()}`)
      
    } catch (keyError: any) {
      console.error(`❌ Failed to get destination address:`, keyError)
      throw new Error(`Failed to derive destination address: ${keyError?.message || 'Unknown error'}`)
    }
    
    console.log(`   Destination: ${destination.toString()}`)
    
    // Execute withdrawal via Drift SDK
    // withdraw(amount, marketIndex, associatedTokenAddress, reduceOnly, subAccountId, txParams, updateFuel)
    const signature = await driftService.getClient().withdraw(
      tokenAmount,
      usdcMarketIndex,
      destination,
      false, // reduceOnly
      0, // subAccountId
      undefined, // txParams
      false // updateFuel
    )
    
    console.log(`✅ Withdrawal successful!`)
    console.log(`   Transaction: ${signature}`)
    console.log(`   Explorer: https://solscan.io/tx/${signature}`)
    
    // Confirm transaction
    console.log('⏳ Confirming transaction on-chain...')
    const connection = driftService.getConnection()
    const confirmation = await connection.confirmTransaction(signature, 'confirmed')
    
    if (confirmation.value.err) {
      throw new Error(`Transaction failed: ${JSON.stringify(confirmation.value.err)}`)
    }
    
    console.log('✅ Transaction confirmed on-chain')
    
    return {
      success: true,
      amount: amountUSD,
      signature,
    }
  } catch (error: any) {
    console.error('❌ Withdrawal failed:', error)
    return {
      success: false,
      error: error.message || 'Unknown error during withdrawal',
    }
  }
}

export async function calculateWithdrawalAmount(): Promise<{
  availableProfit: number
  withdrawalAmount: number
  safeToWithdraw: boolean
  reason?: string
}> {
  try {
    const driftService = await initializeDriftService()
    const health = await driftService.getAccountHealth()
    const currentBalance = health.freeCollateral // Already a number
    
    // Configuration
    const totalInvested = 546 // From roadmap
    const withdrawalPercent = parseFloat(process.env.WITHDRAWAL_PROFIT_PERCENT || '10')
    const minWithdrawalAmount = parseFloat(process.env.MIN_WITHDRAWAL_AMOUNT || '50')
    const minAccountBalance = parseFloat(process.env.MIN_ACCOUNT_BALANCE || '500')
    
    // Calculate available profit
    const availableProfit = Math.max(0, currentBalance - totalInvested)
    const withdrawalAmount = availableProfit * (withdrawalPercent / 100)
    
    console.log(`💰 Withdrawal calculation:`)
    console.log(`   Current balance: $${currentBalance.toFixed(2)}`)
    console.log(`   Total invested: $${totalInvested.toFixed(2)}`)
    console.log(`   Available profit: $${availableProfit.toFixed(2)}`)
    console.log(`   Withdrawal %: ${withdrawalPercent}%`)
    console.log(`   Withdrawal amount: $${withdrawalAmount.toFixed(2)}`)
    console.log(`   Min withdrawal: $${minWithdrawalAmount.toFixed(2)}`)
    console.log(`   Min account balance: $${minAccountBalance.toFixed(2)}`)
    
    // Safety checks
    if (availableProfit <= 0) {
      return {
        availableProfit,
        withdrawalAmount: 0,
        safeToWithdraw: false,
        reason: 'No profits available (current balance ≤ total invested)',
      }
    }
    
    if (withdrawalAmount < minWithdrawalAmount) {
      return {
        availableProfit,
        withdrawalAmount,
        safeToWithdraw: false,
        reason: `Withdrawal amount ($${withdrawalAmount.toFixed(2)}) below minimum ($${minWithdrawalAmount.toFixed(2)})`,
      }
    }
    
    const balanceAfterWithdrawal = currentBalance - withdrawalAmount
    if (balanceAfterWithdrawal < minAccountBalance) {
      return {
        availableProfit,
        withdrawalAmount,
        safeToWithdraw: false,
        reason: `Withdrawal would drop balance to $${balanceAfterWithdrawal.toFixed(2)} (below minimum $${minAccountBalance.toFixed(2)})`,
      }
    }
    
    return {
      availableProfit,
      withdrawalAmount,
      safeToWithdraw: true,
    }
  } catch (error: any) {
    console.error('❌ Withdrawal calculation failed:', error)
    return {
      availableProfit: 0,
      withdrawalAmount: 0,
      safeToWithdraw: false,
      reason: error.message,
    }
  }
}