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

import { ParsedTransactionWithMeta, PublicKey } from '@solana/web3.js'
import BN from 'bn.js'
import { getConnection } from '../solana/connection'
import { getUserAddress, signerCallback } from '../solana/wallet'
import { config } from '../config'
import {
  addCopyHistory,
  updateCopyHistory,
  upsertPositionMapping,
  getPositionMappingByTargetNft,
  getPositionMappingByTargetPosition,
  updatePositionMappingStatus,
  getWatchedAddressByAddress,
  getSetting,
} from '../db/queries'
import type { ParsedOperation } from '../monitor/types'
import { scaleAmount } from './ratio'
import { getTokenPrices, calculateCopyScale } from './price'
import { ensureSufficientBalances, getUsdcBalance, swapTokensBackToUsdc, USDC_MINT } from './swap'

// We import from the built SDK dist using relative path
// eslint-disable-next-line @typescript-eslint/no-require-imports
const { Chain, BYREAL_CLMM_PROGRAM_ID } = require('../clmm-sdk/dist/index.js')

function sleep(ms: number) {
  return new Promise((resolve) => setTimeout(resolve, ms))
}

/**
 * Extract our NFT mint from a createPosition TX.
 * The NFT mint is the second signer (non-wallet signer) in the transaction.
 */
async function extractOurNftMint(txid: string, walletAddress: string): Promise<string> {
  const connection = getConnection()
  // Wait a bit for TX to be confirmed
  await sleep(2000)

  for (let attempt = 0; attempt < 3; attempt++) {
    try {
      const tx: ParsedTransactionWithMeta | null = await connection.getParsedTransaction(txid, {
        commitment: 'confirmed',
        maxSupportedTransactionVersion: 0,
      })
      if (!tx) {
        if (attempt < 2) {
          await sleep(2000)
          continue
        }
        return ''
      }

      const accountKeys = tx.transaction.message.accountKeys
      for (const acc of accountKeys) {
        const pubkey = typeof acc === 'string' ? acc : 'pubkey' in acc ? acc.pubkey.toBase58() : ''
        const isSigner = typeof acc === 'string' ? false : 'signer' in acc ? (acc as { signer: boolean }).signer : false
        if (isSigner && pubkey && pubkey !== walletAddress) {
          return pubkey
        }
      }
    } catch (e) {
      if (attempt < 2) {
        await sleep(2000)
        continue
      }
    }
  }
  return ''
}

export class CopyEngine {
  private chain: InstanceType<typeof Chain>
  private connection = getConnection()

  constructor() {
    this.chain = new Chain({
      connection: this.connection,
      programId: BYREAL_CLMM_PROGRAM_ID,
    })
  }

  async executeCopy(operation: ParsedOperation): Promise<void> {
    console.log(`[CopyEngine] Executing ${operation.type} copy for tx ${operation.signature}`)

    switch (operation.type) {
      case 'open_position':
        return this.copyOpenPosition(operation)
      case 'add_liquidity':
        return this.copyAddLiquidity(operation)
      case 'decrease_liquidity':
        return this.copyDecreaseLiquidity(operation)
      case 'close_position':
        return this.copyClosePosition(operation)
    }
  }

  /**
   * 手动关仓（供 API 调用）
   */
  async manualClosePosition(ourNftMint: string, poolId: string): Promise<string> {
    const nftMint = new PublicKey(ourNftMint)
    const poolInfo = await this.chain.getRawPoolInfoByPoolId(new PublicKey(poolId))
    const mintA = poolInfo.mintA.toBase58()
    const mintB = poolInfo.mintB.toBase58()

    const txid = await this.chain.decreaseFullLiquidity({
      userAddress: getUserAddress(),
      nftMint,
      closePosition: true,
      slippage: 0.99,
      signerCallback,
    })

    console.log(`[CopyEngine] Manual close position TX: ${txid}`)

    // Swap received tokens back to USDC (if enabled)
    if (this.isSwapAfterCloseEnabled()) {
      await sleep(3000)
      await swapTokensBackToUsdc({
        connection: this.connection,
        mints: [mintA, mintB],
      })
    }

    return txid
  }

  /**
   * 检查关仓后是否需要 swap 回 USDC
   */
  private isSwapAfterCloseEnabled(): boolean {
    const val = getSetting('swap_after_close')
    return val !== 'false'
  }

