Convert snUSD into nUSD

export function quoteSnusdInNusd(input: {
  shares: bigint;
  totalShares: bigint;
  stakingVaultNusd6: bigint;
  reservedPendingClaimsNusd6: bigint;
}) {
  if (input.shares === 0n || input.totalShares === 0n) return 0n;
  const redeemable = input.stakingVaultNusd6 > input.reservedPendingClaimsNusd6
    ? input.stakingVaultNusd6 - input.reservedPendingClaimsNusd6
    : 0n;
  return input.shares * redeemable / input.totalShares;
}

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

function keys(client: NestClient) {
  const owner = client.wallet.publicKey;
  const stakeProgram = new PublicKey(client.deployment.programs.nestStake);
  const stakingState = new PublicKey(client.deployment.accounts.stakingState);
  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);
  return {
    owner,
    stakeProgram,
    stakingState,
    nusdMint,
    snusdMint,
    nusdTokenProgram,
    snusdTokenProgram,
    ownerSnusdAccount: ata(snusdMint, owner, snusdTokenProgram),
    pendingWithdrawal: pendingWithdrawalPda(stakeProgram, owner, stakingState),
  };
}

export async function requestUnstake(client: NestClient, snusdAmount: string) {
  const account = keys(client);
  const shares = toBaseUnits(snusdAmount, 6);
  if (shares <= 0n) throw new Error("Unstake shares must be positive");
  if (await client.connection.getAccountInfo(account.pendingWithdrawal, "confirmed")) {
    throw new Error("Finish or recover the existing pending unstake first");
  }
  const ownerSnusd = await getAccount(
    client.connection,
    account.ownerSnusdAccount,
    "confirmed",
    account.snusdTokenProgram,
  );
  if (ownerSnusd.amount < shares) throw new Error("Insufficient snUSD balance");

  const transaction = new Transaction().add(await client.stake.methods
    .requestUnstake(new BN(shares.toString()))
    .accountsStrict({
      stakingState: account.stakingState,
      snusdMint: account.snusdMint,
      ownerSnusdAccount: account.ownerSnusdAccount,
      pendingWithdrawal: account.pendingWithdrawal,
      owner: account.owner,
      snusdTokenProgram: account.snusdTokenProgram,
      systemProgram: SystemProgram.programId,
    })
    .instruction());
  const signature = await simulateSignSend(client, transaction);

  const [pending, snusdAfter, state] = await Promise.all([
    client.stake.account.pendingWithdrawalAccount.fetch(account.pendingWithdrawal),
    getAccount(
      client.connection,
      account.ownerSnusdAccount,
      "confirmed",
      account.snusdTokenProgram,
    ),
    client.stake.account.stakingState.fetch(account.stakingState),
  ]);
  if (
    BigInt(pending.shares.toString()) !== shares
    || ownerSnusd.amount - snusdAfter.amount !== shares
  ) throw new Error(`Unstake request verification failed for ${signature}`);

  const requestTs = BigInt(pending.requestTs.toString());
  return {
    signature,
    pendingWithdrawal: account.pendingWithdrawal,
    unlockTs: requestTs + BigInt(state.cooldownSeconds.toString()),
    claimDeadlineTs: BigInt(pending.claimDeadlineTs.toString()),
  };
}

export async function claimUnstakedNusd(client: NestClient) {
  const account = keys(client);
  const [pending, state, protocol] = await Promise.all([
    client.stake.account.pendingWithdrawalAccount.fetch(account.pendingWithdrawal),
    client.stake.account.stakingState.fetch(account.stakingState),
    client.core.account.protocol.fetch(new PublicKey(client.deployment.accounts.protocol)),
  ]);
  const now = await confirmedChainTime(client.connection);
  const unlock = BigInt(pending.requestTs.toString()) + BigInt(state.cooldownSeconds.toString());
  const deadline = BigInt(pending.claimDeadlineTs.toString());
  if (now < unlock) throw new Error(`Cooldown active until ${unlock}`);
  if (now > deadline) throw new Error("Claim window expired; restore snUSD instead");
  if (BigInt(protocol.badDebtNusd.toString()) !== 0n) {
    throw new Error("Claims are unavailable while protocol bad debt is outstanding");
  }

  const transaction = new Transaction();
  const ownerNusdAccount = await addAtaIfMissing(
    client,
    transaction,
    account.nusdMint,
    account.owner,
    account.nusdTokenProgram,
  );
  const nUsdBefore = await client.connection.getAccountInfo(ownerNusdAccount, "confirmed")
    ? (await getAccount(
        client.connection,
        ownerNusdAccount,
        "confirmed",
        account.nusdTokenProgram,
      )).amount
    : 0n;
  transaction.add(await client.stake.methods
    .completeUnstake()
    .accountsStrict({
      stakingState: account.stakingState,
      pendingWithdrawal: account.pendingWithdrawal,
      nusdMint: account.nusdMint,
      stakingNusdVault: new PublicKey(client.deployment.vaults.stakingNusdVault),
      ownerNusdAccount,
      protocol: new PublicKey(client.deployment.accounts.protocol),
      nestCoreProgram: new PublicKey(client.deployment.programs.nestCore),
      owner: account.owner,
      nusdTokenProgram: account.nusdTokenProgram,
    })
    .instruction());

  const signature = await simulateSignSend(client, transaction);
  const [nUsdAfter, pendingAfter] = await Promise.all([
    getAccount(client.connection, ownerNusdAccount, "confirmed", account.nusdTokenProgram),
    client.stake.account.pendingWithdrawalAccount.fetchNullable(account.pendingWithdrawal),
  ]);
  if (pendingAfter !== null || nUsdAfter.amount <= nUsdBefore) {
    throw new Error(`Unstake claim verification failed for ${signature}`);
  }
  return { signature, receivedNusd6: nUsdAfter.amount - nUsdBefore };
}

export async function restoreExpiredSnusd(client: NestClient) {
  const account = keys(client);
  const pending = await client.stake.account.pendingWithdrawalAccount.fetch(
    account.pendingWithdrawal,
  );
  const deadline = BigInt(pending.claimDeadlineTs.toString());
  if (await confirmedChainTime(client.connection) <= deadline) {
    throw new Error("Claim window is still active");
  }

  const transaction = new Transaction();
  const ownerSnusdAccount = await addAtaIfMissing(
    client,
    transaction,
    account.snusdMint,
    account.owner,
    account.snusdTokenProgram,
  );
  transaction.add(await client.stake.methods
    .cancelExpiredUnstake()
    .accountsStrict({
      stakingState: account.stakingState,
      pendingWithdrawal: account.pendingWithdrawal,
      snusdMint: account.snusdMint,
      ownerSnusdAccount,
      owner: account.owner,
      snusdTokenProgram: account.snusdTokenProgram,
    })
    .instruction());
  const signature = await simulateSignSend(client, transaction);
  const pendingAfter = await client.stake.account.pendingWithdrawalAccount.fetchNullable(
    account.pendingWithdrawal,
  );
  if (pendingAfter !== null) throw new Error(`Recovery verification failed for ${signature}`);
  return signature;
}