trading_bot_v3/v4/lib/drift/client.ts

/**
 * Drift Protocol Client
 * 
 * Handles connection to Drift Protocol and basic operations
 */

import { Connection, PublicKey, Keypair } from '@solana/web3.js'
import { Wallet } from '@coral-xyz/anchor'
import { DriftClient, initialize, User, PerpMarkets } from '@drift-labs/sdk'

export interface DriftConfig {
  rpcUrl: string
  walletPrivateKey: string
  env: 'mainnet-beta' | 'devnet'
}

export class DriftService {
  private connection: Connection
  private wallet: Wallet
  private driftClient: DriftClient | null = null
  private user: User | null = null
  private isInitialized: boolean = false

  constructor(private config: DriftConfig) {
    this.connection = new Connection(config.rpcUrl, 'confirmed')
    
    // Create wallet from private key
    const keypair = Keypair.fromSecretKey(
      Buffer.from(config.walletPrivateKey, 'base58')
    )
    this.wallet = new Wallet(keypair)
    
    console.log('✅ Drift service created for wallet:', this.wallet.publicKey.toString())
  }

  /**
   * Initialize Drift client and subscribe to account updates
   */
  async initialize(): Promise<void> {
    if (this.isInitialized) {
      console.log('⚠️ Drift service already initialized')
      return
    }

    try {
      console.log('🚀 Initializing Drift Protocol client...')

      // Initialize Drift SDK
      await initialize(this.config.env === 'devnet' ? 'devnet' : 'mainnet-beta')

      // Create Drift client
      this.driftClient = new DriftClient({
        connection: this.connection,
        wallet: this.wallet,
        env: this.config.env,
        // Optional: add subaccount ID if using multiple accounts
        // subAccountId: 0,
      })

      // Subscribe to Drift account updates
      await this.driftClient.subscribe()
      console.log('✅ Drift client subscribed to account updates')

      // Get user account
      this.user = this.driftClient.getUser()
      
      this.isInitialized = true
      console.log('✅ Drift service initialized successfully')
      
    } catch (error) {
      console.error('❌ Failed to initialize Drift service:', error)
      throw error
    }
  }

  /**
   * Get current USDC balance
   */
  async getUSDCBalance(): Promise<number> {
    this.ensureInitialized()
    
    try {
      const accountData = this.user!.getUserAccount()
      
      // USDC spot balance (in quote currency)
      const spotBalance = this.user!.getSpotMarketAssetValue(0) // 0 = USDC market
      
      return Number(spotBalance) / 1e6 // USDC has 6 decimals
      
    } catch (error) {
      console.error('❌ Failed to get USDC balance:', error)
      throw error
    }
  }

  /**
   * Get current position for a market
   */
  async getPosition(marketIndex: number): Promise<{
    size: number
    entryPrice: number
    unrealizedPnL: number
    side: 'long' | 'short' | 'none'
  } | null> {
    this.ensureInitialized()
    
    try {
      const position = this.user!.getPerpPosition(marketIndex)
      
      if (!position || position.baseAssetAmount.eq(0)) {
        return null
      }
      
      const baseAssetAmount = Number(position.baseAssetAmount) / 1e9 // 9 decimals
      const quoteAssetAmount = Number(position.quoteAssetAmount) / 1e6 // 6 decimals
      
      // Calculate entry price
      const entryPrice = Math.abs(quoteAssetAmount / baseAssetAmount)
      
      // Get unrealized P&L
      const unrealizedPnL = Number(this.user!.getUnrealizedPNL(false, marketIndex)) / 1e6
      
      const side = baseAssetAmount > 0 ? 'long' : baseAssetAmount < 0 ? 'short' : 'none'
      
      return {
        size: Math.abs(baseAssetAmount),
        entryPrice,
        unrealizedPnL,
        side,
      }
      
    } catch (error) {
      console.error(`❌ Failed to get position for market ${marketIndex}:`, error)
      return null
    }
  }

  /**
   * Get all active positions
   */
  async getAllPositions(): Promise<Array<{
    marketIndex: number
    symbol: string
    size: number
    entryPrice: number
    unrealizedPnL: number
    side: 'long' | 'short'
  }>> {
    this.ensureInitialized()
    
    const positions = []
    
    // Check common markets (SOL, BTC, ETH)
    const markets = [
      { index: 0, symbol: 'SOL-PERP' },
      { index: 1, symbol: 'BTC-PERP' },
      { index: 2, symbol: 'ETH-PERP' },
    ]
    
    for (const market of markets) {
      const position = await this.getPosition(market.index)
      if (position && position.side !== 'none') {
        positions.push({
          marketIndex: market.index,
          symbol: market.symbol,
          ...position,
          side: position.side as 'long' | 'short',
        })
      }
    }
    
    return positions
  }

  /**
   * Get current oracle price for a market
   */
  async getOraclePrice(marketIndex: number): Promise<number> {
    this.ensureInitialized()
    
    try {
      const oracleData = this.driftClient!.getOracleDataForPerpMarket(marketIndex)
      return Number(oracleData.price) / 1e6
      
    } catch (error) {
      console.error(`❌ Failed to get oracle price for market ${marketIndex}:`, error)
      throw error
    }
  }

  /**
   * Get account health (margin ratio)
   */
  async getAccountHealth(): Promise<{
    totalCollateral: number
    totalLiability: number
    freeCollateral: number
    marginRatio: number
  }> {
    this.ensureInitialized()
    
    try {
      const totalCollateral = Number(this.user!.getTotalCollateral()) / 1e6
      const totalLiability = Number(this.user!.getTotalLiability()) / 1e6
      const freeCollateral = Number(this.user!.getFreeCollateral()) / 1e6
      
      const marginRatio = totalLiability > 0 
        ? totalCollateral / totalLiability 
        : Infinity
      
      return {
        totalCollateral,
        totalLiability,
        freeCollateral,
        marginRatio,
      }
      
    } catch (error) {
      console.error('❌ Failed to get account health:', error)
      throw error
    }
  }

  /**
   * Get Drift client instance
   */
  getClient(): DriftClient {
    this.ensureInitialized()
    return this.driftClient!
  }

  /**
   * Get user instance
   */
  getUser(): User {
    this.ensureInitialized()
    return this.user!
  }

  /**
   * Disconnect from Drift
   */
  async disconnect(): Promise<void> {
    if (this.driftClient) {
      await this.driftClient.unsubscribe()
      console.log('✅ Drift client disconnected')
    }
    this.isInitialized = false
  }

  /**
   * Ensure service is initialized
   */
  private ensureInitialized(): void {
    if (!this.isInitialized || !this.driftClient || !this.user) {
      throw new Error('Drift service not initialized. Call initialize() first.')
    }
  }
}

// Singleton instance
let driftServiceInstance: DriftService | null = null

export function getDriftService(): DriftService {
  if (!driftServiceInstance) {
    const config: DriftConfig = {
      rpcUrl: process.env.SOLANA_RPC_URL || 'https://api.mainnet-beta.solana.com',
      walletPrivateKey: process.env.DRIFT_WALLET_PRIVATE_KEY || '',
      env: (process.env.DRIFT_ENV as 'mainnet-beta' | 'devnet') || 'mainnet-beta',
    }
    
    if (!config.walletPrivateKey) {
      throw new Error('DRIFT_WALLET_PRIVATE_KEY not set in environment')
    }
    
    driftServiceInstance = new DriftService(config)
  }
  
  return driftServiceInstance
}

export async function initializeDriftService(): Promise<DriftService> {
  const service = getDriftService()
  await service.initialize()
  return service
}