DEX Shares Calculation Utils

import { borrow, deposit, payback, withdraw } from "./internal/utils";
import type { ISharesMinMaxParams } from "../../types";

export interface ColReserves {
  token0ImaginaryReserves: bigint;
  token0RealReserves: bigint;
  token1ImaginaryReserves: bigint;
  token1RealReserves: bigint;
}

export interface Pex {
  geometricMean: bigint;
  lowerRange: bigint;
  upperRange: bigint;
}

export interface DebtReserves extends ColReserves {
  token0Debt: bigint;
  token1Debt: bigint;
}

export interface SharesMinMaxReturnType {
  success: boolean;
  shares: string;
  sharesWithSlippage: string;
}

function isZero(value: bigint): boolean {
  return value === 0n;
}

function deriveColParams(params: ISharesMinMaxParams) {
  const {
    supply0AmountInWei,
    supply1AmountInWei,
    supplySlippage,
    depositAndBorrow,
    vault,
  } = params;
  const supplyToken = vault.supplyToken.token0;
  const supplyToken1 = vault.supplyToken.token1;

  const dex = vault.supplyDexData;
  const token0AmountWei = BigInt(supply0AmountInWei || "0");
  const token1AmountWei = BigInt(supply1AmountInWei || "0");
  const token0Decimals = BigInt(supplyToken.decimals);
  const token1Decimals = BigInt(supplyToken1?.decimals ?? supplyToken.decimals);
  const dexFee = BigInt(dex.fee);
  const totalShares = BigInt(dex.totalShares);
  const colReserves: ColReserves = {
    token0ImaginaryReserves: BigInt(dex.token0ImaginaryReserves),
    token0RealReserves: BigInt(dex.token0RealReserves),
    token1ImaginaryReserves: BigInt(dex.token1ImaginaryReserves),
    token1RealReserves: BigInt(dex.token1RealReserves),
  };
  const pex: Pex = {
    geometricMean: BigInt(dex.pex.geometricMean),
    lowerRange: BigInt(dex.pex.lowerRange),
    upperRange: BigInt(dex.pex.upperRange),
  };
  return {
    token0AmountWei,
    token1AmountWei,
    token0Decimals,
    token1Decimals,
    slippage: supplySlippage,
    dexFee,
    totalShares,
    colReserves,
    pex,
    depositAndBorrow,
  };
}

function deriveDebtParams(params: ISharesMinMaxParams) {
  const {
    borrow0AmountInWei,
    borrow1AmountInWei,
    borrowSlippage,
    depositAndBorrow,
    vault,
  } = params;
  const borrowToken = vault.borrowToken.token0;
  const borrowToken1 = vault.borrowToken.token1;

  const dex = vault.borrowDexData;
  const token0AmountWei = BigInt(borrow0AmountInWei || "0");
  const token1AmountWei = BigInt(borrow1AmountInWei || "0");
  const token0Decimals = BigInt(borrowToken.decimals);
  const token1Decimals = BigInt(borrowToken1?.decimals ?? borrowToken.decimals);
  const dexFee = BigInt(dex.fee);
  const totalShares = BigInt(dex.totalShares);
  const debtReserves: DebtReserves = {
    token0Debt: BigInt(dex.token0Debt),
    token1Debt: BigInt(dex.token1Debt),
    token0ImaginaryReserves: BigInt(dex.token0ImaginaryReserves),
    token0RealReserves: BigInt(dex.token0RealReserves),
    token1ImaginaryReserves: BigInt(dex.token1ImaginaryReserves),
    token1RealReserves: BigInt(dex.token1RealReserves),
  };
  const pex: Pex = {
    geometricMean: BigInt(dex.pex.geometricMean),
    lowerRange: BigInt(dex.pex.lowerRange),
    upperRange: BigInt(dex.pex.upperRange),
  };
  return {
    token0AmountWei,
    token1AmountWei,
    token0Decimals,
    token1Decimals,
    slippage: borrowSlippage,
    dexFee,
    totalShares,
    debtReserves,
    pex,
    depositAndBorrow,
  };
}

