trading_bot_v3/app/api/drift/balance/route.js

import { NextResponse } from 'next/server'
import { executeWithFailover, getRpcStatus } from '../../../../lib/rpc-failover.js'

export async function GET() {
  try {
    console.log('💰 Getting Drift account balance...')
    
    // Log RPC status
    const rpcStatus = getRpcStatus()
    console.log('🌐 RPC Status:', rpcStatus)

    // Check if environment is configured
    if (!process.env.SOLANA_PRIVATE_KEY) {
      return NextResponse.json({
        success: false,
        error: 'Drift not configured - missing SOLANA_PRIVATE_KEY'
      }, { status: 400 })
    }

    // Execute balance check with RPC failover
    const result = await executeWithFailover(async (connection) => {
      // Import Drift SDK components
      const { DriftClient, initialize, calculateFreeCollateral, calculatePositionPNL, QUOTE_PRECISION } = await import('@drift-labs/sdk')
      const { Keypair } = await import('@solana/web3.js')
      const { AnchorProvider, BN } = await import('@coral-xyz/anchor')

      const privateKeyArray = JSON.parse(process.env.SOLANA_PRIVATE_KEY)
      const keypair = Keypair.fromSecretKey(new Uint8Array(privateKeyArray))
      
      // Use the correct Wallet class from @coral-xyz/anchor/dist/cjs/nodewallet
      const { default: NodeWallet } = await import('@coral-xyz/anchor/dist/cjs/nodewallet.js')
      const wallet = new NodeWallet(keypair)

      // Initialize Drift SDK
      const env = 'mainnet-beta'
      const sdkConfig = initialize({ env })
      
      const driftClient = new DriftClient({
        connection,
        wallet,
        programID: sdkConfig.DRIFT_PROGRAM_ID,
        opts: {
          commitment: 'confirmed',
        },
      })

      try {
        await driftClient.subscribe()
        console.log('✅ Connected to Drift for balance check')

        // Check if user has account
        let userAccount
        try {
          userAccount = await driftClient.getUserAccount()
        } catch (accountError) {
          await driftClient.unsubscribe()
          throw new Error('No Drift user account found. Please initialize your account first.')
        }

        // Get account balances and positions using Drift's built-in methods
        const spotBalances = userAccount.spotPositions || []
        const perpPositions = userAccount.perpPositions || []

        // Use Drift's built-in calculation methods for accuracy
        let totalCollateral = 0
        let unrealizedPnl = 0
        let marginRequirement = 0

        try {
          // Calculate total collateral using Drift's method
          totalCollateral = await driftClient.getUser().getTotalCollateral() / 1e6 // Convert to USDC
        } catch (collateralError) {
          console.warn('⚠️ Could not get total collateral, calculating manually:', collateralError.message)
          
          // Fallback to manual USDC balance calculation
          const usdcBalance = spotBalances.find(pos => pos.marketIndex === 0)
          if (usdcBalance) {
            totalCollateral = Number(usdcBalance.scaledBalance) / 1e6 // Assume 6 decimal precision for USDC
          }
        }

        try {
          // Calculate unrealized PnL using Drift's method
          unrealizedPnl = await driftClient.getUser().getUnrealizedPNL() / 1e6 // Convert to USDC
        } catch (pnlError) {
          console.warn('⚠️ Could not get unrealized PnL, calculating manually:', pnlError.message)
          unrealizedPnl = 0 // Default to 0 if we can't calculate
        }

        try {
          // Calculate margin requirement using Drift's method
          marginRequirement = await driftClient.getUser().getTotalPerpPositionValue() / 1e6 * 0.1 // 10% margin
        } catch (marginError) {
          console.warn('⚠️ Could not get margin requirement, calculating manually:', marginError.message)
          marginRequirement = 0 // Default to 0 if we can't calculate
        }

        // Calculate free collateral and other derived values
        const freeCollateral = totalCollateral - marginRequirement + unrealizedPnl
        const accountValue = totalCollateral + unrealizedPnl
        const leverage = marginRequirement > 0 ? (marginRequirement / accountValue) : 0
        const availableBalance = Math.max(0, freeCollateral)

        // Count active positions
        const activePositions = perpPositions.filter(pos => 
          pos.baseAssetAmount && !pos.baseAssetAmount.isZero()
        )

        const balanceResult = {
          success: true,
          totalCollateral: totalCollateral,
          freeCollateral: freeCollateral,
          marginRequirement: marginRequirement,
          unrealizedPnl: unrealizedPnl,
          accountValue: accountValue,
          leverage: leverage,
          availableBalance: availableBalance,
          activePositionsCount: activePositions.length,
          timestamp: Date.now(),
          rpcEndpoint: getRpcStatus().currentEndpoint,
          details: {
            spotBalances: spotBalances.length,
            perpPositions: activePositions.length,
            wallet: keypair.publicKey.toString()
          }
        }

        await driftClient.unsubscribe()
        
        console.log('💰 Balance retrieved:', {
          totalCollateral: totalCollateral.toFixed(2),
          availableBalance: availableBalance.toFixed(2),
          positions: activePositions.length,
          rpcEndpoint: getRpcStatus().currentEndpoint
        })

        return balanceResult

      } catch (driftError) {
        console.error('❌ Drift balance error:', driftError)
        
        try {
          await driftClient.unsubscribe()
        } catch (cleanupError) {
          console.warn('⚠️ Cleanup error:', cleanupError.message)
        }
        
        throw driftError
      }
    }, 3) // Max 3 retries across different RPCs

    return NextResponse.json(result)

  } catch (error) {
    console.error('❌ Balance API error:', error)
    
    return NextResponse.json({
      success: false,
      error: 'Failed to get Drift account balance',
      details: error.message,
      rpcStatus: getRpcStatus()
    }, { status: 500 })
  }
}

export async function POST() {
  return NextResponse.json({
    message: 'Use GET method to retrieve Drift account balance'
  }, { status: 405 })
}