#!/usr/bin/env node
/**
 * Adjust Pool Prices Script
 * Rebalances an existing LiqP pool's prices to match market prices.
 * Uses IPartyInfo.swapToLimitAmounts() to compute the required swap size,
 * then executes via pool.swap() (swapToLimit has a known implementation bug).
 * Runs two passes to compensate for LMSR price coupling across assets.
 * Optionally basket-stakes a configurable USD amount via mint().
 */

import { ethers } from 'ethers';
import { readFile } from 'fs/promises';
import { config } from 'dotenv';
import { execSync } from 'child_process';

config({ path: new URL('../.env', import.meta.url).pathname });

// ============================================================================
// CONFIGURATION
// ============================================================================

const NETWORK_CONFIG = {
  mockchain: {
    rpcUrl: 'http://localhost:8545',
    chainId: '31337',
    tokens: {
      USDT: { address: '0xbdEd0D2bf404bdcBa897a74E6657f1f12e5C6fb6', coingeckoId: 'tether', decimals: 6 },
      USDC: { address: '0x2910E325cf29dd912E3476B61ef12F49cb931096', coingeckoId: 'usd-coin', decimals: 6 }
    }
  },
  mainnet: {
    rpcUrl: process.env.MAINNET_RPC_URL,
    chainId: '1',
    tokens: {
      USDT: { address: '0xdAC17F958D2ee523a2206206994597C13D831ec7', coingeckoId: 'tether', decimals: 6 },
      USDC: { address: '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', coingeckoId: 'usd-coin', decimals: 6 },
      WBTC: { address: '0x2260FAC5E5542a773Aa44fBCfeDf7C193bc2C599', coingeckoId: 'wrapped-bitcoin', decimals: 8 },
      WETH: { address: '0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2', coingeckoId: 'weth', decimals: 18 },
      UNI:  { address: '0x1f9840a85d5aF5bf1D1762F925BDADdC4201F984', coingeckoId: 'uniswap', decimals: 18 },
      WSOL: { address: '0xD31a59c85aE9D8edEFeC411D448f90841571b89c', coingeckoId: 'solana', decimals: 9 },
      TRX:  { address: '0x50327c6c5a14DCaDE707ABad2E27eB517df87AB5', coingeckoId: 'tron', decimals: 6 },
      AAVE: { address: '0x7Fc66500c84A76Ad7e9c93437bFc5Ac33E2DDaE9', coingeckoId: 'aave', decimals: 18 },
      PEPE: { address: '0x6982508145454Ce325dDbE47a25d4ec3d2311933', coingeckoId: 'pepe', decimals: 18 },
      SHIB: { address: '0x95aD61b0a150d79219dCF64E1E6Cc01f0B64C4cE', coingeckoId: 'shiba-inu', decimals: 18 }
    },
    uniswap: {
      swapRouter02: '0x68b3465833fb72A70ecDF485E0e4C7bD8665Fc45',
      factory: '0x1F98431c8aD98523631AE4a59f267346ea31F984',
      // Only 0.05% (500) and 0.3% (3000) tiers are used — 1% pools are avoided.
      defaultFeeTier: 3000,
      feeTierOverrides: { USDT: 500, USDC: 500, WBTC: 3000, WETH: 500 }
    }
  }
};

const NETWORK = process.env.NETWORK || 'mainnet';
const currentConfig = NETWORK_CONFIG[NETWORK];
if (!currentConfig) {
  console.error(`[!] Invalid NETWORK: ${NETWORK}. Must be 'mockchain' or 'mainnet'`);
  process.exit(1);
}

const TOKEN_CONFIG = currentConfig.tokens;

// Price deviation threshold: only swap if price is off by more than this
const PRICE_DEVIATION_THRESHOLD = 0.001; // 0.1%

// ============================================================================
// ABIs (minimal subset needed)
// ============================================================================

const ERC20ABI = [
  { type: 'function', name: 'balanceOf', stateMutability: 'view', inputs: [{ name: 'account', type: 'address' }], outputs: [{ name: '', type: 'uint256' }] },
  { type: 'function', name: 'decimals', stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'uint8' }] },
  { type: 'function', name: 'symbol', stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'string' }] },
  { type: 'function', name: 'approve', stateMutability: 'nonpayable', inputs: [{ name: 'spender', type: 'address' }, { name: 'amount', type: 'uint256' }], outputs: [{ name: '', type: 'bool' }] },
  { type: 'function', name: 'allowance', stateMutability: 'view', inputs: [{ name: 'owner', type: 'address' }, { name: 'spender', type: 'address' }], outputs: [{ name: '', type: 'uint256' }] }
];

const IPartyPoolABI = [
  { type: 'function', name: 'allTokens', stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'address[]' }] },
  { type: 'function', name: 'totalSupply', stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'uint256' }] },
  { type: 'function', name: 'balances', stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'uint256[]' }] },
  { type: 'function', name: 'balanceOf', stateMutability: 'view', inputs: [{ name: 'account', type: 'address' }], outputs: [{ name: '', type: 'uint256' }] },
  { type: 'function', name: 'denominators', stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'uint256[]' }] }
];

const IPartyInfoABI = [
  {
    type: 'function', name: 'price', stateMutability: 'view',
    inputs: [{ name: 'pool', type: 'address' }, { name: 'baseTokenIndex', type: 'uint256' }, { name: 'quoteTokenIndex', type: 'uint256' }],
    outputs: [{ name: '', type: 'uint256' }]
  },
  {
    type: 'function', name: 'swapToLimitAmounts', stateMutability: 'view',
    inputs: [{ name: 'pool', type: 'address' }, { name: 'inputTokenIndex', type: 'uint256' }, { name: 'outputTokenIndex', type: 'uint256' }, { name: 'limitPrice', type: 'int128' }],
    outputs: [{ name: 'amountIn', type: 'uint256' }, { name: 'amountOut', type: 'uint256' }, { name: 'inFee', type: 'uint256' }]
  },
  {
    type: 'function', name: 'mintAmounts', stateMutability: 'view',
    inputs: [{ name: 'pool', type: 'address' }, { name: 'lpTokenAmount', type: 'uint256' }],
    outputs: [{ name: 'depositAmounts', type: 'uint256[]' }]
  }
];

