trading_bot_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 { DriftClient, initialize, User, PerpMarkets } from '@drift-labs/sdk'
import bs58 from 'bs58'

// Manual wallet interface (more compatible than SDK Wallet class)
interface ManualWallet {
  publicKey: PublicKey
  signTransaction: (tx: any) => Promise<any>
  signAllTransactions: (txs: any[]) => Promise<any[]>
}

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

export class DriftService {
  private connection: Connection
  private wallet: ManualWallet
  private keypair: Keypair
  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
    // Support both formats:
    // 1. JSON array: [91,24,199,...] (from Phantom export as array)
    // 2. Base58 string: "5Jm7X..." (from Phantom export as string)
    let secretKey: Uint8Array
    
    if (config.walletPrivateKey.startsWith('[')) {
      // JSON array format
      const keyArray = JSON.parse(config.walletPrivateKey)
      secretKey = new Uint8Array(keyArray)
    } else {
      // Base58 string format
      secretKey = bs58.decode(config.walletPrivateKey)
    }
    
    this.keypair = Keypair.fromSecretKey(secretKey)
    
    // Create manual wallet interface (more reliable than SDK Wallet)
    this.wallet = {
      publicKey: this.keypair.publicKey,
      signTransaction: async (tx) => {
        if (typeof tx.partialSign === 'function') {
          tx.partialSign(this.keypair)
        } else if (typeof tx.sign === 'function') {
          tx.sign([this.keypair])
        }
        return tx
      },
      signAllTransactions: async (txs) => {
        return txs.map(tx => {
          if (typeof tx.partialSign === 'function') {
            tx.partialSign(this.keypair)
          } else if (typeof tx.sign === 'function') {
            tx.sign([this.keypair])
          }
          return tx
        })
      }
    }
    
    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 (gets program IDs and config)
      const sdkConfig = initialize({
        env: this.config.env === 'devnet' ? 'devnet' : 'mainnet-beta'
      })

      // Create Drift client with manual wallet and SDK config
      this.driftClient = new DriftClient({
        connection: this.connection,
        wallet: this.wallet as any, // Type assertion for compatibility
        programID: new PublicKey(sdkConfig.DRIFT_PROGRAM_ID),
        opts: {
          commitment: 'confirmed',
        },
      })

      // 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.eqn(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!.getTotalLiabilityValue()) / 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 with better persistence
let driftServiceInstance: DriftService | null = null
let initializationPromise: Promise<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)
    console.log('🔄 Created new Drift service singleton')
  } else {
    console.log('♻️  Reusing existing Drift service instance')
  }
  
  return driftServiceInstance
}

export async function initializeDriftService(): Promise<DriftService> {
  // If already initializing, return the same promise to avoid multiple concurrent inits
  if (initializationPromise) {
    console.log('⏳ Waiting for ongoing initialization...')
    return initializationPromise
  }
  
  const service = getDriftService()
  
  // If already initialized, return immediately
  if (service['isInitialized']) {
    console.log('✅ Drift service already initialized')
    return service
  }
  
  // Start initialization and cache the promise
  initializationPromise = service.initialize().then(() => {
    initializationPromise = null // Clear after completion
    return service
  }).catch((error) => {
    initializationPromise = null // Clear on error so it can be retried
    throw error
  })
  
  return initializationPromise
}