Convert nUSD into snUSD

export function quoteStakeShares(input: {
  amountNusd6: bigint;
  totalShares: bigint;
  stakingVaultNusd6: bigint;
  reservedPendingClaimsNusd6: bigint;
  revenueVaultBalanceNusd6: bigint;
  revenueBaselineNusd6: bigint;
}) {
  if (input.amountNusd6 <= 0n) throw new Error("Stake amount must be positive");
  if (input.totalShares === 0n) return input.amountNusd6;

  const redeemable = input.stakingVaultNusd6 > input.reservedPendingClaimsNusd6
    ? input.stakingVaultNusd6 - input.reservedPendingClaimsNusd6
    : 0n;
  const pendingRevenue = input.revenueVaultBalanceNusd6 > input.revenueBaselineNusd6
    ? input.revenueVaultBalanceNusd6 - input.revenueBaselineNusd6
    : 0n;
  const entryAssets = redeemable + pendingRevenue;
  if (entryAssets <= 0n) throw new Error("Staking pool has no entry assets");
  return input.amountNusd6 * input.totalShares / entryAssets;
}

export function minimumSharesOut(quote: bigint, slippageBps = 50n) {
  if (quote <= 0n || slippageBps < 0n || slippageBps >= 10_000n) {
    throw new Error("Invalid quote or slippage");
  }
  const minimum = quote * (10_000n - slippageBps) / 10_000n;
  return minimum > 0n ? minimum : 1n;
}

import { BN } from "@coral-xyz/anchor";
import { PublicKey, Transaction } from "@solana/web3.js";
import { getAccount } from "@solana/spl-token";
import {
  addAtaIfMissing,
  ata,
  simulateSignSend,
  toBaseUnits,
  type NestClient,
} from "../nest-client";
import { minimumSharesOut, quoteStakeShares } from "./stake-quote";

export async function stakeNusd(
  client: NestClient,
  userAmount: string,
  slippageBps = 50n,
) {
  const owner = client.wallet.publicKey;
  const amount = toBaseUnits(userAmount, 6);
  const stakingState = new PublicKey(client.deployment.accounts.stakingState);
  const protocol = new PublicKey(client.deployment.accounts.protocol);
  const nestCoreProgram = new PublicKey(client.deployment.programs.nestCore);
  const nusdMint = new PublicKey(client.deployment.mints.nusd);
  const snusdMint = new PublicKey(client.deployment.mints.snusd);
  const nusdTokenProgram = new PublicKey(client.deployment.tokenPrograms.nusd);
  const snusdTokenProgram = new PublicKey(client.deployment.tokenPrograms.snusd);
  const stakingNusdVault = new PublicKey(client.deployment.vaults.stakingNusdVault);
  const revenueNusdVault = new PublicKey(client.deployment.vaults.stakerRevenueNusdVault);
  const ownerNusdAccount = ata(nusdMint, owner, nusdTokenProgram);

  const [state, protocolAccount, revenueVault, ownerNusd] = await Promise.all([
    client.stake.account.stakingState.fetch(stakingState),
    client.core.account.protocol.fetch(protocol),
    getAccount(client.connection, revenueNusdVault, "confirmed", nusdTokenProgram),
    getAccount(client.connection, ownerNusdAccount, "confirmed", nusdTokenProgram),
  ]);
  if (state.paused) throw new Error("Staking is paused");
  if (BigInt(protocolAccount.badDebtNusd.toString()) !== 0n) {
    throw new Error("Staking is unavailable while bad debt is outstanding");
  }
  if (ownerNusd.amount < amount) throw new Error("Insufficient nUSD balance");

  const quote = quoteStakeShares({
    amountNusd6: amount,
    totalShares: BigInt(state.totalShares.toString()),
    stakingVaultNusd6: BigInt(state.stakingVaultNusd.toString()),
    reservedPendingClaimsNusd6: BigInt(state.reservedPendingClaims.toString()),
    revenueVaultBalanceNusd6: revenueVault.amount,
    revenueBaselineNusd6: BigInt(state.revenueBaselineNusd.toString()),
  });
  const minShares = minimumSharesOut(quote, slippageBps);

  const transaction = new Transaction();
  const ownerSnusdAccount = await addAtaIfMissing(
    client,
    transaction,
    snusdMint,
    owner,
    snusdTokenProgram,
  );
  const sharesBefore = await client.connection.getAccountInfo(ownerSnusdAccount, "confirmed")
    ? (await getAccount(
        client.connection,
        ownerSnusdAccount,
        "confirmed",
        snusdTokenProgram,
      )).amount
    : 0n;

  transaction.add(await client.stake.methods
    .stake(new BN(amount.toString()), new BN(minShares.toString()))
    .accountsStrict({
      stakingState,
      protocol,
      nestCoreProgram,
      nusdMint,
      snusdMint,
      ownerNusdAccount,
      stakingNusdVault,
      ownerSnusdAccount,
      revenueNusdVault,
      owner,
      nusdTokenProgram,
      snusdTokenProgram,
    })
    .instruction());

  const signature = await simulateSignSend(client, transaction);
  const [nusdAfter, snusdAfter] = await Promise.all([
    getAccount(client.connection, ownerNusdAccount, "confirmed", nusdTokenProgram),
    getAccount(client.connection, ownerSnusdAccount, "confirmed", snusdTokenProgram),
  ]);
  const receivedShares = snusdAfter.amount - sharesBefore;
  if (ownerNusd.amount - nusdAfter.amount !== amount || receivedShares < minShares) {
    throw new Error(`Stake verification failed for ${signature}`);
  }
  return { signature, paidNusd6: amount, receivedSnusd6: receivedShares };
}