bundle.ts

import {
  PublicKey,
  Transaction,
  TransactionInstruction,
  Keypair,
  SystemProgram,
  VersionedTransaction,
  TransactionMessage,
  AddressLookupTableAccount,
} from '@solana/web3.js'
import { connection } from '../clients/lp'
import { privy } from '../clients/privy'
import { log } from './log'
import { Bundle } from 'jito-ts/dist/sdk/block-engine/types'
import { searcherClient } from 'jito-ts/dist/sdk/block-engine/searcher'
import { status } from '@grpc/grpc-js'
import { getTipAccount } from './tipAccounts'
import { JITO_BUNDLE_TIP_LAMPORTS_ZAP_IN, HELIUS_RPC_URL, HELIUS_API_KEY, LOG_SIMULATION, JITO_ENDPOINTS, SOLANA_NETWORK } from './constants'

// Default bundle API endpoint (will be overridden when retrying)
const getJitoBundleApiUrl = (endpoint?: string) => {
  const baseEndpoint = endpoint || JITO_ENDPOINTS[0]
  return `https://${baseEndpoint}/api/v1/bundles`
}

export interface BundleResult {
  bundleId: string
  confirmed: boolean
  txSignatures: string[]
  error?: string
}

/**
 * Check if an error is a rate limit error
 * Checks both gRPC status codes and error messages
 */
function isRateLimitError(result: any): boolean {
  if (!result || result.ok) {
    return false
  }
  
  // Check gRPC status code (RESOURCE_EXHAUSTED for rate limiting)
  if (result.err?.code !== undefined) {
    // gRPC status.RESOURCE_EXHAUSTED = 8 (rate limiting)
    if (result.err.code === status.RESOURCE_EXHAUSTED) {
      return true
    }
  }
  
  // Also check HTTP status code if present (429 = Too Many Requests)
  if (result.err?.status === 429 || result.status === 429) {
    return true
  }
  
  // Fallback: Check for rate limit in error message or details
  const errorStr = JSON.stringify(result).toLowerCase()
  return (
    errorStr.includes('rate limit') ||
    errorStr.includes('ratelimit') ||
    errorStr.includes('429') ||
    errorStr.includes('too many requests') ||
    errorStr.includes('resource exhausted')
  )
}

/**
 * Send a bundle with auto-retry on rate limit errors
 * Automatically rotates through Jito endpoints if rate limited
 */
export async function sendBundleWithRetry(
  bundle: Bundle,
  logPrefix: string = '[bundle]'
): Promise<string> {
  let lastError: any = null
  let lastEndpoint: string | undefined = undefined
  
  for (let i = 0; i < JITO_ENDPOINTS.length; i++) {
    const endpoint = JITO_ENDPOINTS[i]
    const client = searcherClient(endpoint)
    
    try {
      log.info(`${logPrefix} Sending bundle to endpoint: ${endpoint} (attempt ${i + 1}/${JITO_ENDPOINTS.length})`)
      const bundleResult = await client.sendBundle(bundle)

      if (bundleResult.ok && bundleResult.value) {
        const bundleId = bundleResult.value
        log.info(`${logPrefix} Bundle ID: ${bundleId} (sent via ${endpoint})`)
        return bundleId
      } else {
        // Check if it's a rate limit error
        if (isRateLimitError(bundleResult)) {
          log.warn(`${logPrefix} Rate limited by ${endpoint}, trying next endpoint...`)
          lastError = bundleResult
          lastEndpoint = endpoint
          // Continue to next endpoint immediately
          continue
        } else {
          // Non-rate-limit error, log and throw
          log.error(`${logPrefix} Failed to send bundle to ${endpoint}:`, bundleResult)
          throw new Error(`Failed to send bundle: ${JSON.stringify(bundleResult)}`)
        }
      }
    } catch (error: any) {
      // Check if error is rate limit related
      const errorStr = error?.message?.toLowerCase() || JSON.stringify(error).toLowerCase()
      if (
        errorStr.includes('rate limit') ||
        errorStr.includes('ratelimit') ||
        errorStr.includes('429') ||
        errorStr.includes('too many requests')
      ) {
        log.warn(`${logPrefix} Rate limit error from ${endpoint}, trying next endpoint...`)
        lastError = error
        lastEndpoint = endpoint
        // Continue to next endpoint
        continue
      } else {
        // Non-rate-limit error, rethrow
        log.error(`${logPrefix} Error sending bundle to ${endpoint}:`, error)
        throw error
      }
    }
  }

  // All endpoints exhausted
  log.error(`${logPrefix} Failed to send bundle to all endpoints. Last error:`, lastError)
  throw new Error(`Failed to send bundle: rate limited on all endpoints. Last endpoint: ${lastEndpoint}`)
}

