import BN from "bn.js";
import { fromWadFloor, fromWadCeil, multiplyWad } from "./wad.js";
import { simulateTrade, tokenPrice, resolveLmsrIndices } from "./lmsr.js";
import { wadToFloat } from "./conversions.js";
import { PoolType } from "../types.js";
import { TRADE_FEE_PERCENTAGE, MAX_Q_VECTOR_DIFF_MULTIPLIER, WAD } from "../constants.js";

/**
 * Trade Simulation Module
 *
 * Frontend helpers for calculating token amounts and costs
 */

const WAD_BN = new BN(WAD.toString());

export interface TradeSimulationResult {
  grossCost: BN;
  fee: BN;
  netCost: BN;
  tokensAmount: BN;
  pricePerToken: number;
  priceImpact: number;
}

/**
 * Calculate price per token for display
 */
function calculatePricePerToken(costLamports: BN, tokenLamports: BN): number {
  if (tokenLamports.isZero()) return 0;
  const priceWad = costLamports.mul(WAD_BN).div(tokenLamports);
  return wadToFloat(priceWad);
}

/**
 * Calculate price impact percentage
 */
function calculatePriceImpact(oldPriceWad: BN, newPriceWad: BN): number {
  if (oldPriceWad.isZero()) return 0;
  const diff = newPriceWad.sub(oldPriceWad);
  const impactWad = diff.mul(WAD_BN).div(oldPriceWad);
  return wadToFloat(impactWad) * 100;
}

/**
 * Binary search for optimal token amount given SOL budget
 */
export function tokensForSolBudget(
  qVector: BN[],
  b: BN,
  poolType: PoolType,
  strikeIndex: number,
  isHit: boolean,
  solBudget: BN
): { tokens: BN; simulation: TradeSimulationResult } {
  const affectedIndices = resolveLmsrIndices(poolType, strikeIndex, isHit);
  const feeRate = new BN(TRADE_FEE_PERCENTAGE.toString());

  // Token amount upper bound based on on-chain q_vector constraint
  // On-chain constraint: max(q_i) - min(q_i) < MAX_Q_VECTOR_DIFF_MULTIPLIER × b
  // Note: This is TOKEN amount limit, not SOL limit
  const maxTokenAmount = b.muln(MAX_Q_VECTOR_DIFF_MULTIPLIER);

  let low = new BN(0);
  let high = maxTokenAmount; // Search up to max tradeable token amount
  let bestTokens = new BN(0);
  let bestSimulation: TradeSimulationResult | null = null;

  for (let i = 0; i < 50; i++) {
    const mid = low.add(high).div(new BN(2));

    if (mid.eq(low)) {
      if (mid.eq(bestTokens)) break;
      if (high.sub(low).lte(new BN(1))) break;
    }

    try {
      const { cost, newQVector } = simulateTrade(qVector, b, affectedIndices, mid, true);
      const fee = multiplyWad(cost, feeRate);
      const netCost = fromWadCeil(cost.add(fee));

      if (netCost.lte(solBudget)) {
        bestTokens = mid;
        const currentPrice = tokenPrice(qVector, b, poolType, strikeIndex, isHit);
        const newPrice = tokenPrice(newQVector, b, poolType, strikeIndex, isHit);
        bestSimulation = {
          grossCost: fromWadFloor(cost),
          fee: fromWadCeil(fee),
          netCost,
          tokensAmount: mid,
          pricePerToken: calculatePricePerToken(netCost, mid),
          priceImpact: calculatePriceImpact(currentPrice, newPrice),
        };
        low = mid;
      } else {
        high = mid;
      }
    } catch (err: unknown) {
      // Numerical overflow or other simulation error - reduce upper bound
      const errorMessage = err instanceof Error ? err.message : String(err);
      if (errorMessage.includes("expWad: input too large")) {
        high = mid;
      } else {
        // Log unexpected errors for debugging but continue binary search
        console.warn(
          "tokensForSolBudget: Unexpected simulation error at",
          mid.toString(),
          "lamports:",
          err
        );
        high = mid; // Try smaller amount
      }
    }
  }

  return {
    tokens: bestTokens,
    simulation: bestSimulation || {
      grossCost: new BN(0),
      fee: new BN(0),
      netCost: new BN(0),
      tokensAmount: new BN(0),
      pricePerToken: 0,
      priceImpact: 0,
    },
  };
}

/**
 * Calculate SOL cost for buying tokens
 */
export function solCostForTokens(
  qVector: BN[],
  b: BN,
  poolType: PoolType,
  strikeIndex: number,
  isHit: boolean,
  tokenAmount: BN
): TradeSimulationResult {
  const affectedIndices = resolveLmsrIndices(poolType, strikeIndex, isHit);
  const feeRate = new BN(TRADE_FEE_PERCENTAGE.toString());

  const { cost, newQVector } = simulateTrade(qVector, b, affectedIndices, tokenAmount, true);
  const fee = multiplyWad(cost, feeRate);
  const netCost = fromWadCeil(cost.add(fee));

  const currentPrice = tokenPrice(qVector, b, poolType, strikeIndex, isHit);
  const newPrice = tokenPrice(newQVector, b, poolType, strikeIndex, isHit);

  return {
    grossCost: fromWadFloor(cost),
    fee: fromWadCeil(fee),
    netCost,
    tokensAmount: tokenAmount,
    pricePerToken: calculatePricePerToken(netCost, tokenAmount),
    priceImpact: calculatePriceImpact(currentPrice, newPrice),
  };
}

/**
 * Calculate SOL refund for selling tokens
 */
export function solRefundForTokens(
  qVector: BN[],
  b: BN,
  poolType: PoolType,
  strikeIndex: number,
  isHit: boolean,
  tokenAmount: BN
): TradeSimulationResult {
  const affectedIndices = resolveLmsrIndices(poolType, strikeIndex, isHit);
  const feeRate = new BN(TRADE_FEE_PERCENTAGE.toString());

  try {
    const { cost, newQVector } = simulateTrade(qVector, b, affectedIndices, tokenAmount, false);
    const fee = multiplyWad(cost, feeRate);
    const netRefund = fromWadFloor(cost.sub(fee));

    const currentPrice = tokenPrice(qVector, b, poolType, strikeIndex, isHit);
    const newPrice = tokenPrice(newQVector, b, poolType, strikeIndex, isHit);

    return {
      grossCost: fromWadFloor(cost),
      fee: fromWadCeil(fee),
      netCost: netRefund,
      tokensAmount: tokenAmount,
      pricePerToken: calculatePricePerToken(netRefund, tokenAmount),
      priceImpact: calculatePriceImpact(currentPrice, newPrice),
    };
  } catch (err: unknown) {
    const errorMessage = err instanceof Error ? err.message : String(err);
    if (errorMessage.includes("expWad: input too large")) {
      return {
        grossCost: new BN(0),
        fee: new BN(0),
        netCost: new BN(0),
        tokensAmount: new BN(0),
        pricePerToken: 0,
        priceImpact: 0,
      };
    }
    throw err;
  }
}