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

import { NextResponse } from 'next/server'

export async function GET() {
  try {
    console.log('💰 Getting Drift account balance...')

    // 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 })
    }

    // Import Drift SDK components
    const { DriftClient, initialize, calculateFreeCollateral, QUOTE_PRECISION } = await import('@drift-labs/sdk')
    const { Connection, Keypair } = await import('@solana/web3.js')
    const { AnchorProvider, Wallet, BN } = await import('@coral-xyz/anchor')

    // Initialize connection and wallet
    const connection = new Connection(
      process.env.SOLANA_RPC_URL || 'https://api.mainnet-beta.solana.com',
      'confirmed'
    )

    const privateKeyArray = JSON.parse(process.env.SOLANA_PRIVATE_KEY)
    const keypair = Keypair.fromSecretKey(new Uint8Array(privateKeyArray))
    const wallet = new Wallet(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()
        return NextResponse.json({
          success: false,
          error: 'No Drift user account found. Please initialize your account first.',
          needsInitialization: true
        }, { status: 404 })
      }

      // Get account balances and positions
      const spotBalances = userAccount.spotPositions || []
      const perpPositions = userAccount.perpPositions || []

      // Calculate key metrics
      let totalCollateral = 0
      let unrealizedPnl = 0
      let marginRequirement = 0

      // Process spot balances (USDC collateral)
      const usdcBalance = spotBalances.find(pos => pos.marketIndex === 0) // USDC is typically index 0
      if (usdcBalance) {
        totalCollateral = Number(usdcBalance.scaledBalance) / Math.pow(10, 6) // USDC has 6 decimals
      }

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

      for (const position of activePositions) {
        const baseAmount = Number(position.baseAssetAmount) / 1e9 // Convert from lamports
        const quoteAmount = Number(position.quoteAssetAmount) / 1e6 // Convert from micro-USDC
        
        unrealizedPnl += quoteAmount
        marginRequirement += Math.abs(baseAmount * 100) // Simplified margin calculation
      }

      // Calculate free collateral (simplified)
      const freeCollateral = totalCollateral - marginRequirement + unrealizedPnl
      
      // Calculate account value and leverage
      const accountValue = totalCollateral + unrealizedPnl
      const leverage = marginRequirement > 0 ? (marginRequirement / accountValue) : 0
      
      // Available balance for new positions
      const availableBalance = Math.max(0, freeCollateral)

      const result = {
        success: true,
        totalCollateral: totalCollateral,
        freeCollateral: freeCollateral,
        marginRequirement: marginRequirement,
        unrealizedPnl: unrealizedPnl,
        accountValue: accountValue,
        leverage: leverage,
        availableBalance: availableBalance,
        activePositionsCount: activePositions.length,
        timestamp: Date.now(),
        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
      })

      return NextResponse.json(result)

    } catch (driftError) {
      console.error('❌ Drift balance error:', driftError)
      
      try {
        await driftClient.unsubscribe()
      } catch (cleanupError) {
        console.warn('⚠️ Cleanup error:', cleanupError.message)
      }
      
      return NextResponse.json({
        success: false,
        error: 'Failed to get Drift account balance',
        details: driftError.message
      }, { status: 500 })
    }

  } catch (error) {
    console.error('❌ Balance API error:', error)
    
    return NextResponse.json({
      success: false,
      error: 'Internal server error getting balance',
      details: error.message
    }, { status: 500 })
  }
}

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