byreal-table/src/lib/jupiter.ts

import {
	Connection,
	PublicKey,
	Keypair,
	VersionedTransaction,
} from '@solana/web3.js'
import ky from 'ky'
import { Token } from '@/lib/byreal-clmm-sdk/src/client/token'

// USDC 和 USDT 作为默认的输入代币（用于兑换）
const USDC_MINT = 'EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v'
const USDT_MINT = 'Es9vMFrzaCERmJfrF4H2FYD4KCoNkY11McCe8BenwNYB'

/** 稳定币 mint：用 USDC 换这些币时无需 swap */
const STABLECOIN_MINTS = [
	USDC_MINT,
	USDT_MINT,
	'USD1ttGY1N17NEEHLmELoaybftRBUSErhqYiQzvEmuB', // USD1
]

export interface JupiterQuote {
	inputMint: string
	inAmount: string
	outputMint: string
	outAmount: string
	swapMode: string
	slippageBps: number
	priceImpactPct: string
	routePlan: Array<{
		swapInfo: {
			label?: string
			[key: string]: unknown
		}
	}>
}

export interface JupiterSwapResponse {
	swapTransaction: string
}

/**
 * 获取代币余额（兼容 Token-2022 和标准 SPL Token）
 * 使用 Token.detectTokenTypeAndGetBalance，避免 Token-2022 代币查询返回 0
 */
export async function getTokenBalance(
	connection: Connection,
	walletAddress: PublicKey,
	mintAddress: string
): Promise<number> {
	try {
		const token = new Token(connection)
		const result = await token.detectTokenTypeAndGetBalance(
			walletAddress.toBase58(),
			mintAddress
		)
		return result.balance
	} catch (error) {
		console.error(`Error getting balance for ${mintAddress}:`, error)
		return 0
	}
}

/**
 * 获取 Jupiter API 请求头
 */
function getJupiterHeaders(): Record<string, string> {
	const headers: Record<string, string> = {
		Accept: 'application/json',
	}

	const apiKey = process.env.JUPITER_API_KEY
	if (apiKey) {
		headers['x-api-key'] = apiKey
	}

	return headers
}

/** Jupiter Price API v3 返回的单个代币价格（price.jup.ag 已弃用，请用 api.jup.ag/price/v3） */
const JUPITER_PRICE_API_URL =
	process.env.JUPITER_PRICE_API_URL || 'https://api.jup.ag/price/v3'

/**
 * 从 Jupiter Price API v3 获取代币 USD 价格（备用，避免使用已弃用的 price.jup.ag）
 * @returns 成功时返回 { tokenAPriceUsd, tokenBPriceUsd }，失败返回 null
 */
export async function getTokenPricesFromJupiter(
	mintA: string,
	mintB: string
): Promise<{ tokenAPriceUsd: number; tokenBPriceUsd: number } | null> {
	try {
		const ids = [mintA, mintB].filter((m, i, arr) => arr.indexOf(m) === i).join(',')
		const data = await ky
			.get(`${JUPITER_PRICE_API_URL}?ids=${ids}`, {
				headers: getJupiterHeaders(),
				timeout: 10000,
			})
			.json<Record<string, { price?: number; usdPrice?: number }>>()
		const priceOf = (mint: string): number | undefined => {
			const row = data[mint]
			if (!row) return undefined
			return row.usdPrice ?? row.price
		}
		const pa = priceOf(mintA)
		const pb = priceOf(mintB)
		if (pa != null && pa > 0 && pb != null && pb > 0) {
			return { tokenAPriceUsd: pa, tokenBPriceUsd: pb }
		}
		return null
	} catch (error) {
		console.warn('Jupiter price v3 fetch failed:', error)
		return null
	}
}

/**
 * 从 Jupiter API 获取 quote
 * @param restrictIntermediate - 为 false 时允许更多中间代币路由（用于 NO_ROUTES_FOUND 时重试）
 */
