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
}

export interface UltraOrderResponse {
	mode: string
	inputMint: string
	outputMint: string
	inAmount: string
	outAmount: string
	inUsdValue?: number
	outUsdValue?: number
	priceImpact?: number
	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
		bps: number
		usdValue?: number
	}>
	feeMint?: string
	feeBps?: number
	platformFee?: {
		amount: string
		feeBps: number
	}
	signatureFeeLamports: number
	signatureFeePayer: string | null
	prioritizationFeeLamports: number
	prioritizationFeePayer: string | null
	rentFeeLamports: number
	rentFeePayer: string | null
	router: string
	transaction: string | null
	gasless: boolean
	requestId: string
	totalTime: number
	taker: string | null
	quoteId?: string
	maker?: string
	expireAt?: string
	errorCode?: number
	errorMessage?: string
}

export interface UltraExecuteResponse {
	status: 'Success' | 'Failed'
	signature?: string
	slot?: string
	error?: string
	code: number
	totalInputAmount?: string
	totalOutputAmount?: string
	inputAmountResult?: string
	outputAmountResult?: string
	swapEvents?: Array<{
		inputMint: string
		inputAmount: string
		outputMint: string
		outputAmount: 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 返回的单个代币价格（需 API Key，从 https://portal.jup.ag/ 获取） */
const JUPITER_PRICE_API_URL =
	process.env.JUPITER_PRICE_API_URL || 'https://api.jup.ag/price/v3'

/**
 * 从 Jupiter Price API v3 获取代币 USD 价格（需配置环境变量 JUPITER_API_KEY）
 * @returns 成功时返回 { tokenAPriceUsd, tokenBPriceUsd }，失败返回 null
 */
export async function getTokenPricesFromJupiter(
	mintA: string,
	mintB: string
): Promise<{ tokenAPriceUsd: number; tokenBPriceUsd: number } | null> {
	try {
		if (!process.env.JUPITER_API_KEY) {
			console.warn(
				'Jupiter Price API v3 需要 API Key，请在 .env 中设置 JUPITER_API_KEY（从 https://portal.jup.ag/ 获取）'
			)
			return null
		}
		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: unknown) {
		const is401 =
			error &&
			typeof error === 'object' &&
			'response' in error &&
			(error as { response?: { status?: number } }).response?.status === 401
		if (is401) {
			console.warn(
				'Jupiter Price API v3 返回 401：请在 .env 中设置有效的 JUPITER_API_KEY（从 https://portal.jup.ag/ 获取）'
			)
		} else {
			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
}

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

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

	const searchParams = new URLSearchParams({
		inputMint,
		outputMint,
		amount: rawAmount,
		taker,
		slippageBps: String(slippageBps),
	})

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

	return response
}

/**
 * 执行 Jupiter Ultra swap
 */
export async function executeUltraSwap(
	signedTransaction: string,
	requestId: string
): Promise<UltraExecuteResponse> {
	const headers = {
		...getJupiterHeaders(),
		'Content-Type': 'application/json',
	}

	const response = await ky
		.post(`${ULTRA_API_URL}/execute`, {
			json: {
				signedTransaction,
				requestId,
			},
			timeout: 30000,
			headers,
		})
		.json<UltraExecuteResponse>()

	return response
}

/**
 * 按 USD 金额执行 swap（使用 Jupiter Ultra API）
 * @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()
	if (mint === USDC_MINT) {
		console.log(`Skip swap: output is USDC, no need to swap USDC -> USDC`)
		return { success: true }
	}

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

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

	const slippageBps = 200

	try {
		const orderData = await fetchUltraOrder(
			inputMint,
			outputMint,
			inputAmountRaw,
			walletAddress.toBase58(),
			slippageBps
		)

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

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

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

		const signedTransaction = Buffer.from(transaction.serialize()).toString(
			'base64'
		)

		const executeResult = await executeUltraSwap(
			signedTransaction,
			orderData.requestId
		)

		if (executeResult.status !== 'Success') {
			throw new Error(executeResult.error || 'Execute failed')
		}

		const signature = executeResult.signature!
		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
 * 会先检查用户现有余额，如有则扣除相应的 swap 金额
 */
export async function ensureSufficientBalances(
	connection: Connection,
	keypair: Keypair,
	tokenA: {
		mint: string
		valueUsd: number
		priceUsd: number
		decimals: number
	},
	tokenB: { mint: string; valueUsd: number; priceUsd: number; decimals: number }
): Promise<{
	success: boolean
	swapTxids: string[]
	error?: string
}> {
	const swapTxids: string[] = []
	const walletAddress = keypair.publicKey

	// Token A：检查现有余额，计算需要 swap 的金额
	const balanceA = await getTokenBalance(connection, walletAddress, tokenA.mint)
	const existingValueA = balanceA * tokenA.priceUsd
	const needSwapA = Math.max(0, tokenA.valueUsd - existingValueA)

	console.log(
		`\n--- Token A (${tokenA.mint.slice(0, 8)}...): need $${tokenA.valueUsd}, have $${existingValueA.toFixed(4)}, swap $${needSwapA.toFixed(4)} ---`
	)

	const resultA = await swapIfNeeded(
		connection,
		keypair,
		tokenA.mint,
		needSwapA,
		USDC_MINT
	)

	if (!resultA.success) {
		console.log('USDC swap failed for Token A, trying USDT...')
		const resultA_USDT = await swapIfNeeded(
			connection,
			keypair,
			tokenA.mint,
			needSwapA,
			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：检查现有余额，计算需要 swap 的金额
	const balanceB = await getTokenBalance(connection, walletAddress, tokenB.mint)
	const existingValueB = balanceB * tokenB.priceUsd
	const needSwapB = Math.max(0, tokenB.valueUsd - existingValueB)

	console.log(
		`\n--- Token B (${tokenB.mint.slice(0, 8)}...): need $${tokenB.valueUsd}, have $${existingValueB.toFixed(4)}, swap $${needSwapB.toFixed(4)} ---`
	)

	const resultB = await swapIfNeeded(
		connection,
		keypair,
		tokenB.mint,
		needSwapB,
		USDC_MINT
	)

	if (!resultB.success) {
		console.log('USDC swap failed for Token B, trying USDT...')
		const resultB_USDT = await swapIfNeeded(
			connection,
			keypair,
			tokenB.mint,
			needSwapB,
			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 }
}