const UniswapV3FactoryABI = [
  { type: 'function', name: 'getPool', stateMutability: 'view',
    inputs: [{ name: 'tokenA', type: 'address' }, { name: 'tokenB', type: 'address' }, { name: 'fee', type: 'uint24' }],
    outputs: [{ name: 'pool', type: 'address' }] }
];

const UniswapV3PoolABI = [
  { type: 'function', name: 'liquidity', stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'uint128' }] }
];

const UniswapSwapRouter02ABI = [
  {
    type: 'function', name: 'exactOutputSingle', stateMutability: 'payable',
    inputs: [{
      name: 'params', type: 'tuple',
      components: [
        { name: 'tokenIn', type: 'address' },
        { name: 'tokenOut', type: 'address' },
        { name: 'fee', type: 'uint24' },
        { name: 'recipient', type: 'address' },
        { name: 'amountOut', type: 'uint256' },
        { name: 'amountInMaximum', type: 'uint256' },
        { name: 'sqrtPriceLimitX96', type: 'uint160' }
      ]
    }],
    outputs: [{ name: 'amountIn', type: 'uint256' }]
  }
];

// ============================================================================
// SIGNER ABSTRACTION
// ============================================================================

/**
 * Build a signer from parsed CLI options.
 *
 * mode 'privateKey': all writes go through ethers.Wallet (no cast required).
 * mode 'trezor':     all writes go through `cast send --trezor` (no private key stored).
 */
function buildSigner(opts, provider) {
  if (opts.privateKey) {
    const wallet = new ethers.Wallet(opts.privateKey, provider);
    return { mode: 'privateKey', address: wallet.address, wallet, rpcUrl: currentConfig.rpcUrl };
  } else {
    return { mode: 'trezor', address: opts.trezorAddress, wallet: null, rpcUrl: currentConfig.rpcUrl };
  }
}

/** Execute a cast send command (Trezor mode). Throws on failure. */
function runCast(signer, to, sigStr, argStr) {
  const nonce = signer.nextNonce();
  const cmd = [
    `cast send ${to}`,
    `  "${sigStr}"`,
    `  ${argStr}`,
    `  --rpc-url '${signer.rpcUrl}'`,
    `  --from ${signer.address}`,
    `  --nonce ${nonce}`,
    `  --trezor --mnemonic-index 0`
  ].join(' \\\n');
  console.log(`    [cast] ${cmd.replace(/\s+/g, ' ')}`);
  const output = execSync(cmd, { encoding: 'utf-8' });
  if (output) console.log(output.trim());
  return output;
}

// ============================================================================
// HELPERS
// ============================================================================

const Q64 = 2n ** 64n;

/**
 * Convert a floating-point value to Q64.64 BigInt.
 * Uses 1e9-scaled integer arithmetic to avoid float precision loss.
 */
function floatToQ64(x) {
  return BigInt(Math.round(x * 1e9)) * Q64 / 1_000_000_000n;
}

/**
 * Compute the limitPrice argument for swapToLimitAmounts.
 *
 * Derived from swapAmountsForPriceLimit: setting r0_new = r0_target in
 *   r0_new = r0² / (L×(1+r0) - r0²)  and solving for L gives:
 *
 *   L = r0i² × (1 + r0t) / (r0t × (1 + r0i))
 *
 * Requires r0tQ64 < r0iQ64 (target price lower than current in swap direction).
 * Both arguments are Q64.64 BigInts (raw internal LMSR price ratios, no denom scaling).
 */
function computeLimitPrice(r0iQ64, r0tQ64) {
  return r0iQ64 * r0iQ64 * (r0tQ64 + Q64) / (r0tQ64 * (r0iQ64 + Q64));
}

async function fetchCoinGeckoPrices(symbols) {
  const relevantTokens = Object.entries(TOKEN_CONFIG).filter(([sym]) => symbols.includes(sym));
  const ids = relevantTokens.map(([, t]) => t.coingeckoId).join(',');
  const url = `https://api.coingecko.com/api/v3/simple/price?ids=${ids}&vs_currencies=usd`;

  console.log(`[~] Fetching prices from CoinGecko...`);
  const response = await fetch(url);
  if (!response.ok) throw new Error(`CoinGecko API error: ${response.statusText}`);
  const data = await response.json();

  const prices = {};
  for (const [symbol, tokenInfo] of relevantTokens) {
    prices[symbol] = data[tokenInfo.coingeckoId]?.usd;
    if (!prices[symbol]) throw new Error(`No price returned for ${symbol} (${tokenInfo.coingeckoId})`);
  }
  return prices;
}

async function diagnoseGasError(err, signer, provider) {
  if (err.code !== 'UNPREDICTABLE_GAS_LIMIT') return;
  const [balance, feeData] = await Promise.all([
    provider.getBalance(signer.address),
    provider.getFeeData(),
  ]);
  const maxFee = feeData.maxFeePerGas || feeData.gasPrice || ethers.BigNumber.from(0);
  if (maxFee.isZero()) return;
  const affordableGas = balance.div(maxFee).toNumber();
  if (affordableGas < 100_000) {
    const needed = maxFee.mul(100_000).sub(balance);
    throw new Error(
      `Out of ETH for gas: wallet has ${ethers.utils.formatEther(balance)} ETH ` +
      `(~${affordableGas.toLocaleString()} gas units at ${ethers.utils.formatUnits(maxFee, 'gwei')} gwei maxFee). ` +
      `Top up at least ${ethers.utils.formatEther(needed)} ETH.`
    );
  }
}