export async function fetchJupiterQuote(
	inputMint: string,
	outputMint: string,
	amount: string | number,
	swapMode: 'ExactIn' | 'ExactOut' = 'ExactIn',
	slippageBps: number = 200, // 2%
	restrictIntermediate: boolean = true
): Promise<JupiterQuote> {
	const jupiterBaseUrl =
		process.env.JUPITER_API_URL || 'https://api.jup.ag/swap/v1'

	const rawAmount =
		typeof amount === 'string' ? amount : String(Math.floor(Number(amount)))

	const searchParams = new URLSearchParams({
		inputMint,
		outputMint,
		amount: rawAmount,
		swapMode,
		slippageBps: String(slippageBps),
		onlyDirectRoutes: 'false',
		restrictIntermediateTokens: restrictIntermediate ? 'true' : 'false',
	})

	const response = await ky
		.get(`${jupiterBaseUrl}/quote`, {
			searchParams,
			timeout: 30000,
			headers: getJupiterHeaders(),
		})
		.json<JupiterQuote>()

	return response
}

/**
 * 执行 Jupiter swap
 */
export async function executeJupiterSwap(
	quoteData: JupiterQuote,
	walletAddress: string
): Promise<JupiterSwapResponse> {
	const jupiterBaseUrl =
		process.env.JUPITER_API_URL || 'https://api.jup.ag/swap/v1'

	const headers = {
		...getJupiterHeaders(),
		'Content-Type': 'application/json',
	}

	const response = await ky
		.post(`${jupiterBaseUrl}/swap`, {
			json: {
				quoteResponse: quoteData,
				userPublicKey: walletAddress,
				wrapAndUnwrapSol: true,
				prioritizationFeeLamports: 10000,
				dynamicSlippage: false,
			},
			timeout: 30000,
			headers,
		})
		.json<JupiterSwapResponse>()

	return response
}

/**
 * 按 USD 金额执行 swap（ExactIn：花掉指定美元 USDC，换回目标代币）
 * 不使用 ExactOut，不使用代币数量 amount，仅使用 valueUsd。
 * @param usdValue - 要换入的美元数（例如 5.5 表示花 $5.5 USDC）
 */
export async function swapIfNeeded(
	connection: Connection,
	keypair: Keypair,
	outputMint: string,
	usdValue: number,
	inputMint: string = USDC_MINT
): Promise<{ success: boolean; txid?: string; error?: string }> {
	const walletAddress = keypair.publicKey

	if (!usdValue || Number(usdValue) <= 0) {
		console.log(`Skip swap: USD value is ${usdValue}`)
		return { success: true }
	}

	const mint = String(outputMint ?? '').trim()
	const isStablecoin =
		STABLECOIN_MINTS.includes(mint) || mint === USDC_MINT
	if (isStablecoin) {
		console.log(
			`Skip swap: output is stablecoin (USDC/USDT), no need to swap USDC -> same`
		)
		return { success: true }
	}

	const inputAmountRaw = Math.floor(Number(usdValue) * 1e6)
	if (inputAmountRaw < 1e6) {
		console.log(`Skip swap: USD value too small (${usdValue})`)
		return { success: true }
	}

	console.log(
		`Swap: $${usdValue} USDC -> ${outputMint.slice(0, 8)}... (ExactIn)`
	)

	const slippageBps = 200 // 2%

	const fetchQuoteWithRetry = async (
		restrictIntermediate: boolean
	): Promise<JupiterQuote> => {
		try {
			return await fetchJupiterQuote(
				inputMint,
				outputMint,
				inputAmountRaw,
				'ExactIn',
				slippageBps,
				restrictIntermediate
			)
		} catch (e) {
			const msg = e instanceof Error ? e.message : String(e)
			if (
				msg.includes('No routes found') &&
				restrictIntermediate
			) {
				console.log(
					'No routes with restrictIntermediate=true, retrying with allow all intermediates...'
				)
				return fetchJupiterQuote(
					inputMint,
					outputMint,
					inputAmountRaw,
					'ExactIn',
					slippageBps,
					false
				)
			}
			throw e
		}
	}

	try {
		const quoteData = await fetchQuoteWithRetry(true)

		if (quoteData.routePlan && quoteData.routePlan.length > 0) {
			const routeLabels = quoteData.routePlan
				.map((r) => r.swapInfo?.label || 'Unknown')
				.join(' -> ')
			console.log(`Route: ${routeLabels}`)
			console.log(
				`Expected output: ${quoteData.outAmount} (ExactIn), price impact: ${quoteData.priceImpactPct}%`
			)
		}

		const swapData = await executeJupiterSwap(
			quoteData,
			walletAddress.toBase58()
		)

		const swapTransactionBuf = Buffer.from(swapData.swapTransaction, 'base64')
		const transaction = VersionedTransaction.deserialize(swapTransactionBuf)
		transaction.sign([keypair])

		const signature = await connection.sendTransaction(transaction, {
			maxRetries: 3,
			skipPreflight: false,
		})

		console.log(`Swap transaction sent: ${signature}`)

		const confirmation = await connection.confirmTransaction(
			signature,
			'confirmed'
		)

		if (confirmation.value.err) {
			throw new Error(`Transaction failed: ${confirmation.value.err}`)
		}

		console.log(`Swap confirmed: https://solscan.io/tx/${signature}`)
		const newBalance = await getTokenBalance(
			connection,
			walletAddress,
			outputMint
		)
		console.log(`New balance: ${newBalance.toFixed(6)}`)

		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,
		}
	}
}