/**
 * Calculates min/max collateral shares for deposit or withdraw (T2/T4).
 * Deposit: pass any of supply0AmountInWei and supply1AmountInWei for the amounts you wish to deposit.
 * Withdraw: based on your initial deposit, pass one token amount if paying entirely from that token,
 * or pass both amounts in any desired split.
 * Uses deposit when depositAndBorrow is true, otherwise withdraw.
 */
export function calculateColSharesMinMax(
  params: ISharesMinMaxParams,
): SharesMinMaxReturnType | undefined {
  if (!params.vault.supplyToken.token1) return undefined;

  const {
    token0AmountWei,
    token1AmountWei,
    token0Decimals,
    token1Decimals,
    slippage,
    dexFee,
    totalShares,
    colReserves,
    pex,
    depositAndBorrow,
  } = deriveColParams(params);

  if (isZero(token0AmountWei) && isZero(token1AmountWei)) {
    return undefined;
  }

  try {
    const fn = depositAndBorrow ? deposit : withdraw;

    const resp = depositAndBorrow
      ? (fn as typeof deposit)(
          token0AmountWei,
          token1AmountWei,
          token0Decimals,
          token1Decimals,
          slippage,
          dexFee,
          totalShares,
          colReserves,
        )
      : (fn as typeof withdraw)(
          token0AmountWei,
          token1AmountWei,
          token0Decimals,
          token1Decimals,
          slippage,
          dexFee,
          totalShares,
          colReserves,
          pex,
        );

    if (!resp.success) {
      throw new Error("success is false");
    }

    const sharesBn = resp.shares as bigint;
    const sharesWithSlippageBn = resp.sharesWithSlippage as bigint;

    return {
      success: resp.success,
      shares: depositAndBorrow ? String(sharesBn) : String(-sharesBn),
      sharesWithSlippage: depositAndBorrow
        ? String(sharesWithSlippageBn)
        : String(-sharesWithSlippageBn),
    };
  } catch {
    return undefined;
  }
}

/**
 * Calculates min/max debt shares for borrow or payback (T3).
 * Borrow: pass both borrow0AmountInWei and borrow1AmountInWei for the amounts you wish to borrow.
 * Payback: based on your initial borrow, pass one token amount if repaying entirely from that token,
 * or pass both amounts in any desired split.
 * Uses borrow when depositAndBorrow is true, otherwise payback.
 */
export function calculateDebtSharesMinMax(
  params: ISharesMinMaxParams,
): SharesMinMaxReturnType | undefined {
  if (!params.vault.borrowToken.token1) return undefined;

  const {
    token0AmountWei,
    token1AmountWei,
    token0Decimals,
    token1Decimals,
    slippage,
    dexFee,
    totalShares,
    debtReserves,
    pex,
    depositAndBorrow,
  } = deriveDebtParams(params);

  if (isZero(token0AmountWei) && isZero(token1AmountWei)) {
    return undefined;
  }

  try {
    const fn = depositAndBorrow ? borrow : payback;

    const resp = depositAndBorrow
      ? (fn as typeof borrow)(
          token0AmountWei,
          token1AmountWei,
          token0Decimals,
          token1Decimals,
          slippage,
          dexFee,
          totalShares,
          debtReserves,
          pex,
        )
      : (fn as typeof payback)(
          token0AmountWei,
          token1AmountWei,
          token0Decimals,
          token1Decimals,
          slippage,
          dexFee,
          totalShares,
          debtReserves,
        );

    if (!resp.success) {
      throw new Error("success is false");
    }

    const sharesBn = resp.shares as bigint;
    const sharesWithSlippageBn = resp.sharesWithSlippage as bigint;

    return {
      success: resp.success,
      shares: depositAndBorrow ? String(sharesBn) : String(-sharesBn),
      sharesWithSlippage: depositAndBorrow
        ? String(sharesWithSlippageBn)
        : String(-sharesWithSlippageBn),
    };
  } catch {
    return undefined;
  }
}