byreal-trade-copy/src/lib/copier/swap.ts

import { Connection, PublicKey, VersionedTransaction } from '@solana/web3.js'
import ky from 'ky'
import { config } from '../config'
import { getKeypair, getUserAddress } from '../solana/wallet'
import { getJupiterHeaders, withRetry } from './jupiter-client'

export const USDC_MINT = 'EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v'
const USDT_MINT = 'Es9vMFrzaCERmJfrF4H2FYD4KCoNkY11McCe8BenwNYB'
const SOL_MINT = 'So11111111111111111111111111111111111111112'
const USD1_MINT = 'USD1ttGY1N17NEEHLmELoaybftRBUSErhqYiQzvEmuB'

/** 稳定币 mint：不需要 swap 的 */
const STABLECOIN_MINTS = new Set([USDC_MINT, USDT_MINT, USD1_MINT])

/** 不需要 swap 回 USDC 的 mint 列表 */
const SKIP_SWAP_BACK_MINTS = new Set([USDC_MINT, USDT_MINT, SOL_MINT, USD1_MINT])

const ULTRA_API_URL = 'https://api.jup.ag/ultra/v1'
const SWAP_API_URL = 'https://api.jup.ag/swap/v1'

/**
 * Ultra execute negative error codes that are safe to retry with the same
 * signedTransaction + requestId (idempotent within the ~2 min TTL).
 * Ref: https://dev.jup.ag/docs/ultra/response.md
 */
const ULTRA_RETRYABLE_CODES = new Set([-1, -1000, -1001, -1005, -1006, -2000, -2003, -2005])

/** Warn when a signed Ultra order payload approaches its ~2 min TTL. */
const ULTRA_ORDER_TTL_WARN_MS = 90_000

interface UltraOrderResponse {
  requestId: string
  inputMint: string
  outputMint: string
  inAmount: string
  outAmount: string
  otherAmountThreshold: string
  swapMode: string
  slippageBps: number
  priceImpactPct: string
  routePlan: Array<{
    swapInfo: {
      ammKey: string
      label: string
      inputMint: string
      outputMint: string
      inAmount: string
      outAmount: string
    }
    percent: number
  }>
  transaction: string | null
  router: string
  gasless: boolean
  errorCode?: number
  errorMessage?: string
}

interface UltraExecuteResponse {
  status: 'Success' | 'Failed'
  signature?: string
  error?: string
  code: number
}

/**
 * 获取 Jupiter Ultra Order（GET 请求，返回 unsigned transaction）
 */
async function fetchUltraOrder(
  inputMint: string,
  outputMint: string,
  amount: string | number,
  taker: string,
  slippageBps: number = 200,
  swapMode: 'ExactIn' | 'ExactOut' = 'ExactIn',
): Promise<UltraOrderResponse> {
  const rawAmount = typeof amount === 'string' ? amount : String(Math.floor(Number(amount)))

  return withRetry(() =>
    ky
      .get(`${ULTRA_API_URL}/order`, {
        searchParams: {
          inputMint,
          outputMint,
          amount: rawAmount,
          taker,
          slippageBps: String(slippageBps),
          swapMode,
        },
        timeout: 30000,
        headers: getJupiterHeaders(),
      })
      .json<UltraOrderResponse>(),
  )
}

/**
 * 执行 Jupiter Ultra swap（POST 签名后的交易）
 */
async function executeUltraSwap(
  signedTransaction: string,
  requestId: string,
): Promise<UltraExecuteResponse> {
  return withRetry(() =>
    ky
      .post(`${ULTRA_API_URL}/execute`, {
        json: { signedTransaction, requestId },
        timeout: 30000,
        headers: {
          ...getJupiterHeaders(),
          'Content-Type': 'application/json',
        },
      })
      .json<UltraExecuteResponse>(),
  )
}

/**
 * Jupiter Swap v1 API: GET /quote (supports ExactOut)
 */
async function fetchSwapQuote(
  inputMint: string,
  outputMint: string,
  amount: string,
  slippageBps: number,
  swapMode: 'ExactIn' | 'ExactOut',
): Promise<Record<string, unknown>> {
  return withRetry(() =>
    ky
      .get(`${SWAP_API_URL}/quote`, {
        searchParams: {
          inputMint,
          outputMint,
          amount,
          slippageBps: String(slippageBps),
          swapMode,
        },
        timeout: 30000,
        headers: getJupiterHeaders(),
      })
      .json<Record<string, unknown>>(),
  )
}

/**
 * Jupiter Swap v1 API: POST /swap (returns unsigned transaction)
 */