/**
 * 确保 LP 所需的两种代币都有足够余额
 * 按 valueUsd 换币：使用 ExactIn，花 $valueUsd USDC 换目标代币，不使用 amount
 */
export async function ensureSufficientBalances(
	connection: Connection,
	keypair: Keypair,
	tokenA: { mint: string; valueUsd: number },
	tokenB: { mint: string; valueUsd: number }
): Promise<{
	success: boolean
	swapTxids: string[]
	error?: string
}> {
	const swapTxids: string[] = []

	// Token A：按 valueUsd 换
	console.log(
		`\n--- Token A (${tokenA.mint.slice(0, 8)}...): swap $${tokenA.valueUsd} USDC ---`
	)
	const resultA = await swapIfNeeded(
		connection,
		keypair,
		tokenA.mint,
		tokenA.valueUsd,
		USDC_MINT
	)

	if (!resultA.success) {
		console.log('USDC swap failed for Token A, trying USDT...')
		const resultA_USDT = await swapIfNeeded(
			connection,
			keypair,
			tokenA.mint,
			tokenA.valueUsd,
			USDT_MINT
		)

		if (!resultA_USDT.success) {
			return {
				success: false,
				swapTxids,
				error: `Failed to get Token A: ${resultA.error || resultA_USDT.error}`,
			}
		}

		if (resultA_USDT.txid) swapTxids.push(resultA_USDT.txid)
	} else if (resultA.txid) {
		swapTxids.push(resultA.txid)
	}

	// Token B：按 valueUsd 换
	console.log(
		`\n--- Token B (${tokenB.mint.slice(0, 8)}...): swap $${tokenB.valueUsd} USDC ---`
	)
	const resultB = await swapIfNeeded(
		connection,
		keypair,
		tokenB.mint,
		tokenB.valueUsd,
		USDC_MINT
	)

	if (!resultB.success) {
		console.log('USDC swap failed for Token B, trying USDT...')
		const resultB_USDT = await swapIfNeeded(
			connection,
			keypair,
			tokenB.mint,
			tokenB.valueUsd,
			USDT_MINT
		)

		if (!resultB_USDT.success) {
			return {
				success: false,
				swapTxids,
				error: `Failed to get Token B: ${resultB.error || resultB_USDT.error}`,
			}
		}

		if (resultB_USDT.txid) swapTxids.push(resultB_USDT.txid)
	} else if (resultB.txid) {
		swapTxids.push(resultB.txid)
	}

	return { success: true, swapTxids }
}