/**
 * Check bundle status via Jito API
 * Returns transaction signatures if bundle landed
 */
export async function getBundleStatus(
  bundleId: string, 
  logPrefix: string = '[bundle]',
  endpoint?: string
): Promise<{
  status: 'pending' | 'confirmed' | 'finalized' | 'failed' | 'not_found'
  transactions?: string[]
}> {
  const apiUrl = getJitoBundleApiUrl(endpoint)
  
  try {
    const response = await fetch(apiUrl, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        jsonrpc: '2.0',
        id: 1,
        method: 'getBundleStatuses',
        params: [[bundleId]],
      }),
    })

    // Handle HTTP rate limit
    if (response.status === 429) {
      log.warn(`${logPrefix} Rate limited (429) when checking bundle status`)
      // Try next endpoint if available
      if (endpoint) {
        const currentIndex = JITO_ENDPOINTS.indexOf(endpoint)
        if (currentIndex >= 0 && currentIndex < JITO_ENDPOINTS.length - 1) {
          const nextEndpoint = JITO_ENDPOINTS[currentIndex + 1]
          log.info(`${logPrefix} Retrying bundle status check with endpoint: ${nextEndpoint}`)
          return getBundleStatus(bundleId, logPrefix, nextEndpoint)
        }
      }
      return { status: 'not_found' }
    }

    if (!response.ok) {
      log.error(`${logPrefix} Bundle status check failed: ${response.status} ${response.statusText}`)
      return { status: 'not_found' }
    }

    const data = await response.json()
    const result = data?.result?.value?.[0]

    if (!result) {
      // Jito doesn't have this bundle - could be expired or never received
      return { status: 'not_found' }
    }

    const confirmationStatus = result.confirmation_status
    log.info(`${logPrefix} Bundle ${bundleId.slice(0, 8)}... status: ${confirmationStatus || 'unknown'}, err: ${JSON.stringify(result.err)}`)

    if (confirmationStatus === 'finalized' || confirmationStatus === 'confirmed') {
      return {
        status: confirmationStatus,
        transactions: result.transactions || [],
      }
    }

    // Landed status means it was included but may not be confirmed yet
    if (confirmationStatus === 'landed') {
      return {
        status: 'confirmed',
        transactions: result.transactions || [],
      }
    }

    // If err.Ok is null, the bundle is still pending
    if (result.err && result.err.Ok === null) {
      return { status: 'pending' }
    }

    // Any other error means failed
    if (result.err) {
      log.error(`${logPrefix} Bundle error: ${JSON.stringify(result.err)}`)
      return { status: 'failed' }
    }

    // No confirmation status yet, still pending
    return { status: 'pending' }
  } catch (error) {
    log.error('[bundle] Failed to get bundle status:', error)
    return { status: 'not_found' }
  }
}

/**
 * Wait for bundle confirmation with polling
 * Returns transaction signatures when confirmed
 */