  /**
   * 获取地址的倍率和最大值设置（优先用地址单独设置，否则用全局默认）
   */
  private getAddressSettings(signerAddress: string): { multiplier: number; maxUsd: number } {
    const addrRow = getWatchedAddressByAddress(signerAddress)
    return {
      multiplier: addrRow?.copy_multiplier ?? config.copyMultiplier,
      maxUsd: addrRow?.copy_max_usd ?? config.copyMaxUsd,
    }
  }

  /**
   * 检查 USDC 余额是否足够
   */
  private async checkUsdcBalance(requiredUsd: number): Promise<boolean> {
    const balance = await getUsdcBalance(this.connection)
    const balanceUsd = Number(balance) / 1e6
    if (balanceUsd < requiredUsd) {
      console.log(
        `[CopyEngine] Insufficient USDC balance: $${balanceUsd.toFixed(2)} < $${requiredUsd.toFixed(2)}, skipping`,
      )
      return false
    }
    return true
  }

  private async copyOpenPosition(op: ParsedOperation) {
    const historyId = addCopyHistory({
      targetAddress: op.signer,
      targetTxSig: op.signature,
      operation: 'open_position',
      targetNftMint: op.nftMint,
      poolId: op.poolId,
      targetAmountA: op.amountA,
      targetAmountB: op.amountB,
      status: 'executing',
    })

    try {
      // Get pool info
      const poolInfo = await this.chain.getRawPoolInfoByPoolId(new PublicKey(op.poolId))
      const mintA = op.mintA || poolInfo.mintA.toBase58()
      const mintB = op.mintB || poolInfo.mintB.toBase58()
      const decimalsA = poolInfo.mintDecimalsA as number
      const decimalsB = poolInfo.mintDecimalsB as number

      // Get tick range - prefer from event data, fall back to fetching position
      let tickLower = op.tickLower
      let tickUpper = op.tickUpper

      if (tickLower === undefined || tickUpper === undefined) {
        const posInfo = await this.chain.getRawPositionInfoByNftMint(new PublicKey(op.nftMint))
        if (posInfo) {
          tickLower = posInfo.tickLowerIndex
          tickUpper = posInfo.tickUpperIndex
        } else {
          throw new Error('Cannot determine tick range for position')
        }
      }

      // Get token prices and calculate copy scale
      const prices = await getTokenPrices(mintA, mintB)
      const priceA = prices[mintA]
      const priceB = prices[mintB]

      if (!priceA || !priceB) {
        throw new Error(`Cannot get token prices: A=${priceA || 'N/A'}, B=${priceB || 'N/A'}`)
      }

      const addrSettings = this.getAddressSettings(op.signer)

      const { ratio, targetUsd, ourUsd } = calculateCopyScale({
        targetAmountA: op.amountA || '0',
        targetAmountB: op.amountB || '0',
        decimalsA,
        decimalsB,
        priceA,
        priceB,
        multiplier: addrSettings.multiplier,
        maxUsd: addrSettings.maxUsd,
      })

      console.log(
        `[CopyEngine] Target: $${targetUsd.toFixed(2)}, Multiplier: ${addrSettings.multiplier}x, ` +
          `Max: $${addrSettings.maxUsd}, Our: $${ourUsd.toFixed(2)} (ratio: ${ratio.toFixed(4)})`,
      )

      if (ratio <= 0) {
        throw new Error('Copy ratio is zero - target position value is zero or prices unavailable')
      }

      // Check USDC balance before proceeding
      const neededUsd = mintA === USDC_MINT || mintB === USDC_MINT ? ourUsd * 0.6 : ourUsd
      if (!(await this.checkUsdcBalance(neededUsd))) {
        updateCopyHistory(historyId, {
          status: 'skipped',
          errorMessage: `Insufficient USDC balance for $${ourUsd.toFixed(2)} position`,
        })
        return
      }

      // Scale amounts by calculated ratio
      const scaledAmountA = op.amountA ? scaleAmount(new BN(op.amountA), ratio) : new BN(0)
      const scaledAmountB = op.amountB ? scaleAmount(new BN(op.amountB), ratio) : new BN(0)

      // Determine base token and otherAmountMax
      // 直接用 scaledAmount 作为 otherAmountMax（只是上限，SDK 不会多用）
      // 参考 byreal-copy: base='MintA', baseAmount=scaledAmount0, otherAmountMax=scaledAmount1
      const base = scaledAmountA.gt(new BN(0)) ? 'MintA' : 'MintB'
      const baseAmount = base === 'MintA' ? scaledAmountA : scaledAmountB
      const otherAmountMax = base === 'MintA' ? scaledAmountB : scaledAmountA

      if (baseAmount.isZero()) {
        throw new Error('Scaled amount is zero')
      }

      // Ensure token balances (ExactOut: check balance, swap deficit)
      const swapResult = await ensureSufficientBalances({
        connection: this.connection,
        tokenA: { mint: mintA, requiredAmount: scaledAmountA.toString() },
        tokenB: { mint: mintB, requiredAmount: scaledAmountB.toString() },
      })

      if (!swapResult.success) {
        throw new Error(`Token swap failed: ${swapResult.error}`)
      }

      if (swapResult.swapTxids.length > 0) {
        updateCopyHistory(historyId, { swapTxSig: swapResult.swapTxids.join(',') })
        await sleep(2000) // Wait for swap to settle
      }

      // Execute position creation with referer_position memo
      const txid = await this.chain.createPosition({
        userAddress: getUserAddress(),
        poolInfo,
        tickLower: tickLower!,
        tickUpper: tickUpper!,
        base,
        baseAmount,
        otherAmountMax,
        refererPosition: new PublicKey(op.personalPosition),
        signerCallback,
      })

      console.log(`[CopyEngine] Open position TX: ${txid}`)

      // Extract our NFT mint from the confirmed transaction
      const ourNftMint = await extractOurNftMint(txid, getUserAddress().toBase58())
      if (ourNftMint) {
        console.log(`[CopyEngine] Our NFT mint: ${ourNftMint}`)
      } else {
        console.warn(`[CopyEngine] Could not extract our NFT mint from TX ${txid}`)
      }

      updateCopyHistory(historyId, {
        ourNftMint: ourNftMint || undefined,
        ourTxSig: txid,
        ourAmountA: scaledAmountA.toString(),
        ourAmountB: scaledAmountB.toString(),
        status: 'success',
      })

      // Store position mapping with our NFT mint
      upsertPositionMapping({
        targetAddress: op.signer,
        targetNftMint: op.nftMint,
        targetPersonalPosition: op.personalPosition,
        ourNftMint: ourNftMint || undefined,
        poolId: op.poolId,
        tickLower: tickLower!,
        tickUpper: tickUpper!,
      })
    } catch (e) {
      const msg = e instanceof Error ? e.message : String(e)
      console.error(`[CopyEngine] Open position failed:`, msg)
      updateCopyHistory(historyId, { status: 'failed', errorMessage: msg })
    }
  }

