Withdraw collateral

import { BN } from "@coral-xyz/anchor";
import { PublicKey, Transaction } from "@solana/web3.js";
import { getAccount } from "@solana/spl-token";
import {
  addAtaIfMissing,
  borrowVaultPda,
  marketBySymbol,
  oraclePda,
  simulateSignSend,
  toBaseUnits,
  type NestClient,
} from "../nest-client";
import {
  buildOracleRefreshInstructions,
  fetchCurrentOraclePrices,
} from "../nest-oracle";
import { readCurrentDebt } from "./read-current-debt";

export async function withdrawCollateral(
  client: NestClient,
  symbol: string,
  userAmount: string,
) {
  const owner = client.wallet.publicKey;
  const market = marketBySymbol(client, symbol);
  const amount = toBaseUnits(userAmount, market.decimals);
  if (amount <= 0n) throw new Error("Withdrawal amount must be positive");

  const coreProgramId = new PublicKey(client.deployment.programs.nestCore);
  const protocol = new PublicKey(client.deployment.accounts.protocol);
  const collateralConfig = new PublicKey(market.collateralConfig);
  const collateralMint = new PublicKey(market.mint);
  const collateralVault = new PublicKey(market.collateralVault);
  const collateralTokenProgram = new PublicKey(market.tokenProgram);
  const vault = borrowVaultPda(coreProgramId, owner, collateralConfig);
  const oracle = oraclePda(coreProgramId, collateralConfig);

  const [config, vaultBefore, currentDebt] = await Promise.all([
    client.core.account.collateralConfig.fetch(collateralConfig),
    client.core.account.vault.fetch(vault),
    readCurrentDebt(client, symbol),
  ]);
  if (config.withdrawsPaused) throw new Error(`${symbol} withdrawals are paused`);
  const collateralBefore = BigInt(vaultBefore.collateralRaw.toString());
  if (amount > collateralBefore) throw new Error("Withdrawal exceeds deposited collateral");

  const transaction = new Transaction();
  const ownerCollateralAccount = await addAtaIfMissing(
    client,
    transaction,
    collateralMint,
    owner,
    collateralTokenProgram,
  );

  if (currentDebt.projectedDebtNusd6 > 0n) {
    // Calculate maximumWithdrawRaw from the Position health guide first and
    // keep a buffer below it. The program uses the refreshed on-chain price.
    const currentPrice = (await fetchCurrentOraclePrices([market])).get(market.symbol);
    if (!currentPrice) throw new Error(`No current price for ${market.symbol}`);
    const prefixCount = transaction.instructions.length;
    transaction.add(...buildOracleRefreshInstructions(
      client,
      market,
      prefixCount,
      currentPrice.oracleUpdate,
    ));
  }

  transaction.add(await client.core.methods
    .withdrawWithOracle(new BN(amount.toString()))
    .accountsStrict({
      protocol,
      collateralConfig,
      vault,
      oracle,
      collateralMint,
      collateralVault,
      ownerCollateralAccount,
      owner,
      collateralTokenProgram,
    })
    .instruction());

  const ownerBefore = await client.connection.getAccountInfo(ownerCollateralAccount, "confirmed")
    ? (await getAccount(
        client.connection,
        ownerCollateralAccount,
        "confirmed",
        collateralTokenProgram,
      )).amount
    : 0n;
  const signature = await simulateSignSend(client, transaction);

  const [vaultAfter, ownerAfter] = await Promise.all([
    client.core.account.vault.fetch(vault),
    getAccount(client.connection, ownerCollateralAccount, "confirmed", collateralTokenProgram),
  ]);
  const recordedAfter = BigInt(vaultAfter.collateralRaw.toString());
  if (recordedAfter !== collateralBefore - amount || ownerAfter.amount !== ownerBefore + amount) {
    throw new Error(`Withdrawal verification failed for ${signature}`);
  }
  return { signature, remainingCollateralRaw: recordedAfter };
}