export async function waitForBundleConfirmation(
  bundleId: string,
  maxAttempts: number = 10,
  delayMs: number = 2000,
  logPrefix: string = '[bundle]'
): Promise<{ confirmed: boolean; txSignatures: string[] }> {
  let lastStatus = ''
  let notFoundCount = 0

  for (let attempt = 1; attempt <= maxAttempts; attempt++) {
    const status = await getBundleStatus(bundleId, logPrefix)
    lastStatus = status.status

    if (status.status === 'confirmed' || status.status === 'finalized') {
      log.info(`${logPrefix} Bundle confirmed after ${attempt} attempts`)
      return { confirmed: true, txSignatures: status.transactions || [] }
    }

    if (status.status === 'failed') {
      log.error(`${logPrefix} Bundle failed`)
      return { confirmed: false, txSignatures: [] }
    }

    // Track consecutive not_found responses
    // Jito may return not_found briefly while processing, but prolonged not_found is concerning
    if (status.status === 'not_found') {
      notFoundCount++
      // If we get not_found consistently for many attempts, the bundle likely wasn't received
      if (notFoundCount >= 10) {
        log.warn(`${logPrefix} Bundle not found in Jito after ${notFoundCount} checks`)
      }
    } else {
      notFoundCount = 0
    }

    if (attempt < maxAttempts) {
      await new Promise((resolve) => setTimeout(resolve, delayMs))
    }
  }

  log.warn(`${logPrefix} Bundle confirmation timeout after ${maxAttempts} attempts, last status: ${lastStatus}`)
  return { confirmed: false, txSignatures: [] }
}

/**
 * Creates a tip instruction for bundling
 * Returns a random tip account and the transfer instruction
 */
export function createTipInstruction(
  walletPubkey: PublicKey,
  tipAmountLamports?: number,
): { tipAccount: PublicKey; instruction: TransactionInstruction } {
  const tipAccount = getTipAccount() // Random tip account
  const tipAmount = tipAmountLamports ?? JITO_BUNDLE_TIP_LAMPORTS_ZAP_IN
  const tipIx = SystemProgram.transfer({
    fromPubkey: walletPubkey,
    toPubkey: tipAccount,
    lamports: tipAmount,
  })

  return { tipAccount, instruction: tipIx }
}

/**
 * Simulates a bundle using Helius simulateBundle RPC
 * Returns simulation results or null if simulation fails/disabled
 */
export async function simulateBundle(
  encodedTransactions: string[],
  logPrefix: string = '[bundle_sim]',
): Promise<any | null> {
  if (!HELIUS_API_KEY) {
    if (LOG_SIMULATION) {
      log.warn(`${logPrefix} HELIUS_API_KEY not set, skipping bundle simulation`)
    }
    return null
  }

  try {
    const url = `${HELIUS_RPC_URL}?api-key=${HELIUS_API_KEY}`
    
    const requestBody = {
      jsonrpc: '2.0',
      id: '1',
      method: 'simulateBundle',
      params: [
        {
          encodedTransactions,
          transactionEncoding: 'base64',
          replaceRecentBlockhash: true,
          skipSigVerify: true,
        },
      ],
    }

    if (LOG_SIMULATION) {
      log.info(`${logPrefix} Simulating bundle with ${encodedTransactions.length} transactions...`)
    }
    
    const response = await fetch(url, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
      },
      body: JSON.stringify(requestBody),
    })

    if (!response.ok) {
      if (LOG_SIMULATION) {
        log.warn(`${logPrefix} Bundle simulation HTTP error: ${response.status} ${response.statusText}`)
      }
      return null
    }

    const json = await response.json()

    if (LOG_SIMULATION) {
      log.info(`${logPrefix} Complete simulation response:`, JSON.stringify(json, null, 2))
    }

    if (json.error) {
      if (LOG_SIMULATION) {
        log.warn(`${logPrefix} Bundle simulation error:`, json.error)
      }
      return null
    }

    const result = json.result
    if (!result || !result.value) {
      if (LOG_SIMULATION) {
        log.warn(`${logPrefix} Bundle simulation returned no result or value`)
        log.warn(`${logPrefix} Full JSON response:`, JSON.stringify(json, null, 2))
      }
      return null
    }

    const value = result.value
    const summary = value.summary

    // Parse summary - can be "succeeded", or object with "failed" property
    let summaryStatus = 'unknown'
    let failedTx: any = null
    
    if (typeof summary === 'string') {
      summaryStatus = summary
    } else if (summary && typeof summary === 'object') {
      if (summary.failed) {
        summaryStatus = 'failed'
        failedTx = summary.failed
        if (LOG_SIMULATION) {
          log.error(`${logPrefix} Bundle simulation FAILED`)
          log.error(`${logPrefix} Failed transaction signature: ${failedTx.tx_signature || 'N/A'}`)
          if (failedTx.error) {
            log.error(`${logPrefix} Failure error:`, JSON.stringify(failedTx.error, null, 2))
          }
        }
      } else if (summary.succeeded) {
        summaryStatus = 'succeeded'
      }
    }

    if (LOG_SIMULATION) {
      log.info(`${logPrefix} Bundle simulation summary: ${summaryStatus}`)
    }
    
    // Log transaction results
    if (LOG_SIMULATION && value.transactionResults && Array.isArray(value.transactionResults)) {
      log.info(`${logPrefix} Transaction results count: ${value.transactionResults.length}`)
      
      value.transactionResults.forEach((txResult: any, index: number) => {
        const txIndex = index + 1
        if (txResult.err) {
          log.error(`${logPrefix} Transaction ${txIndex} simulation error:`, JSON.stringify(txResult.err, null, 2))
        } else {
          log.info(`${logPrefix} Transaction ${txIndex} simulation: success, units: ${txResult.unitsConsumed || 'N/A'}`)
        }
        
        // Check if this transaction matches the failed signature
        if (failedTx && failedTx.tx_signature) {
          // Note: We can't directly match signatures from transactionResults, but we can check by index
          // The failed transaction might be one that's not in transactionResults (like the tip transaction)
        }
      })
      
      // If summary says failed but we have results, the failure might be in a transaction not in results
      if (summaryStatus === 'failed' && failedTx) {
        log.error(`${logPrefix} Failed transaction details:`)
        log.error(`${logPrefix}   Signature: ${failedTx.tx_signature}`)
        log.error(`${logPrefix}   Error: ${JSON.stringify(failedTx.error, null, 2)}`)
        log.error(`${logPrefix}   Note: This transaction may not be in transactionResults (could be tip transaction or transaction #${value.transactionResults.length + 1})`)
      }
    } else if (LOG_SIMULATION && !value.transactionResults) {
      log.warn(`${logPrefix} No transactionResults in simulation result`)
    }

    return value
  } catch (error) {
    if (LOG_SIMULATION) {
      log.warn(`${logPrefix} Bundle simulation failed:`, error)
    }
    return null
  }
}