  private async copyAddLiquidity(op: ParsedOperation) {
    // Find our corresponding position
    const mapping = op.nftMint
      ? getPositionMappingByTargetNft(op.nftMint)
      : getPositionMappingByTargetPosition(op.personalPosition)

    if (!mapping || !mapping.our_nft_mint) {
      console.log(`[CopyEngine] No matching position for add_liquidity, skipping`)
      addCopyHistory({
        targetAddress: op.signer,
        targetTxSig: op.signature,
        operation: 'add_liquidity',
        targetNftMint: op.nftMint,
        poolId: op.poolId,
        status: 'skipped',
      })
      return
    }

    const historyId = addCopyHistory({
      targetAddress: op.signer,
      targetTxSig: op.signature,
      operation: 'add_liquidity',
      targetNftMint: op.nftMint,
      poolId: op.poolId,
      targetAmountA: op.amountA,
      targetAmountB: op.amountB,
      status: 'executing',
    })

    try {
      const ourNftMint = new PublicKey(mapping.our_nft_mint)
      const poolInfo = await this.chain.getRawPoolInfoByPoolId(new PublicKey(mapping.pool_id))
      const mintA = op.mintA || poolInfo.mintA.toBase58()
      const mintB = op.mintB || poolInfo.mintB.toBase58()
      const decimalsA = poolInfo.mintDecimalsA as number
      const decimalsB = poolInfo.mintDecimalsB as number

      // Get token prices and calculate copy scale
      const prices = await getTokenPrices(mintA, mintB)
      const priceA = prices[mintA]
      const priceB = prices[mintB]

      if (!priceA || !priceB) {
        throw new Error(`Cannot get token prices: A=${priceA || 'N/A'}, B=${priceB || 'N/A'}`)
      }

      const addrSettings = this.getAddressSettings(op.signer)

      const { ratio, targetUsd, ourUsd } = calculateCopyScale({
        targetAmountA: op.amountA || '0',
        targetAmountB: op.amountB || '0',
        decimalsA,
        decimalsB,
        priceA,
        priceB,
        multiplier: addrSettings.multiplier,
        maxUsd: addrSettings.maxUsd,
      })

      console.log(
        `[CopyEngine] Add liquidity - Target: $${targetUsd.toFixed(2)}, Our: $${ourUsd.toFixed(2)} (ratio: ${ratio.toFixed(4)})`,
      )

      if (ratio <= 0) {
        throw new Error('Copy ratio is zero')
      }

      // Check USDC balance before proceeding
      const neededUsd = mintA === USDC_MINT || mintB === USDC_MINT ? ourUsd * 0.6 : ourUsd
      if (!(await this.checkUsdcBalance(neededUsd))) {
        updateCopyHistory(historyId, {
          status: 'skipped',
          errorMessage: `Insufficient USDC balance for $${ourUsd.toFixed(2)} add liquidity`,
        })
        return
      }

      // Scale amounts
      const scaledAmountA = op.amountA ? scaleAmount(new BN(op.amountA), ratio) : new BN(0)
      const scaledAmountB = op.amountB ? scaleAmount(new BN(op.amountB), ratio) : new BN(0)

      const base = scaledAmountA.gt(new BN(0)) ? 'MintA' : 'MintB'
      const baseAmount = base === 'MintA' ? scaledAmountA : scaledAmountB
      const otherAmountMax = base === 'MintA' ? scaledAmountB : scaledAmountA

      if (baseAmount.isZero()) {
        throw new Error('Scaled amount is zero')
      }

      // Ensure token balances (ExactOut)
      const swapResult = await ensureSufficientBalances({
        connection: this.connection,
        tokenA: { mint: mintA, requiredAmount: scaledAmountA.toString() },
        tokenB: { mint: mintB, requiredAmount: scaledAmountB.toString() },
      })

      if (!swapResult.success) {
        throw new Error(`Token swap failed: ${swapResult.error}`)
      }

      if (swapResult.swapTxids.length > 0) {
        await sleep(2000)
      }

      const txid = await this.chain.addLiquidity({
        userAddress: getUserAddress(),
        nftMint: ourNftMint,
        base,
        baseAmount,
        otherAmountMax,
        signerCallback,
      })

      console.log(`[CopyEngine] Add liquidity TX: ${txid}`)
      updateCopyHistory(historyId, {
        ourNftMint: mapping.our_nft_mint,
        ourTxSig: txid,
        ourAmountA: scaledAmountA.toString(),
        ourAmountB: scaledAmountB.toString(),
        status: 'success',
      })
    } catch (e) {
      const msg = e instanceof Error ? e.message : String(e)
      console.error(`[CopyEngine] Add liquidity failed:`, msg)
      updateCopyHistory(historyId, { status: 'failed', errorMessage: msg })
    }
  }