async function fetchSwapTransaction(
  quoteResponse: Record<string, unknown>,
  userPublicKey: string,
): Promise<{ swapTransaction: string; lastValidBlockHeight: number }> {
  return withRetry(() =>
    ky
      .post(`${SWAP_API_URL}/swap`, {
        json: {
          quoteResponse,
          userPublicKey,
          wrapAndUnwrapSol: true,
          dynamicComputeUnitLimit: true,
        },
        timeout: 30000,
        headers: {
          ...getJupiterHeaders(),
          'Content-Type': 'application/json',
        },
      })
      .json<{ swapTransaction: string; lastValidBlockHeight: number }>(),
  )
}

/**
 * 签名并执行 swap
 * - ExactIn: 使用 Jupiter Ultra API（更快，含幂等重试）
 * - ExactOut: 使用 Jupiter Swap v1 API（Ultra 不支持 ExactOut）
 */
async function signAndExecuteSwap(
  inputMint: string,
  outputMint: string,
  amount: string,
  slippageBps: number = 200,
  swapMode: 'ExactIn' | 'ExactOut' = 'ExactIn',
): Promise<{ success: boolean; txid?: string; error?: string }> {
  const keypair = getKeypair()

  try {
    if (swapMode === 'ExactOut') {
      return await signAndExecuteSwapV1(inputMint, outputMint, amount, slippageBps, swapMode)
    }

    // ExactIn: use Ultra API
    const orderData = await fetchUltraOrder(
      inputMint,
      outputMint,
      amount,
      keypair.publicKey.toBase58(),
      slippageBps,
      swapMode,
    )
    const orderFetchedAt = Date.now()

    if (!orderData.transaction) {
      const errorMsg = orderData.errorMessage || 'No transaction returned'
      console.error(`[Swap] Ultra order error: ${errorMsg}`)
      return { success: false, error: errorMsg }
    }

    if (orderData.routePlan?.length > 0) {
      const routeLabels = orderData.routePlan.map((r) => r.swapInfo?.label || '?').join(' -> ')
      console.log(
        `[Swap] ${swapMode} Route: ${routeLabels}, in: ${orderData.inAmount}, out: ${orderData.outAmount}`,
      )
    }

    const txBuf = Buffer.from(orderData.transaction, 'base64')
    const transaction = VersionedTransaction.deserialize(txBuf)
    transaction.sign([keypair])
    const signedTx = Buffer.from(transaction.serialize()).toString('base64')

    // Warn if signed payload is approaching the ~2 min TTL
    const elapsed = Date.now() - orderFetchedAt
    if (elapsed > ULTRA_ORDER_TTL_WARN_MS) {
      console.warn(
        `[Swap] Ultra order payload is ${Math.round(elapsed / 1000)}s old (TTL ~2 min), may be stale`,
      )
    }

    // Execute with retry on retryable negative codes.
    // Ultra /execute is idempotent for the same signedTransaction + requestId within ~2 min TTL.
    let executeResult = await executeUltraSwap(signedTx, orderData.requestId)

    for (
      let retry = 0;
      retry < 2 &&
      executeResult.status !== 'Success' &&
      ULTRA_RETRYABLE_CODES.has(executeResult.code);
      retry++
    ) {
      const delay = 1000 * Math.pow(2, retry)
      console.warn(
        `[Swap] Ultra execute retryable code ${executeResult.code}, retrying in ${delay}ms (retry ${retry + 1}/2)`,
      )
      await new Promise((r) => setTimeout(r, delay))
      executeResult = await executeUltraSwap(signedTx, orderData.requestId)
    }

    if (executeResult.status !== 'Success') {
      throw new Error(
        `Execute failed [code=${executeResult.code}]: ${executeResult.error || 'Unknown error'}`,
      )
    }

    const signature = executeResult.signature!
    console.log(`[Swap] Confirmed: ${signature}`)
    return { success: true, txid: signature }
  } catch (error) {
    const errorMessage = error instanceof Error ? error.message : 'Unknown swap error'
    console.error(`[Swap] Failed:`, errorMessage)
    return { success: false, error: errorMessage }
  }
}

/**
 * ExactOut swap via Jupiter Swap v1 API.
 * If ExactOut is not supported for the pair (400), falls back to ExactIn
 * by estimating the input amount from a reference quote + 15% buffer.
 */