/**
 * Bundles and sends transactions via Jito
 * Returns bundle ID if successful
 */
export async function bundleAndSendTransactions(
  transactions: Array<Transaction | VersionedTransaction>,
  walletPubkey: PublicKey,
  embeddedWalletId: string,
  poolAuthContext: any,
  positionKeypair?: Keypair | null,
  programId?: PublicKey,
  logPrefix: string = '[bundle]',
  transactionNames?: string[],
  tipAmountLamports?: number,
  transactionIndexToKeypairMap?: Map<number, Keypair>,
): Promise<string | null> {
  if (transactions.length === 0) {
    return null
  }

  // Get shared blockhash for all transactions
  const blockhashResp =
    await connection.getLatestBlockhashAndContext('finalized')
  const sharedBlockhash = blockhashResp.value.blockhash

  // Process transactions - handle VersionedTransaction and Transaction differently
  const processedTxs: Array<Transaction | VersionedTransaction> = []
  const positionKeypairSigningMap = new Map<number, Keypair>()
  
  for (let i = 0; i < transactions.length; i++) {
    const tx = transactions[i]
    const txName = transactionNames?.[i] || `Transaction ${i + 1}`
    
    if (tx instanceof VersionedTransaction) {
      // Handle VersionedTransaction (Jupiter swaps)
      // Update blockhash by decompiling, updating, and recompiling
      const message = tx.message
      const altAccountResponses = await Promise.all(
        message.addressTableLookups.map((l) =>
          connection.getAddressLookupTable(l.accountKey),
        ),
      )

      const altAccounts = altAccountResponses.map((item) => {
        if (item.value == null) throw new Error('ALT is null')
        return item.value
      })

      const decompiledMessage = TransactionMessage.decompile(message, {
        addressLookupTableAccounts: altAccounts,
      })

      decompiledMessage.recentBlockhash = sharedBlockhash

      // If this is the last transaction, add tip instruction
      if (i === transactions.length - 1) {
        const { instruction: tipIx } = createTipInstruction(walletPubkey, tipAmountLamports)
        decompiledMessage.instructions.push(tipIx)
        log.info(`${logPrefix} Added tip instruction to last transaction (VersionedTransaction)`)
      }

      const updatedVersionedTx = new VersionedTransaction(
        decompiledMessage.compileToV0Message(altAccounts),
      )
      
      // Don't copy signatures from original - they're invalid after changing blockhash
      processedTxs.push(updatedVersionedTx)
      log.info(`${logPrefix} Updated blockhash for ${txName} (VersionedTransaction)`)
    } else {
      // Handle Transaction (add liquidity, fees)
      tx.recentBlockhash = sharedBlockhash
      tx.feePayer = walletPubkey

      // If this is the last transaction, add tip instruction
      if (i === transactions.length - 1) {
        const { instruction: tipIx } = createTipInstruction(walletPubkey, tipAmountLamports)
        tx.add(tipIx)
        log.info(`${logPrefix} Added tip instruction to last transaction (Transaction)`)
      }

      if (transactionIndexToKeypairMap) {
        const kp = transactionIndexToKeypairMap.get(i)
        if (kp && programId) {
          const isAddLiquidityTx = tx.instructions.some(
            (ix) => ix.programId.equals(programId),
          )
          if (isAddLiquidityTx) {
            const positionPubkey = kp.publicKey
            const isRequiredSigner = tx.instructions.some(ix =>
              ix.keys.some(key => key.pubkey.equals(positionPubkey) && key.isSigner)
            )
            if (isRequiredSigner) {
              positionKeypairSigningMap.set(i, kp)
            }
          }
        }
      } else if (positionKeypair && programId) {
        // Fallback to single position keypair (backward compatibility)
        const isAddLiquidityTx = tx.instructions.some(
          (ix) => ix.programId.equals(programId),
        )
        if (isAddLiquidityTx) {
          const positionPubkey = positionKeypair.publicKey
          const isRequiredSigner = tx.instructions.some(ix =>
            ix.keys.some(key => key.pubkey.equals(positionPubkey) && key.isSigner)
          )
          if (isRequiredSigner) {
            positionKeypairSigningMap.set(i, positionKeypair)
          }
        }
      }
      
      processedTxs.push(tx)
    }
  }

  // Serialize transactions for signing
  const serializedTxs = processedTxs.map((tx) => {
    if (tx instanceof VersionedTransaction) {
      return Buffer.from(tx.serialize()).toString('base64')
    } else {
      try {
        const serialized = tx.serialize({ requireAllSignatures: false })
        if (serialized.length > 1232) {
          throw new Error(`Transaction too large: ${serialized.length} bytes (max 1232 bytes)`)
        }
        return Buffer.from(serialized).toString('base64')
      } catch (error) {
        log.error(`${logPrefix} Failed to serialize transaction:`, error)
        throw error
      }
    }
  })

  // Sign all transactions with Privy
  const signedTxs = await Promise.all(
    serializedTxs.map((serializedTx) =>
      privy.wallets().solana().signTransaction(embeddedWalletId, {
        transaction: serializedTx,
        authorization_context: poolAuthContext,
      }),
    ),
  )

  // Verify all transactions were signed
  for (let i = 0; i < signedTxs.length; i++) {
    if (!signedTxs[i]?.signed_transaction) {
      throw new Error(`Transaction ${i} not signed`)
    }
  }

  // Deserialize signed transactions to VersionedTransaction
  const signedVersionedTxs = signedTxs.map((signedTx) =>
    VersionedTransaction.deserialize(
      Buffer.from(signedTx.signed_transaction!, 'base64'),
    ),
  )

  // Sign transactions that need position keypair AFTER Privy signing
  for (const [txIndex, kp] of positionKeypairSigningMap) {
    const tx = signedVersionedTxs[txIndex]
    if (tx) {
      tx.sign([kp])
    }
  }

  // Prepare encoded transactions for simulation
  const encodedTransactions = signedVersionedTxs.map(tx => 
    Buffer.from(tx.serialize()).toString('base64')
  )

  // Always simulate bundle before sending
  const simulationResult = await simulateBundle(encodedTransactions, logPrefix)
  
  if (!simulationResult) {
    throw new Error("Bundle simulation failed. Make sure you're using the correct pool and have sufficient balance.")
  }
  
  const summary = simulationResult.summary || simulationResult.value?.summary
  const isFailed = typeof summary === 'object' && summary.failed || (typeof summary === 'string' && summary === 'failed')
  
  if (isFailed) {
    const failedTx = typeof summary === 'object' && summary.failed
    let errorMessage = "Bundle simulation failed. "
    
    if (failedTx?.error) {
      const error = failedTx.error
      // Check for common error patterns
      if (Array.isArray(error) && error.length >= 2) {
        const errorMsg = error[1]
        if (typeof errorMsg === 'string') {
          if (errorMsg.includes('custom program error: 0x1')) {
            errorMessage += "Insufficient funds for rent or transaction fees. Each position creation requires SOL for rent. Please ensure you have enough SOL in your wallet."
          } else if (errorMsg.includes('Computational budget exceeded') || errorMsg.includes('computational budget exceeded')) {
            errorMessage += "Transaction exceeded compute budget. This may happen with very large position ranges or many positions. Please try with a smaller position range or fewer positions."
          } else if (errorMsg.includes('InsufficientFundsForFee')) {
            errorMessage += "Insufficient funds for transaction fees. Please ensure you have enough SOL in your wallet."
          } else if (errorMsg.includes('insufficient lamports')) {
            errorMessage += "Insufficient SOL. Please ensure you have enough SOL in your wallet for rent and fees."
          } else {
            errorMessage += `Transaction error: ${errorMsg}`
          }
        } else {
          errorMessage += "Transaction failed. Please check your wallet balance and try again."
        }
      } else {
        errorMessage += "Transaction failed. Please check your wallet balance and try again."
      }
    } else {
      errorMessage += "Make sure you're using the correct pool and have sufficient balance."
    }
    
    throw new Error(errorMessage)
  }
  
  if (LOG_SIMULATION) {
    log.info(`${logPrefix} Bundle simulation succeeded`)
  }

  // Create bundle (tip instruction is already in the last transaction)
  const bundle = new Bundle([], transactions.length)
  bundle.addTransactions(...signedVersionedTxs)

  // Send bundle with auto-retry on rate limit
  const bundleId = await sendBundleWithRetry(bundle, logPrefix)
  return bundleId
}