async function approveToken(signer, provider, tokenAddress, tokenSymbol, spenderAddress, amount) {
  const decimals = TOKEN_CONFIG[tokenSymbol]?.decimals ?? 18;
  const tokenReadOnly = new ethers.Contract(tokenAddress, ERC20ABI, provider);
  const currentAllowance = await tokenReadOnly.allowance(signer.address, spenderAddress);

  if (currentAllowance.gte(amount)) {
    console.log(`    [=] ${tokenSymbol} allowance already sufficient (${ethers.utils.formatUnits(currentAllowance, decimals)})`);
    return;
  }

  if (signer.mode === 'privateKey') {
    try {
      const token = new ethers.Contract(tokenAddress, ERC20ABI, signer.wallet);
      if (tokenSymbol === 'USDT' && currentAllowance.gt(0)) {
        const tx = await token.approve(spenderAddress, 0, { nonce: signer.nextNonce() });
        await tx.wait();
        console.log(`    [+] ${tokenSymbol} allowance reset to 0`);
      }
      const tx = await token.approve(spenderAddress, amount, { nonce: signer.nextNonce() });
      await tx.wait();
      console.log(`    [+] ${tokenSymbol} approved (tx: ${tx.hash})`);
    } catch (err) {
      await diagnoseGasError(err, signer, provider);
      throw err;
    }
  } else {
    if (tokenSymbol === 'USDT' && currentAllowance.gt(0)) {
      runCast(signer, tokenAddress, 'approve(address,uint256)', `${spenderAddress} 0`);
      console.log(`    [+] ${tokenSymbol} allowance reset to 0`);
    }
    runCast(signer, tokenAddress, 'approve(address,uint256)', `${spenderAddress} ${amount.toString()}`);
    console.log(`    [+] ${tokenSymbol} approved`);
  }
}

async function checkPoolLiquidity(provider, tokenA, tokenB, feeTier) {
  const uniswapConfig = currentConfig.uniswap;
  const factory = new ethers.Contract(uniswapConfig.factory, UniswapV3FactoryABI, provider);
  const poolAddress = await factory.getPool(tokenA, tokenB, feeTier);
  if (poolAddress === ethers.constants.AddressZero) {
    throw new Error(
      `No Uniswap V3 pool exists for this pair at fee tier ${feeTier / 10000}% — ` +
      `try a different fee tier or route through an intermediate token`
    );
  }
  const pool = new ethers.Contract(poolAddress, UniswapV3PoolABI, provider);
  const liquidity = await pool.liquidity();
  if (liquidity.isZero()) {
    throw new Error(
      `Uniswap V3 pool ${poolAddress} has no active in-range liquidity at fee tier ${feeTier / 10000}% — ` +
      `the pool exists but all liquidity is out of range`
    );
  }
}

async function buyFromUniswap(signer, provider, tokenOut, tokenOutSymbol, amountOut, tokenIn, tokenInSymbol, maxAmountIn, dryRun) {
  const uniswapConfig = currentConfig.uniswap;
  if (!uniswapConfig) {
    throw new Error(`Uniswap not configured for network '${NETWORK}' — cannot buy ${tokenOutSymbol}`);
  }

  const feeTier    = uniswapConfig.feeTierOverrides[tokenOutSymbol] ?? uniswapConfig.defaultFeeTier;
  const outDecimals = TOKEN_CONFIG[tokenOutSymbol]?.decimals ?? 18;
  const inDecimals  = TOKEN_CONFIG[tokenInSymbol]?.decimals ?? 6;

  console.log(`    [~] Buying ${ethers.utils.formatUnits(amountOut, outDecimals)} ${tokenOutSymbol} from Uniswap V3 (fee tier: ${feeTier / 10000}%)`);
  console.log(`    [~] Max spending: ${ethers.utils.formatUnits(maxAmountIn, inDecimals)} ${tokenInSymbol}`);

  await checkPoolLiquidity(provider, tokenIn, tokenOut, feeTier);

  if (dryRun) {
    console.log(`    [DRY-RUN] Would buy ${tokenOutSymbol} via Uniswap V3`);
    return;
  }

  await approveToken(signer, provider, tokenIn, tokenInSymbol, uniswapConfig.swapRouter02, maxAmountIn);

  if (signer.mode === 'privateKey') {
    const router = new ethers.Contract(uniswapConfig.swapRouter02, UniswapSwapRouter02ABI, signer.wallet);
    let tx;
    try {
      tx = await router.exactOutputSingle({
        tokenIn,
        tokenOut,
        fee: feeTier,
        recipient: signer.address,
        amountOut,
        amountInMaximum: maxAmountIn,
        sqrtPriceLimitX96: 0
      }, { nonce: signer.nextNonce() });
    } catch (err) {
      await diagnoseGasError(err, signer, provider);
      throw err;
    }
    await tx.wait();
    console.log(`    [+] Bought ${tokenOutSymbol} via Uniswap (tx: ${tx.hash})`);
  } else {
    // cast send with tuple arg: "(tokenIn,tokenOut,fee,recipient,amountOut,amountInMaximum,sqrtLimit)"
    const tupleArg = `"(${tokenIn},${tokenOut},${feeTier},${signer.address},${amountOut.toString()},${maxAmountIn.toString()},0)"`;
    runCast(
      signer,
      uniswapConfig.swapRouter02,
      'exactOutputSingle((address,address,uint24,address,uint256,uint256,uint160))',
      tupleArg
    );
    console.log(`    [+] Bought ${tokenOutSymbol} via Uniswap`);
  }
}

