harb/web-app/src/composables/useSnatchSelection.ts

import { ref, watchEffect, computed, type Ref } from 'vue';
import { usePositions, type Position } from './usePositions';
import { useStake } from './useStake';
import { useWallet } from './useWallet';
import { useStatCollection } from './useStatCollection';
import { useAdjustTaxRate } from './useAdjustTaxRates';
import { calculateSnatchShortfall } from 'kraiken-lib/staking';
import { selectSnatchPositions, minimumTaxRate, type SnatchablePosition } from 'kraiken-lib/snatch';
import { DEFAULT_CHAIN_ID } from '@/config';

/**
 * Converts Kraiken token assets to shares using the same formula as Stake.sol:
 * shares = (assets * stakeTotalSupply) / kraikenTotalSupply
 *
 * @param assets - Amount of Kraiken tokens
 * @param kraikenTotalSupply - Total supply of Kraiken tokens
 * @param stakeTotalSupply - Total supply of stake shares (constant from contract)
 * @returns Number of shares corresponding to the assets
 */
function assetsToSharesLocal(assets: bigint, kraikenTotalSupply: bigint, stakeTotalSupply: bigint): bigint {
  if (kraikenTotalSupply === 0n) {
    return 0n;
  }
  // Equivalent to: assets.mulDiv(stakeTotalSupply, kraikenTotalSupply, Math.Rounding.Down)
  return (assets * stakeTotalSupply) / kraikenTotalSupply;
}

export function useSnatchSelection(demo = false, taxRateIndex?: Ref<number>, chainId?: number) {
  const wallet = useWallet();
  const resolvedChainId = chainId ?? wallet.account.chainId ?? DEFAULT_CHAIN_ID;
  const { activePositions } = usePositions(resolvedChainId);
  const stake = useStake();
  const statCollection = useStatCollection(resolvedChainId);
  const adjustTaxRate = useAdjustTaxRate();

  const snatchablePositions = ref<Position[]>([]);
  const shortfallShares = ref<bigint>(0n);
  const floorTax = ref(adjustTaxRate.taxRates[0]?.year ?? 1);
  let selectionRun = 0;

  const openPositionsAvailable = computed(() => shortfallShares.value <= 0n);

  function getMinFloorTax() {
    const minRate = minimumTaxRate(activePositions.value, 0);
    return Math.round(minRate * 100);
  }

  watchEffect(onCleanup => {
    const runId = ++selectionRun;
    let cancelled = false;
    onCleanup(() => {
      cancelled = true;
    });

    // No longer async since we compute shares locally
    const compute = () => {
      if (statCollection.stakeTotalSupply === 0n) {
        shortfallShares.value = 0n;
        if (!cancelled && runId === selectionRun) {
          snatchablePositions.value = [];
          floorTax.value = getMinFloorTax();
        }
        return;
      }

      const stakingShares = stake.stakingAmountShares ?? 0n;
      const shortfall = calculateSnatchShortfall(statCollection.outstandingStake, stakingShares, statCollection.stakeTotalSupply, 2n, 10n);

      shortfallShares.value = shortfall;

      if (shortfall <= 0n) {
        if (!cancelled && runId === selectionRun) {
          snatchablePositions.value = [];
          floorTax.value = getMinFloorTax();
        }
        return;
      }

      const stakeTaxRateIndex = (stake as { taxRateIndex?: number }).taxRateIndex;
      const selectedTaxRateIndex = taxRateIndex?.value ?? stakeTaxRateIndex;
      const maxTaxRateDecimal =
        typeof selectedTaxRateIndex === 'number' && Number.isInteger(selectedTaxRateIndex)
          ? (adjustTaxRate.taxRates[selectedTaxRateIndex]?.decimal ?? Number.POSITIVE_INFINITY)
          : Number.POSITIVE_INFINITY;
      const includeOwned = demo;
      const recipient = wallet.account.address ?? null;

      const eligiblePositions = activePositions.value.filter((position: Position) => {
        // Filter by tax rate index instead of decimal to avoid floating-point issues
        const posIndex = position.taxRateIndex;
        if (typeof posIndex !== 'number' || (typeof selectedTaxRateIndex === 'number' && posIndex >= selectedTaxRateIndex)) {
          return false;
        }
        if (!includeOwned && position.iAmOwner) {
          return false;
        }
        return true;
      });

      if (eligiblePositions.length === 0) {
        if (!cancelled && runId === selectionRun) {
          snatchablePositions.value = [];
          floorTax.value = getMinFloorTax();
        }
        return;
      }

      const candidates: SnatchablePosition[] = [];

      // Compute shares locally using the same formula as Stake.sol
      for (const position of eligiblePositions) {
        const shares = assetsToSharesLocal(position.harbDeposit, statCollection.kraikenTotalSupply, statCollection.stakeTotalSupply);
        candidates.push({
          id: position.positionId,
          owner: position.owner,
          stakeShares: shares,
          taxRate: position.taxRate,
          taxRateIndex: position.taxRateIndex,
        });
      }

      const selection = selectSnatchPositions(candidates, {
        shortfallShares: shortfall,
        maxTaxRate: maxTaxRateDecimal,
        includeOwned,
        recipientAddress: recipient,
      });

      if (cancelled || runId !== selectionRun) {
        return;
      }

      if (selection.remainingShortfall > 0n) {
        snatchablePositions.value = [];
        floorTax.value = getMinFloorTax();
        return;
      }

      const positionById = new Map<bigint, Position>();
      for (const position of activePositions.value) {
        positionById.set(position.positionId, position);
      }

      const selectedPositions = selection.selected
        .map(candidate => positionById.get(candidate.id))
        .filter((value): value is Position => Boolean(value));

      snatchablePositions.value = selectedPositions;

      if (selection.maxSelectedTaxRateIndex !== undefined) {
        const nextIndex = selection.maxSelectedTaxRateIndex + 1;
        const option = adjustTaxRate.taxRates[nextIndex] ?? adjustTaxRate.taxRates[selection.maxSelectedTaxRateIndex];
        floorTax.value = option ? option.year : getMinFloorTax();
      } else {
        floorTax.value = getMinFloorTax();
      }
    };

    compute();
  });

  return {
    snatchablePositions,
    shortfallShares,
    floorTax,
    openPositionsAvailable,
  };
}