async function signAndExecuteSwapV1(
  inputMint: string,
  outputMint: string,
  amount: string,
  slippageBps: number,
  swapMode: 'ExactIn' | 'ExactOut',
): Promise<{ success: boolean; txid?: string; error?: string }> {
  const keypair = getKeypair()
  const connection = (await import('../solana/connection')).getConnection()

  try {
    let quoteResponse: Record<string, unknown>

    if (swapMode === 'ExactOut') {
      try {
        // Try ExactOut first
        quoteResponse = await fetchSwapQuote(inputMint, outputMint, amount, slippageBps, 'ExactOut')
      } catch {
        // ExactOut not supported for this pair, fall back to ExactIn
        console.log('[Swap] ExactOut not available, estimating ExactIn amount via reference quote')
        quoteResponse = await estimateExactInQuote(inputMint, outputMint, amount, slippageBps)
      }
    } else {
      quoteResponse = await fetchSwapQuote(inputMint, outputMint, amount, slippageBps, swapMode)
    }

    console.log(
      `[Swap] v1 quote: in=${quoteResponse.inAmount}, out=${quoteResponse.outAmount}, mode=${quoteResponse.swapMode}`,
    )

    // Get swap transaction
    const { swapTransaction } = await fetchSwapTransaction(
      quoteResponse,
      keypair.publicKey.toBase58(),
    )

    // Fetch blockhash BEFORE sending so the confirmation window is correctly anchored.
    // Using maxRetries: 0 (per Jupiter quickstart) to avoid duplicate submissions.
    const latestBlockhash = await connection.getLatestBlockhash('confirmed')

    // Sign and send
    const txBuf = Buffer.from(swapTransaction, 'base64')
    const transaction = VersionedTransaction.deserialize(txBuf)
    transaction.sign([keypair])

    const signature = await connection.sendRawTransaction(transaction.serialize(), {
      skipPreflight: true,
      maxRetries: 0,
    })

    // Confirm using the blockhash fetched before send
    await connection.confirmTransaction({ signature, ...latestBlockhash }, 'confirmed')

    console.log(`[Swap] v1 Confirmed: ${signature}`)
    return { success: true, txid: signature }
  } catch (error) {
    const errorMessage = error instanceof Error ? error.message : 'Unknown swap error'
    console.error(`[Swap] v1 Failed:`, errorMessage)
    return { success: false, error: errorMessage }
  }
}

/**
 * ExactOut 不可用时，通过参考报价估算所需 ExactIn 输入量。
 * 1. 用小额 ExactIn 报价获取汇率
 * 2. 按汇率反推所需输入量，加 15% buffer
 * 3. 用计算出的输入量获取正式 ExactIn 报价
 *
 * @param inputDecimals - input token decimals (default 6 for USDC/USDT)
 */
async function estimateExactInQuote(
  inputMint: string,
  outputMint: string,
  desiredOutputAmount: string,
  slippageBps: number,
  inputDecimals: number = 6,
): Promise<Record<string, unknown>> {
  // Use 10 units of input token as reference to get the exchange rate.
  // Respects actual token decimals (e.g. 10 USDC = 10_000_000 for 6-decimal; 10 SOL = 10_000_000_000 for 9-decimal).
  const refInputAmount = String(Math.round(10 * Math.pow(10, inputDecimals)))
  const refQuote = await fetchSwapQuote(inputMint, outputMint, refInputAmount, slippageBps, 'ExactIn')

  const refIn = BigInt(refQuote.inAmount as string)
  const refOut = BigInt(refQuote.outAmount as string)

  if (refOut === 0n) {
    throw new Error('Reference quote returned 0 output, cannot estimate rate')
  }

  // estimatedInput = desiredOutput * (refIn / refOut) * 1.15
  const desiredOut = BigInt(desiredOutputAmount)
  const estimatedInput = (desiredOut * refIn * 115n) / (refOut * 100n)

  console.log(
    `[Swap] Rate estimate: ${refIn}→${refOut}, need ${desiredOut} out, estimated input ${estimatedInput} (+15% buffer)`,
  )

  // Get actual ExactIn quote with estimated input amount
  return fetchSwapQuote(inputMint, outputMint, estimatedInput.toString(), slippageBps, 'ExactIn')
}

/**
 * 获取代币余额（raw amount）
 */
async function getTokenBalance(connection: Connection, mint: string): Promise<bigint> {
  const userAddr = getUserAddress()
  if (mint === SOL_MINT) {
    const balance = await connection.getBalance(userAddr)
    return BigInt(balance)
  }
  try {
    const accounts = await connection.getParsedTokenAccountsByOwner(userAddr, {
      mint: new PublicKey(mint),
    })
    let balance = 0n
    for (const acc of accounts.value) {
      balance += BigInt(acc.account.data.parsed.info.tokenAmount.amount)
    }
    return balance
  } catch {
    return 0n
  }
}

/**
 * 确保某个代币有足够余额，不足则通过 USDC -> token swap 补齐（ExactOut 模式）
 * 参考 byreal-copy: 查余额 → 算差额 → 加 5% buffer → ExactOut swap
 * @returns true 如果执行了 swap
 */
