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

import { NextResponse } from 'next/server'

// Helper function to get market index from symbol
function getMarketIndex(symbol) {
  const marketMap = {
    'SOL': 0,
    'BTC': 1,
    'ETH': 2,
    'APT': 3,
    'AVAX': 4,
    'BNB': 5,
    'MATIC': 6,
    'ARB': 7,
    'DOGE': 8,
    'OP': 9
  }
  
  const index = marketMap[symbol.toUpperCase()]
  if (index === undefined) {
    throw new Error(`Unsupported symbol: ${symbol}`)
  }
  
  return index
}

// Helper function to get symbol from market index
function getSymbolFromMarketIndex(marketIndex) {
  const symbols = ['SOL', 'BTC', 'ETH', 'APT', 'AVAX', 'BNB', 'MATIC', 'ARB', 'DOGE', 'OP']
  return symbols[marketIndex] || `UNKNOWN_${marketIndex}`
}

// Helper function to get trading balance with better error handling
async function getTradingBalance(driftClient) {
  try {
    const userAccount = await driftClient.getUserAccount()
    
    if (!userAccount) {
      throw new Error('User account is null')
    }
    
    console.log('📊 Raw user account data keys:', Object.keys(userAccount))
    
    // Get all spot positions
    const spotPositions = userAccount.spotPositions || []
    const usdcPosition = spotPositions.find(pos => pos.marketIndex === 0) // USDC is usually index 0
    
    // Convert BigNumber values to regular numbers
    const BN = (await import('bn.js')).default
    
    // Get collateral info - convert from BN to number
    const totalCollateral = userAccount.totalCollateral ? 
      (userAccount.totalCollateral instanceof BN ? userAccount.totalCollateral.toNumber() / 1e6 : 
       parseFloat(userAccount.totalCollateral.toString()) / 1e6) : 0
    
    const freeCollateral = userAccount.freeCollateral ? 
      (userAccount.freeCollateral instanceof BN ? userAccount.freeCollateral.toNumber() / 1e6 : 
       parseFloat(userAccount.freeCollateral.toString()) / 1e6) : 0
    
    // Get USDC balance
    const usdcBalance = usdcPosition && usdcPosition.scaledBalance ? 
      (usdcPosition.scaledBalance instanceof BN ? usdcPosition.scaledBalance.toNumber() / 1e6 : 
       parseFloat(usdcPosition.scaledBalance.toString()) / 1e6) : 0
    
    console.log('💰 Parsed balances:', {
      totalCollateral,
      freeCollateral,
      usdcBalance,
      spotPositionsCount: spotPositions.length
    })
    
    return {
      totalCollateral: totalCollateral.toString(),
      freeCollateral: freeCollateral.toString(),
      usdcBalance: usdcBalance.toString(),
      marginRatio: userAccount.marginRatio ? userAccount.marginRatio.toString() : '0',
      accountExists: true,
      spotPositions: spotPositions.map(pos => ({
        marketIndex: pos.marketIndex,
        balance: pos.scaledBalance ? 
          (pos.scaledBalance instanceof BN ? pos.scaledBalance.toNumber() / 1e6 : 
           parseFloat(pos.scaledBalance.toString()) / 1e6) : 0
      }))
    }
  } catch (error) {
    throw new Error(`Balance retrieval failed: ${error.message}`)
  }
}

