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

import { ref, watchEffect, computed } 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'

/**
 * 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) {
  const { activePositions } = usePositions()
  const stake = useStake()
  const wallet = useWallet()
  const statCollection = useStatCollection()
  const adjustTaxRate = useAdjustTaxRate()

  const snatchablePositions = ref<Position[]>([])
  const shortfallShares = ref<bigint>(0n)
  const floorTax = ref(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 maxTaxRateDecimal = (stake.taxRate ?? 0) / 100
      const includeOwned = demo
      const recipient = wallet.account.address ?? null

      const eligiblePositions = activePositions.value.filter((position: Position) => {
        if (position.taxRate >= maxTaxRateDecimal) {
          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,
  }
}