  private async copyDecreaseLiquidity(op: ParsedOperation) {
    const mapping = op.nftMint
      ? getPositionMappingByTargetNft(op.nftMint)
      : getPositionMappingByTargetPosition(op.personalPosition)

    if (!mapping || !mapping.our_nft_mint) {
      console.log(`[CopyEngine] No matching position for decrease_liquidity, skipping`)
      addCopyHistory({
        targetAddress: op.signer,
        targetTxSig: op.signature,
        operation: 'decrease_liquidity',
        status: 'skipped',
      })
      return
    }

    const historyId = addCopyHistory({
      targetAddress: op.signer,
      targetTxSig: op.signature,
      operation: 'decrease_liquidity',
      targetNftMint: op.nftMint,
      poolId: op.poolId,
      targetAmountA: op.amountA,
      targetAmountB: op.amountB,
      status: 'executing',
    })

    try {
      const ourNftMint = new PublicKey(mapping.our_nft_mint)

      // Get our position info
      const ourPositionInfo = await this.chain.getRawPositionInfoByNftMint(ourNftMint)
      if (!ourPositionInfo) {
        throw new Error('Our position not found on chain')
      }

      // 简单方式（参考 byreal-copy）：直接用 BN.min(目标减少量, 我们的流动性)
      let liquidityToDecrease: BN
      if (op.liquidity) {
        liquidityToDecrease = BN.min(new BN(op.liquidity), ourPositionInfo.liquidity)
      } else {
        // 没有流动性数据，减少 50%
        liquidityToDecrease = ourPositionInfo.liquidity.div(new BN(2))
      }

      if (liquidityToDecrease.isZero()) {
        throw new Error('Nothing to decrease')
      }

      console.log(
        `[CopyEngine] Decrease: target=${op.liquidity || 'N/A'}, ours=${ourPositionInfo.liquidity.toString()}, decreasing=${liquidityToDecrease.toString()}`,
      )

      // Use generous slippage for decrease — we're removing our own liquidity,
      // price can move significantly between detection and execution
      const txid = await this.chain.decreaseLiquidity({
        userAddress: getUserAddress(),
        nftMint: ourNftMint,
        liquidity: liquidityToDecrease,
        slippage: 0.99,
        signerCallback,
      })

      console.log(`[CopyEngine] Decrease liquidity TX: ${txid}`)
      updateCopyHistory(historyId, {
        ourNftMint: mapping.our_nft_mint,
        ourTxSig: txid,
        status: 'success',
      })
    } catch (e) {
      const msg = e instanceof Error ? e.message : String(e)
      console.error(`[CopyEngine] Decrease liquidity failed:`, msg)
      updateCopyHistory(historyId, { status: 'failed', errorMessage: msg })
    }
  }

