import { MergeMiner } from "@quarryprotocol/quarry-sdk";
import { useSail } from "@saberhq/sail";
import type { TokenAmount } from "@saberhq/token-utils";
import { useCallback } from "react";
import invariant from "tiny-invariant";

import { useQuarrySDK } from "../..";
import { useMergedQuarries } from "./context";

/**
 * Withdraws all staked tokens from the merge pool.
 * @returns
 */
export const useMergedWithdraw = (): ((
  amount: TokenAmount
) => Promise<void>) => {
  const { sdkMut } = useQuarrySDK();
  const { handleTX, handleTXs } = useSail();
  const {
    mergePool,
    mergeMiner,
    replicaRewarderKeys,
    primaryQuarry,
    quarries,
  } = useMergedQuarries();

  return useCallback(
    async (amount: TokenAmount) => {
      invariant(sdkMut, "sdk not connected");
      invariant(primaryQuarry, "primary quarry");
      invariant(mergePool, "merge pool");
      invariant(mergeMiner, "merge miner");

      const mm = new MergeMiner(
        sdkMut.mergeMine,
        {
          key: mergePool.publicKey,
          data: mergePool.account,
        },
        {
          key: mergeMiner.publicKey,
          data: mergeMiner.account,
        }
      );

      const activeQuarries = quarries.filter((q) => q.miner);
      const replicaUnstakeTXs = await Promise.all(
        activeQuarries
          .filter((q) => q.isReplica)
          .map(async (q) => {
            return await mm.unstakeAllReplica(q.rewarder.publicKey);
          })
      );

      const shouldRestake = primaryQuarry.userStakedAmount?.greaterThan(amount);
      const numSteps =
        (replicaUnstakeTXs.length > 0 ? 1 : 0) + (shouldRestake ? 1 : 0) + 1;
      let nextStepNumber = 1;

      if (replicaUnstakeTXs.length > 0) {
        const unstakes = await handleTXs(
          replicaUnstakeTXs,
          `Unstake tokens (${nextStepNumber++}/${numSteps})`
        );
        if (!unstakes.success) {
          throw new Error(
            `Error unstaking replica tokens: ${
              unstakes.errors?.toString() ?? "unkown error"
            }`
          );
        }
        await Promise.all(unstakes.pending.map((p) => p.wait()));
      }

      const primaryUnstakeTX = await mm.withdraw({
        amount,
        rewarder: primaryQuarry.rewarder.publicKey,
      });
      const unstakeResult = await handleTX(
        primaryUnstakeTX,
        `Unstake tokens (${nextStepNumber++}/${numSteps})`
      );
      if (!unstakeResult.success || !unstakeResult.pending) {
        throw new Error(
          `Error unstaking primary tokens: ${
            unstakeResult.errors?.toString() ?? "unknown errors"
          }`
        );
      }
      await unstakeResult.pending.wait();

      // restake the replicas if the balance is still non-zero
      if (shouldRestake) {
        const replicaStakeTXs = await Promise.all(
          replicaRewarderKeys.map(async (replicaRewarder) => {
            return await mm.stakeReplicaMiner(replicaRewarder);
          })
        );
        const { pending, success, errors } = await handleTXs(
          replicaStakeTXs,
          `Unstake tokens (${nextStepNumber++}/${numSteps})`
        );
        if (!success) {
          throw new Error(
            `Error restaking replica tokens: ${
              errors?.toString() ?? "unknown error"
            }`
          );
        }
        await Promise.all(pending.map((p) => p.wait()));
      }
    },
    [
      handleTX,
      handleTXs,
      mergeMiner,
      mergePool,
      primaryQuarry,
      quarries,
      replicaRewarderKeys,
      sdkMut,
    ]
  );
};