export async function POST(request) {
  try {
    console.log('🌊 Drift leverage trading endpoint...')

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

    const { action = 'get_balance', symbol = 'SOL', amount, side, leverage = 1 } = await request.json()

    // Import Drift SDK components
    const { DriftClient, initialize } = await import('@drift-labs/sdk')
    const { Connection, Keypair } = await import('@solana/web3.js')

    // Initialize connection with Helius
    const heliusApiKey = '5e236449-f936-4af7-ae38-f15e2f1a3757'
    const rpcUrl = `https://mainnet.helius-rpc.com/?api-key=${heliusApiKey}`
    const connection = new Connection(rpcUrl, 'confirmed')
    
    console.log('🌐 Using mainnet with Helius RPC')

    const privateKeyArray = JSON.parse(process.env.SOLANA_PRIVATE_KEY)
    const keypair = Keypair.fromSecretKey(new Uint8Array(privateKeyArray))
    
    // Create wallet using manual interface (most reliable)
    const wallet = {
      publicKey: keypair.publicKey,
      signTransaction: async (tx) => {
        if (typeof tx.partialSign === 'function') {
          tx.partialSign(keypair)
        } else if (typeof tx.sign === 'function') {
          tx.sign([keypair])
        }
        return tx
      },
      signAllTransactions: async (txs) => {
        return txs.map(tx => {
          if (typeof tx.partialSign === 'function') {
            tx.partialSign(keypair)
          } else if (typeof tx.sign === 'function') {
            tx.sign([keypair])
          }
          return tx
        })
      }
    }

    console.log('🔐 Connecting to Drift with wallet:', keypair.publicKey.toString())

    // 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 {
      // Subscribe to drift client
      await driftClient.subscribe()
      console.log('✅ Connected to Drift successfully')

      // Handle action
      let result = {}

      if (action === 'get_balance') {
        try {
          // Simple and direct approach
          console.log('🔍 Getting user account...')
          const userAccount = await driftClient.getUserAccount()
          
          if (userAccount) {
            console.log('✅ User account found, getting balance...')
            result = await getTradingBalance(driftClient)
            console.log('✅ Balance retrieved successfully')
          } else {
            console.log('❌ User account is null')
            result = {
              message: 'User account exists but returns null',
              accountExists: false
            }
          }
        } catch (error) {
          console.log('❌ Error getting user account:', error.message)
          
          // Check wallet SOL balance as fallback
          const walletBalance = await connection.getBalance(keypair.publicKey)
          const solBalance = walletBalance / 1e9
          
          result = {
            message: 'Cannot access user account data',
            error: error.message,
            solBalance: solBalance,
            walletAddress: keypair.publicKey.toString(),
            suggestion: 'Account may need to be accessed through Drift UI first or deposit USDC directly'
          }
        }
      } else if (action === 'place_order') {
        // Place a leverage order
        if (!amount || !side) {
          result = {
            error: 'Missing required parameters: amount and side'
          }
        } else {
          try {
            const { OrderType, PositionDirection } = await import('@drift-labs/sdk')
            const BN = (await import('bn.js')).default
            
            const marketIndex = getMarketIndex(symbol)
            
            // Convert amount to base units (SOL uses 9 decimals)
            const baseAssetAmount = new BN(Math.floor(amount * 1e9))
            
            console.log(`💰 Amount conversion:`, {
              inputAmount: amount,
              calculatedAmount: amount * 1e9,
              flooredAmount: Math.floor(amount * 1e9),
              baseAssetAmount: baseAssetAmount.toString()
            })
            
            // Determine direction
            const direction = side.toLowerCase() === 'buy' ? PositionDirection.LONG : PositionDirection.SHORT
            
            console.log(`📊 Placing ${side} order:`, {
              symbol,
              marketIndex,
              amount,
              leverage,
              baseAssetAmount: baseAssetAmount.toString()
            })

            // Place perpetual order
            const txSig = await driftClient.placePerpOrder({
              orderType: OrderType.MARKET,
              marketIndex,
              direction,
              baseAssetAmount,
              reduceOnly: false,
            })

            console.log('✅ Order placed:', txSig)
            
            // Wait for confirmation
            await new Promise(resolve => setTimeout(resolve, 3000))

            // Get position after order
            const userAccount = await driftClient.getUserAccount()
            const position = userAccount.perpPositions.find(pos => pos.marketIndex === marketIndex)

            result = {
              transactionId: txSig,
              symbol,
              side,
              amount,
              leverage,
              position: position ? {
                marketIndex: position.marketIndex,
                baseAssetAmount: position.baseAssetAmount.toString(),
                quoteAssetAmount: position.quoteAssetAmount.toString()
              } : null
            }
          } catch (orderError) {
            console.log('❌ Failed to place order:', orderError.message)
            result = {
              error: 'Failed to place order',
              details: orderError.message
            }
          }
        }
      } else {
        result = { message: `Action ${action} not yet implemented` }
      }

      // Clean up connection
      await driftClient.unsubscribe()

      return NextResponse.json({
        success: true,
        action,
        result,
        timestamp: Date.now()
      })

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

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