  private async copyClosePosition(op: ParsedOperation) {
    const mapping = op.nftMint
      ? getPositionMappingByTargetNft(op.nftMint)
      : getPositionMappingByTargetPosition(op.personalPosition)

    if (!mapping || !mapping.our_nft_mint) {
      console.log(`[CopyEngine] No matching position for close_position, skipping`)
      addCopyHistory({
        targetAddress: op.signer,
        targetTxSig: op.signature,
        operation: 'close_position',
        status: 'skipped',
      })
      return
    }

    const historyId = addCopyHistory({
      targetAddress: op.signer,
      targetTxSig: op.signature,
      operation: 'close_position',
      targetNftMint: op.nftMint,
      poolId: mapping.pool_id,
      status: 'executing',
    })

    try {
      const ourNftMint = new PublicKey(mapping.our_nft_mint)
      const poolInfo = await this.chain.getRawPoolInfoByPoolId(new PublicKey(mapping.pool_id))
      const mintA = poolInfo.mintA.toBase58()
      const mintB = poolInfo.mintB.toBase58()

      const txid = await this.chain.decreaseFullLiquidity({
        userAddress: getUserAddress(),
        nftMint: ourNftMint,
        closePosition: true,
        slippage: 0.99,
        signerCallback,
      })

      console.log(`[CopyEngine] Close position TX: ${txid}`)
      updateCopyHistory(historyId, {
        ourNftMint: mapping.our_nft_mint,
        ourTxSig: txid,
        status: 'success',
      })
      updatePositionMappingStatus(mapping.target_nft_mint, 'closed')

      // Swap received tokens back to USDC (if enabled)
      if (this.isSwapAfterCloseEnabled()) {
        await sleep(3000)
        const swapBack = await swapTokensBackToUsdc({
          connection: this.connection,
          mints: [mintA, mintB],
        })
        if (swapBack.swapTxids.length > 0) {
          console.log(`[CopyEngine] Swapped back to USDC: ${swapBack.swapTxids.join(', ')}`)
        }
      }
    } catch (e) {
      const msg = e instanceof Error ? e.message : String(e)
      console.error(`[CopyEngine] Close position failed:`, msg)
      updateCopyHistory(historyId, { status: 'failed', errorMessage: msg })
    }
  }
}