// ============================================================================
// REBALANCING LOGIC
// ============================================================================

/** Query the current pool price deviation for a single token vs CoinGecko. */
async function getTokenDeviation(partyInfo, poolAddr, token, quoteToken, usdPrices) {
  const currentQ128BN = await partyInfo.price(poolAddr, token.idx, quoteToken.idx);
  const currentQ128 = BigInt(currentQ128BN.toString());
  // partyInfo.price returns r0 * 10^(quoteDec-baseDec) in Q128 format.
  const currentQ64n = currentQ128 / Q64;
  const poolPriceRaw = Number(currentQ64n) / Number(Q64);
  const poolPriceHuman = poolPriceRaw * Math.pow(10, token.decimals - quoteToken.decimals);
  const poolUSD = poolPriceHuman * usdPrices[quoteToken.symbol];
  const targetPriceHuman = usdPrices[token.symbol] / usdPrices[quoteToken.symbol];
  const deviation = targetPriceHuman === 0 ? 0
    : (poolPriceHuman - targetPriceHuman) / targetPriceHuman * 100;
  return { token, deviation, deviationAbs: Math.abs(deviation), poolUSD };
}

/**
 * Execute a rebalancing swap for a single token via pool.swap().
 * Returns true if a swap was executed, false if price is already within threshold.
 */