async function ensureTokenBalance(params: {
  connection: Connection
  tokenMint: string
  requiredAmount: bigint
}): Promise<{ swapped: boolean; txid?: string; error?: string }> {
  const { connection, tokenMint, requiredAmount } = params

  if (requiredAmount <= 0n) {
    return { swapped: false }
  }

  // 稳定币不需要 swap（USDC 本身、USDT、USD1）
  if (STABLECOIN_MINTS.has(tokenMint)) {
    return { swapped: false }
  }

  const currentBalance = await getTokenBalance(connection, tokenMint)
  if (currentBalance >= requiredAmount) {
    console.log(
      `[Swap] ${tokenMint.slice(0, 8)}...: have ${currentBalance}, need ${requiredAmount}, sufficient`,
    )
    return { swapped: false }
  }

  const deficit = requiredAmount - currentBalance
  // 加 2% buffer 防止余额微小波动
  const swapAmount = (deficit * 102n) / 100n

  console.log(
    `[Swap] ${tokenMint.slice(0, 8)}...: have ${currentBalance}, need ${requiredAmount}, deficit ${deficit}, swap(ExactOut) ${swapAmount}`,
  )

  // 优先 USDC → token（ExactOut: amount = 需要获得的代币数量）
  let result = await signAndExecuteSwap(USDC_MINT, tokenMint, swapAmount.toString(), 300, 'ExactOut')
  if (result.success) {
    return { swapped: true, txid: result.txid }
  }

  // 失败则尝试 USDT → token
  console.log('[Swap] USDC failed, trying USDT...')
  result = await signAndExecuteSwap(USDT_MINT, tokenMint, swapAmount.toString(), 300, 'ExactOut')
  if (result.success) {
    return { swapped: true, txid: result.txid }
  }

  return { swapped: false, error: result.error }
}

/**
 * 确保 LP 所需的两种代币都有足够余额（ExactOut 模式，按实际代币数量 swap）
 * 查当前余额 → 计算差额 → 加 buffer → USDC/USDT swap 补齐
 */
export async function ensureSufficientBalances(params: {
  connection: Connection
  tokenA: { mint: string; requiredAmount: string }
  tokenB: { mint: string; requiredAmount: string }
}): Promise<{ success: boolean; swapTxids: string[]; error?: string }> {
  const { connection, tokenA, tokenB } = params
  const swapTxids: string[] = []

  // Token A
  const resultA = await ensureTokenBalance({
    connection,
    tokenMint: tokenA.mint,
    requiredAmount: BigInt(tokenA.requiredAmount || '0'),
  })
  if (resultA.error) {
    return { success: false, swapTxids, error: `Failed to get Token A: ${resultA.error}` }
  }
  if (resultA.txid) swapTxids.push(resultA.txid)

  // Token B
  const resultB = await ensureTokenBalance({
    connection,
    tokenMint: tokenB.mint,
    requiredAmount: BigInt(tokenB.requiredAmount || '0'),
  })
  if (resultB.error) {
    return { success: false, swapTxids, error: `Failed to get Token B: ${resultB.error}` }
  }
  if (resultB.txid) swapTxids.push(resultB.txid)

  return { success: true, swapTxids }
}

/**
 * 获取 USDC 余额（raw amount, 6 decimals）
 */
export async function getUsdcBalance(connection: Connection): Promise<bigint> {
  return getTokenBalance(connection, USDC_MINT)
}

/**
 * 关仓后将收到的代币换回 USDC
 * 跳过 SOL、USDC、USDT、USD1
 */
export async function swapTokensBackToUsdc(params: {
  connection: Connection
  mints: string[]
}): Promise<{ swapTxids: string[] }> {
  const { connection, mints } = params
  const swapTxids: string[] = []

  for (const mint of mints) {
    if (SKIP_SWAP_BACK_MINTS.has(mint)) continue

    try {
      const balance = await getTokenBalance(connection, mint)
      if (balance <= 0n) continue

      console.log(`[Swap] Swapping ${balance.toString()} of ${mint.slice(0, 8)}... back to USDC`)

      const result = await signAndExecuteSwap(mint, USDC_MINT, balance.toString(), 300)
      if (result.success && result.txid) {
        console.log(`[Swap] Swapped ${mint.slice(0, 8)}... -> USDC: ${result.txid}`)
        swapTxids.push(result.txid)
      } else {
        console.warn(`[Swap] Swap ${mint.slice(0, 8)}... -> USDC failed: ${result.error}`)
      }
    } catch (e) {
      console.warn(`[Swap] Error swapping ${mint.slice(0, 8)}... -> USDC:`, e)
    }
  }

  return { swapTxids }
}