/**
 * Bundle transactions, send, and wait for confirmation
 * Returns bundle result with confirmation status and tx signatures
 */
export async function bundleAndConfirm(
  transactions: Array<Transaction | VersionedTransaction>,
  walletPubkey: PublicKey,
  embeddedWalletId: string,
  poolAuthContext: any,
  positionKeypair?: Keypair | null,
  programId?: PublicKey,
  logPrefix: string = '[bundle]',
  transactionNames?: string[],
  tipAmountLamports?: number,
  transactionIndexToKeypairMap?: Map<number, Keypair>,
): Promise<BundleResult> {
  try {
    const bundleId = await bundleAndSendTransactions(
      transactions,
      walletPubkey,
      embeddedWalletId,
      poolAuthContext,
      positionKeypair,
      programId,
      logPrefix,
      transactionNames,
      tipAmountLamports,
      transactionIndexToKeypairMap,
    )

    if (!bundleId) {
      return {
        bundleId: '',
        confirmed: false,
        txSignatures: [],
        error: 'Failed to send bundle',
      }
    }

    // Wait for confirmation
    const confirmation = await waitForBundleConfirmation(bundleId, 30, 2000, logPrefix)

    return {
      bundleId,
      confirmed: confirmation.confirmed,
      txSignatures: confirmation.txSignatures,
      error: confirmation.confirmed ? undefined : 'Bundle not confirmed',
    }
  } catch (error) {
    const errorMessage = error instanceof Error ? error.message : String(error)
    return {
      bundleId: '',
      confirmed: false,
      txSignatures: [],
      error: errorMessage,
    }
  }
}