async function rebalanceTokenIfNeeded(
  { signer, provider, partyInfo, poolAddr, quoteToken, usdPrices, dryRun, denominators },
  token
) {
  const currentQ128BN = await partyInfo.price(poolAddr, token.idx, quoteToken.idx);
  const currentQ128 = BigInt(currentQ128BN.toString());
  const currentQ64n = currentQ128 / Q64;
  const poolPriceRaw = Number(currentQ64n) / Number(Q64);
  const poolPriceHuman = poolPriceRaw * Math.pow(10, token.decimals - quoteToken.decimals);

  const targetPriceHuman = usdPrices[token.symbol] / usdPrices[quoteToken.symbol];
  const deviation = targetPriceHuman === 0 ? 0
    : (poolPriceHuman - targetPriceHuman) / targetPriceHuman * 100;

  if (Math.abs(deviation) < PRICE_DEVIATION_THRESHOLD * 100) {
    return false;
  }

  // LMSR behaves like a regular AMM: injecting a token makes it cheaper.
  // Overpriced base: inject TOKEN, extract QUOTE → lowers TOKEN price.
  // Underpriced base: inject QUOTE, extract TOKEN → raises TOKEN price.
  const overpriced  = poolPriceHuman > targetPriceHuman;
  const inputToken  = overpriced ? token      : quoteToken;
  const outputToken = overpriced ? quoteToken : token;
  const inputIdx    = overpriced ? token.idx  : quoteToken.idx;
  const outputIdx   = overpriced ? quoteToken.idx : token.idx;
  const direction   = overpriced ? 'lower-price' : 'raise-price';

  console.log(`    → ${direction}: inject ${inputToken.symbol}, extract ${outputToken.symbol}`);

  // Compute raw internal Q64.64 r0 for the swap direction using actual pool denominators.
  // partyInfo.price(i,j) = r0_internal(i,j) × D[j] / D[i]  in Q128.128
  // → r0_internal Q64.64 = price_Q128 × D[i] / D[j] / Q64
  const r0PriceQ128 = BigInt((await partyInfo.price(poolAddr, inputIdx, outputIdx)).toString());
  const r0InitialQ64 = r0PriceQ128 * denominators[inputIdx] / denominators[outputIdx] / Q64;

  // Scale r0_initial to get r0_target using the human-price ratio in the swap direction.
  // overpriced (TOKEN→QUOTE): target price is lower → ratio = target/current < 1
  // underpriced (QUOTE→TOKEN): we want lower r0(QUOTE,TOKEN) → ratio = current/target < 1
  const numPrice = overpriced ? targetPriceHuman : poolPriceHuman;
  const denPrice = overpriced ? poolPriceHuman   : targetPriceHuman;
  const r0TargetQ64 = r0InitialQ64 * BigInt(Math.round(numPrice * 1e9)) / BigInt(Math.round(denPrice * 1e9));

  const limitPriceQ64 = computeLimitPrice(r0InitialQ64, r0TargetQ64);
  const r0iFloat = Number(r0InitialQ64) / Number(Q64);
  const r0tFloat = Number(r0TargetQ64) / Number(Q64);
  console.log(`    r0Initial=${r0iFloat.toExponential(4)}  r0Target=${r0tFloat.toExponential(4)}  limitPrice=${limitPriceQ64.toString()}`);

  const REDUCE_SIZE_RETRIES = 5;
  let amountIn, amountOut, inFee, effectiveLimitPrice;
  let r0TargetRetry = r0TargetQ64;
  for (let attempt = 0; ; attempt++) {
    const retryLimitPrice = attempt === 0 ? limitPriceQ64 : computeLimitPrice(r0InitialQ64, r0TargetRetry);
    try {
      const result = await partyInfo.swapToLimitAmounts(poolAddr, inputIdx, outputIdx, retryLimitPrice.toString());
      amountIn  = result.amountIn;
      amountOut = result.amountOut;
      inFee     = result.inFee;
      effectiveLimitPrice = retryLimitPrice;
      break;
    } catch (err) {
      if (
        attempt < REDUCE_SIZE_RETRIES &&
        err.message.includes('arithmetic underflow or overflow')
      ) {
        r0TargetRetry = r0InitialQ64 - (r0InitialQ64 - r0TargetRetry) / 10n;
        console.log(`    [!] swapToLimitAmounts overflow — retrying with reduced target (attempt ${attempt + 2}/${REDUCE_SIZE_RETRIES + 1})`);
        continue;
      }
      console.log(`    [!] swapToLimitAmounts reverted: ${err.message} — skipping`);
      return false;
    }
  }

  if (amountIn.isZero()) {
    console.log(`    [+] amountIn=0, nothing to do`);
    return false;
  }

  console.log(`    amountIn:  ${ethers.utils.formatUnits(amountIn, inputToken.decimals)} ${inputToken.symbol}`);
  console.log(`    amountOut: ${ethers.utils.formatUnits(amountOut, outputToken.decimals)} ${outputToken.symbol}`);
  console.log(`    fee:       ${ethers.utils.formatUnits(inFee, inputToken.decimals)} ${inputToken.symbol}`);

  // Ensure the signer holds enough of the input token; buy shortfall from Uniswap if needed
  const inputERC20   = new ethers.Contract(inputToken.address, ERC20ABI, provider);
  const inputBalance = await inputERC20.balanceOf(signer.address);

  if (inputBalance.lt(amountIn)) {
    const deficit = amountIn.sub(inputBalance);
    console.log(`    [!] Insufficient ${inputToken.symbol}: ` +
      `have ${ethers.utils.formatUnits(inputBalance, inputToken.decimals)}, ` +
      `need ${ethers.utils.formatUnits(amountIn, inputToken.decimals)}, ` +
      `deficit ${ethers.utils.formatUnits(deficit, inputToken.decimals)}`);

    if (inputToken.symbol === quoteToken.symbol) {
      console.log(`    [!] Input is the quote token — no way to acquire more, skipping`);
      return false;
    }

    const deficitUSD = Number(ethers.utils.formatUnits(deficit, inputToken.decimals)) * usdPrices[inputToken.symbol];
    const maxQuoteIn = ethers.utils.parseUnits(
      (deficitUSD * 1.02 / usdPrices[quoteToken.symbol]).toFixed(quoteToken.decimals),
      quoteToken.decimals
    );

    await buyFromUniswap(
      signer, provider,
      inputToken.address, inputToken.symbol, deficit,
      quoteToken.address, quoteToken.symbol, maxQuoteIn,
      dryRun
    );
  }

  console.log(`    [~] Approving ${inputToken.symbol} for pool...`);
  if (!dryRun) {
    await approveToken(signer, provider, inputToken.address, inputToken.symbol, poolAddr, amountIn.mul(101).div(100));
  } else {
    console.log(`    [DRY-RUN] Would approve ${inputToken.symbol}`);
  }

  // Execute via pool.swap() with the computed amountIn and the effective limit price.
  const deadline = Math.floor(Date.now() / 1000) + 300;

  if (!dryRun) {
    if (signer.mode === 'privateKey') {
      const poolContract = new ethers.Contract(poolAddr, [
        { type: 'function', name: 'swap', stateMutability: 'payable',
          inputs: [
            { name: 'payer', type: 'address' }, { name: 'fundingSelector', type: 'bytes4' },
            { name: 'receiver', type: 'address' }, { name: 'inputTokenIndex', type: 'uint256' },
            { name: 'outputTokenIndex', type: 'uint256' }, { name: 'maxAmountIn', type: 'uint256' },
            { name: 'limitPrice', type: 'int128' }, { name: 'deadline', type: 'uint256' },
            { name: 'unwrap', type: 'bool' }, { name: 'cbData', type: 'bytes' }
          ],
          outputs: [{ name: 'amountIn', type: 'uint256' }, { name: 'amountOut', type: 'uint256' }, { name: 'inFee', type: 'uint256' }]
        }
      ], signer.wallet);
      let tx;
      let swapAmountIn = amountIn;
      for (let attempt = 0; ; attempt++) {
        const nonce = signer.nextNonce();
        try {
          tx = await poolContract.swap(
            signer.address, '0x00000000', signer.address,
            inputIdx, outputIdx, swapAmountIn, effectiveLimitPrice.toString(), deadline, false, '0x',
            { nonce }
          );
          break;
        } catch (err) {
          if (
            attempt < REDUCE_SIZE_RETRIES &&
            err.code === 'UNPREDICTABLE_GAS_LIMIT' &&
            err.message.includes('arithmetic underflow or overflow')
          ) {
            signer._nonce--; // tx never submitted; reuse nonce
            swapAmountIn = swapAmountIn.div(10);
            console.log(`    [!] Arithmetic overflow — retrying with ${ethers.utils.formatUnits(swapAmountIn, inputToken.decimals)} ${inputToken.symbol} (attempt ${attempt + 2}/${REDUCE_SIZE_RETRIES})`);
            continue;
          }
          await diagnoseGasError(err, signer, provider);
          throw err;
        }
      }
      await tx.wait();
      console.log(`    [+] Swap executed (tx: ${tx.hash})`);
    } else {
      runCast(
        signer, poolAddr,
        'swap(address,bytes4,address,uint256,uint256,uint256,int128,uint256,bool,bytes)',
        `${signer.address} "0x00000000" ${signer.address} ${inputIdx} ${outputIdx} ${amountIn.toString()} ${effectiveLimitPrice.toString()} ${deadline} false "0x"`
      );
      console.log(`    [+] Swap executed`);
    }

    const newQ128BN  = await partyInfo.price(poolAddr, token.idx, quoteToken.idx);
    const newQ64n    = BigInt(newQ128BN.toString()) / Q64;
    const newPoolUSD = Number(newQ64n) / Number(Q64)
      * Math.pow(10, token.decimals - quoteToken.decimals)
      * usdPrices[quoteToken.symbol];
    const residual = (newPoolUSD - usdPrices[token.symbol]) / usdPrices[token.symbol] * 100;
    console.log(`    [+] Updated pool price: $${newPoolUSD.toPrecision(6)} (market: $${usdPrices[token.symbol].toPrecision(6)}  residual: ${residual.toFixed(3)}%)`);
  } else {
    console.log(`    [DRY-RUN] Would swap ${inputToken.symbol} → ${outputToken.symbol} via pool.swap()`);
  }

  return true;
}

