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 ".";

/**
 * Allows staking into a merged pool.
 * @returns
 */
export const useMergedStake = (): ((amount: TokenAmount) => Promise<void>) => {
  const { sdkMut } = useQuarrySDK();
  const { handleTX, handleTXs } = useSail();
  const {
    primaryToken,
    mergePool,
    mergeMiner,
    primaryRewarderKey,
    replicaRewarderKeys,
    primaryQuarry,
  } = useMergedQuarries();

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

      const hasReplicaRewarders = replicaRewarderKeys.length > 0;
      const numSteps = (mergeMiner ? 0 : 1) + (hasReplicaRewarders ? 1 : 0) + 1;

      let stepIndex = 1;
      let mmKey = mergeMiner?.publicKey;
      if (!mergeMiner) {
        const { tx: initMMTX, key: mergeMinerKey } =
          await sdkMut.mergeMine.newMM({
            pool: {
              key: mergePool.publicKey,
              data: { primaryMint: primaryToken.mintAccount },
            },
            rewarder: primaryRewarderKey,
            rewardsMint: primaryQuarry.rewardToken.mintAccount,
          });
        if (initMMTX) {
          const { success, pending, errors } = await handleTX(
            initMMTX,
            `Set up miner (1/${numSteps})`
          );
          if (!success || !pending) {
            throw new Error(
              `Error setting up miner: ${errors?.toString() ?? "unknown error"}`
            );
          }
          await pending.wait();
        }
        mmKey = mergeMinerKey;
      }

      invariant(mmKey, "mmKey");
      const mm = await sdkMut.mergeMine.loadMM({ mmKey });

      const depositTX = await mm.deposit({
        amount,
        rewarder: primaryQuarry.rewarder.publicKey,
      });

      const stakeResult = await handleTX(
        depositTX,
        `Stake tokens${numSteps > 1 ? ` (${stepIndex++}/${numSteps})` : ""}`
      );
      if (!stakeResult.pending || !stakeResult.success) {
        throw new Error(
          `Error staking tokens (primary): ${
            stakeResult.errors?.toString() ?? "unknown error"
          }`
        );
      }
      await stakeResult.pending.wait();

      if (replicaRewarderKeys.length > 0) {
        const replicaStakeTXs = await Promise.all(
          replicaRewarderKeys.map(async (replicaRewarder) => {
            return await mm.stakeReplicaMiner(replicaRewarder);
          })
        );
        const { pending, success, errors } = await handleTXs(
          replicaStakeTXs,
          `Stake tokens (${stepIndex++}/${numSteps})`
        );
        if (!success) {
          throw new Error(
            `Error staking tokens (replica): ${
              errors?.toString() ?? "unknown error"
            }`
          );
        }
        await Promise.all(pending.map((p) => p.wait()));
      }
    },
    [
      handleTX,
      handleTXs,
      mergeMiner,
      mergePool,
      primaryQuarry,
      primaryRewarderKey,
      primaryToken.mintAccount,
      replicaRewarderKeys,
      sdkMut,
    ]
  );
};