// ============================================================================
// CLI ARGUMENT PARSING
// ============================================================================

function parseArgs() {
  const args = process.argv.slice(2);

  if (args.includes('--help') || args.includes('-h')) {
    console.log(`
Usage: node adjust_pool_prices_and_stake.js --pool <address> (--private-key <key> | --trezor <address>) [OPTIONS]

Signer (exactly one required):
  --private-key <key>   Sign all transactions with this private key
  --trezor <address>    Sign all transactions via Trezor hardware wallet at given address

Options:
  --pool <address>      Target LiqP pool address (required)
  --quote <symbol>      Quote token symbol (default: auto-detect USDT/USDC)
  --iterations <n>      Max rebalancing iterations, best-first (default: 10)
  --stake <usd>         USD amount to basket-stake via mint() after rebalancing
  --skip-rebalance      Skip price adjustment; only run staking (requires --stake)
  --dry-run             Preview only, no transactions executed
  --help, -h            Show this help message

Environment:
  NETWORK               'mainnet' (default) or 'mockchain'
  MAINNET_RPC_URL       RPC endpoint (required for mainnet)

Examples:
  node adjust_pool_prices_and_stake.js --pool 0x1234... --private-key 0xabc... --dry-run
  node adjust_pool_prices_and_stake.js --pool 0x1234... --trezor 0xdef... --stake 500
  node adjust_pool_prices_and_stake.js --pool 0x1234... --trezor 0xdef... --stake 500 --skip-rebalance
`);
    process.exit(0);
  }

  const opts = {
    pool: null, quote: null, stake: null, iterations: 10,
    privateKey: null, trezorAddress: null,
    skipRebalance: false, dryRun: false
  };

  for (let i = 0; i < args.length; i++) {
    if      (args[i] === '--pool'        && i + 1 < args.length) { opts.pool         = args[++i]; }
    else if (args[i] === '--quote'       && i + 1 < args.length) { opts.quote        = args[++i].toUpperCase(); }
    else if (args[i] === '--iterations'  && i + 1 < args.length) { opts.iterations   = parseInt(args[++i], 10); }
    else if (args[i] === '--stake'       && i + 1 < args.length) { opts.stake        = parseFloat(args[++i]); }
    else if (args[i] === '--private-key' && i + 1 < args.length) { opts.privateKey   = args[++i]; }
    else if (args[i] === '--trezor'      && i + 1 < args.length) { opts.trezorAddress = args[++i]; }
    else if (args[i] === '--skip-rebalance') { opts.skipRebalance = true; }
    else if (args[i] === '--dry-run')        { opts.dryRun = true; }
  }

  if (!opts.pool) {
    console.error('[!] --pool <address> is required.');
    process.exit(1);
  }
  if (!opts.privateKey && !opts.trezorAddress) {
    console.error('[!] Signer required: provide --private-key <key> or --trezor <address>.');
    process.exit(1);
  }
  if (opts.privateKey && opts.trezorAddress) {
    console.error('[!] Provide --private-key OR --trezor, not both.');
    process.exit(1);
  }
  if (opts.skipRebalance && !opts.stake) {
    console.error('[!] --skip-rebalance requires --stake <usd>.');
    process.exit(1);
  }

  return opts;
}

// ============================================================================
// MAIN
// ============================================================================

async function main() {
  const opts = parseArgs();
  const { pool: poolAddr, dryRun } = opts;

  const signerLabel = opts.trezorAddress ? `trezor:${opts.trezorAddress}` : 'private-key';
  console.log(`${'='.repeat(70)}`);
  console.log(`Adjust Pool Prices`);
  console.log(`Network: ${NETWORK}  Pool: ${poolAddr}  Signer: ${signerLabel}${opts.skipRebalance ? '  [STAKE-ONLY]' : ''}${dryRun ? '  [DRY-RUN]' : ''}`);
  console.log(`${'='.repeat(70)}\n`);

  if (!currentConfig.rpcUrl) {
    console.error(`[!] Missing RPC URL for network '${NETWORK}' (set MAINNET_RPC_URL)`);
    process.exit(1);
  }

  // ── Phase 0: Setup ─────────────────────────────────────────────────────────

  const provider = new ethers.providers.JsonRpcProvider(currentConfig.rpcUrl);
  const signer   = buildSigner(opts, provider);
  console.log(`[+] Signer:    ${signer.address} (${signer.mode})`);

  if (!dryRun) {
    signer._nonce = await provider.getTransactionCount(signer.address, "pending");
    console.log(`[+] Nonce:     ${signer._nonce} (pending)`);
  } else {
    signer._nonce = 0;
  }
  signer.nextNonce = () => signer._nonce++;

  const chainInfoData = JSON.parse(
    await readFile(new URL('../src/contracts/liqp-deployments.json', import.meta.url), 'utf-8')
  );
  const chainContracts = chainInfoData[currentConfig.chainId]?.v1;
  if (!chainContracts) {
    console.error(`[!] No contract addresses for chainId ${currentConfig.chainId}`);
    process.exit(1);
  }
  const PARTY_INFO_ADDRESS = chainContracts.PartyInfo;
  console.log(`[+] PartyInfo: ${PARTY_INFO_ADDRESS}`);

  const pool      = new ethers.Contract(poolAddr, IPartyPoolABI, provider);
  const partyInfo = new ethers.Contract(PARTY_INFO_ADDRESS, IPartyInfoABI, provider);

  const denomsBN = await pool.denominators();
  const denominators = denomsBN.map(d => BigInt(d.toString()));
  console.log(`[+] Denominators: [${denominators.join(', ')}]`);

  const tokenAddresses = await pool.allTokens();
  console.log(`[+] Pool tokens: ${tokenAddresses.length}`);

  const tokenConfigByAddr = {};
  for (const [sym, info] of Object.entries(TOKEN_CONFIG)) {
    tokenConfigByAddr[info.address.toLowerCase()] = { symbol: sym, ...info };
  }

  const poolTokens = tokenAddresses.map((addr, idx) => {
    const meta = tokenConfigByAddr[addr.toLowerCase()];
    if (!meta) throw new Error(`Unknown token at pool index ${idx}: ${addr}`);
    return { idx, address: addr, symbol: meta.symbol, decimals: meta.decimals, coingeckoId: meta.coingeckoId };
  });

  let quoteToken;
  if (opts.quote) {
    quoteToken = poolTokens.find(t => t.symbol === opts.quote);
    if (!quoteToken) throw new Error(`Quote token ${opts.quote} not found in pool`);
  } else {
    quoteToken = poolTokens.find(t => t.symbol === 'USDT') ?? poolTokens.find(t => t.symbol === 'USDC');
    if (!quoteToken) throw new Error('No USDT or USDC in pool; specify --quote <symbol>');
  }
  console.log(`[+] Quote:     ${quoteToken.symbol} (index ${quoteToken.idx})\n`);

  const nonQuoteTokens = poolTokens.filter(t => t.idx !== quoteToken.idx);

  // ── Phase 1: Fetch Prices ───────────────────────────────────────────────────

  const symbols   = poolTokens.map(t => t.symbol);
  const usdPrices = await fetchCoinGeckoPrices(symbols);

  console.log(`\n[+] Market prices:`);
  for (const sym of symbols) {
    console.log(`    ${sym.padEnd(6)}: $${usdPrices[sym].toLocaleString()}`);
  }

  const ctx = { signer, provider, partyInfo, poolAddr, poolTokens, quoteToken, usdPrices, dryRun, denominators };

  // ── Phase 2: Rebalancing — best-first iterative (skipped with --skip-rebalance) ─

  let swapCount = 0;

  if (opts.skipRebalance) {
    console.log(`\n[~] Skipping price adjustment (--skip-rebalance).`);
  } else {
    console.log(`\n${'─'.repeat(70)}`);
    console.log(`Rebalancing — best-first, up to ${opts.iterations} iterations (threshold: ${PRICE_DEVIATION_THRESHOLD * 100}%)`);
    console.log(`${'─'.repeat(70)}`);

    let converged = false;
    for (let iter = 1; iter <= opts.iterations; iter++) {
      // Query all pool prices in parallel, find the most deviant token
      const deviations = await Promise.all(
        nonQuoteTokens.map(t => getTokenDeviation(partyInfo, poolAddr, t, quoteToken, usdPrices))
      );

      console.log(`\n[iter ${iter}/${opts.iterations}]`);
      for (const d of deviations) {
        const marker = d.deviationAbs >= PRICE_DEVIATION_THRESHOLD * 100 ? '!' : '=';
        console.log(`  [${marker}] ${d.token.symbol.padEnd(6)} pool=$${d.poolUSD.toPrecision(6)}  market=$${usdPrices[d.token.symbol].toPrecision(6)}  deviation=${d.deviation.toFixed(3)}%`);
      }

      const worst = deviations.reduce((a, b) => a.deviationAbs > b.deviationAbs ? a : b);

      if (worst.deviationAbs < PRICE_DEVIATION_THRESHOLD * 100) {
        console.log(`\n[+] All prices within threshold — converged after ${swapCount} swap(s).`);
        converged = true;
        break;
      }

      console.log(`\n  Swapping ${worst.token.symbol} (${worst.deviation.toFixed(3)}% off)`);
      const swapped = await rebalanceTokenIfNeeded(ctx, worst.token);
      if (swapped) swapCount++;
    }

    if (!converged) {
      console.log(`\n[~] Reached iteration limit (${opts.iterations}) with ${swapCount} swap(s) executed.`);
    }
  }

  // ── Phase 4: Optional Basket Stake via mint() ──────────────────────────────

  if (opts.stake) {
    console.log(`\n${'='.repeat(70)}`);
    console.log(`Phase 4: Basket Stake ($${opts.stake} USD)`);
    console.log(`${'='.repeat(70)}`);

    const totalSupplyBN = await pool.totalSupply();
    const balancesBN    = await pool.balances();

    let tvlUSD = 0;
    for (let i = 0; i < poolTokens.length; i++) {
      tvlUSD += Number(ethers.utils.formatUnits(balancesBN[i], poolTokens[i].decimals)) * usdPrices[poolTokens[i].symbol];
    }

    const totalSupplyFloat = Number(ethers.utils.formatUnits(totalSupplyBN, 18));
    const lpPriceUSD       = tvlUSD / totalSupplyFloat;
    console.log(`\n[~] Pool TVL: $${tvlUSD.toFixed(2)}, LP price: $${lpPriceUSD.toFixed(6)}`);

    const lpAmountFloat = opts.stake / lpPriceUSD;
    const lpAmount      = ethers.utils.parseUnits(lpAmountFloat.toFixed(18), 18);
    console.log(`[~] Target LP amount: ${lpAmountFloat.toFixed(6)} LP tokens (~$${opts.stake})`);

    const depositAmounts = await partyInfo.mintAmounts(poolAddr, lpAmount);
    console.log(`\n[~] Required deposits for mint:`);

    const deficits = [];
    let totalQuoteNeeded = ethers.BigNumber.from(0);

    for (let i = 0; i < poolTokens.length; i++) {
      const t          = poolTokens[i];
      const required   = depositAmounts[i];
      const tokenERC20 = new ethers.Contract(t.address, ERC20ABI, provider);
      const balance    = await tokenERC20.balanceOf(signer.address);
      const sufficient = balance.gte(required);

      console.log(
        `    ${t.symbol.padEnd(6)}: ` +
        `${ethers.utils.formatUnits(required, t.decimals).padStart(22)} required  ` +
        `${ethers.utils.formatUnits(balance,  t.decimals).padStart(22)} in wallet  ` +
        `${sufficient ? '✓' : '✗'}`
      );

      if (!sufficient) {
        const deficit    = required.sub(balance);
        const deficitUSD = Number(ethers.utils.formatUnits(deficit, t.decimals)) * usdPrices[t.symbol];
        const quoteNeeded = ethers.utils.parseUnits(
          (deficitUSD * 1.02 / usdPrices[quoteToken.symbol]).toFixed(quoteToken.decimals),
          quoteToken.decimals
        );
        totalQuoteNeeded = totalQuoteNeeded.add(quoteNeeded);
        deficits.push({ token: t, deficit, quoteNeeded });
      }
    }

    // Pre-flight: fail before any transaction if quote balance is insufficient
    if (deficits.length > 0) {
      const quoteERC20   = new ethers.Contract(quoteToken.address, ERC20ABI, provider);
      const quoteBalance = await quoteERC20.balanceOf(signer.address);

      if (quoteBalance.lt(totalQuoteNeeded)) {
        console.error(`\n[!] Pre-flight check FAILED — insufficient ${quoteToken.symbol} to cover token deficits:`);
        for (const d of deficits) {
          console.error(`    ${d.token.symbol}: deficit ${ethers.utils.formatUnits(d.deficit, d.token.decimals)}` +
            ` → costs ~${ethers.utils.formatUnits(d.quoteNeeded, quoteToken.decimals)} ${quoteToken.symbol}`);
        }
        console.error(`    Total needed: ${ethers.utils.formatUnits(totalQuoteNeeded, quoteToken.decimals)} ${quoteToken.symbol}`);
        console.error(`    Wallet has:   ${ethers.utils.formatUnits(quoteBalance,     quoteToken.decimals)} ${quoteToken.symbol}`);
        process.exit(1);
      }
      console.log(`\n[+] Pre-flight check passed`);
    }

    if (dryRun) {
      console.log(`\n[DRY-RUN] Would buy ${deficits.length} deficit token(s) and mint ${lpAmountFloat.toFixed(6)} LP`);
    } else {
      for (const d of deficits) {
        console.log(`\n[~] Buying ${d.token.symbol} deficit from Uniswap...`);
        await buyFromUniswap(
          signer, provider,
          d.token.address, d.token.symbol, d.deficit,
          quoteToken.address, quoteToken.symbol, d.quoteNeeded,
          false
        );
      }

      console.log(`\n[~] Approving tokens for pool mint...`);
      for (let i = 0; i < poolTokens.length; i++) {
        const t        = poolTokens[i];
        const required = depositAmounts[i];
        if (required.gt(0)) {
          await approveToken(signer, provider, t.address, t.symbol, poolAddr, required.mul(101).div(100));
        }
      }

      const deadline = Math.floor(Date.now() / 1000) + 300;

      if (signer.mode === 'privateKey') {
        const poolContract = new ethers.Contract(poolAddr, [
          { type: 'function', name: 'mint', stateMutability: 'payable',
            inputs: [
              { name: 'payer', type: 'address' }, { name: 'receiver', type: 'address' },
              { name: 'lpTokenAmount', type: 'uint256' }, { name: 'deadline', type: 'uint256' }
            ],
            outputs: [{ name: 'lpMinted', type: 'uint256' }]
          }
        ], signer.wallet);
        let tx;
        try {
          tx = await poolContract.mint(signer.address, signer.address, lpAmount, deadline, { nonce: signer.nextNonce() });
        } catch (err) {
          await diagnoseGasError(err, signer, provider);
          throw err;
        }
        await tx.wait();
        console.log(`[+] Mint executed (tx: ${tx.hash})`);
      } else {
        runCast(
          signer, poolAddr,
          'mint(address,address,uint256,uint256)',
          `${signer.address} ${signer.address} ${lpAmount.toString()} ${deadline}`
        );
        console.log(`[+] Mint executed`);
      }

      const newLpBalance = await pool.balanceOf(signer.address);
      console.log(`[+] Wallet LP balance: ${ethers.utils.formatUnits(newLpBalance, 18)}`);
    }
  }

  // ── Summary ────────────────────────────────────────────────────────────────

  console.log(`\n${'='.repeat(70)}`);
  console.log(`Summary`);
  console.log(`${'='.repeat(70)}`);
  if (opts.skipRebalance) {
    console.log(`Rebalance:    skipped (--skip-rebalance)`);
  } else {
    console.log(`Rebalance:    ${swapCount} swap(s), up to ${opts.iterations} iterations${dryRun ? ' (dry-run)' : ''}`);
  }
  if (opts.stake) {
    console.log(`Stake:        $${opts.stake} USD${dryRun ? ' (dry-run)' : ''}`);
  }
  console.log(`Done.`);
}

main().catch(err => {
  console.error('[!] Fatal error:', err.message || err);
  process.exit(1);
});