import { Program, BN } from "@project-serum/anchor";
import { PublicKey } from "@solana/web3.js";
import { PsyStake } from "./psy_stake";

export const deriveStakePool = async (
  programId: PublicKey,
  stakingTokenMint: PublicKey,
  uniqueSeed: BN
) => {
  return await PublicKey.findProgramAddress(
    [
      stakingTokenMint.toBuffer(),
      uniqueSeed.toArrayLike(Buffer, "le", 2),
      Buffer.from("StakePool"),
    ],
    programId
  );
};

export const deriveStakingTokenAccount = async (
  programId: PublicKey,
  stakePool: PublicKey
) => {
  return await PublicKey.findProgramAddress(
    [stakePool.toBuffer(), Buffer.from("StakingTokenPool")],
    programId
  );
};

export const deriveRewardPool = async (
  programId: PublicKey,
  stakePool: PublicKey,
  poolId: number
) => {
  return await PublicKey.findProgramAddress(
    [
      new BN(poolId).toArrayLike(Buffer, "le", 1),
      stakePool.toBuffer(),
      Buffer.from("RewardPool"),
    ],
    programId
  );
};

export const deriveRewardRecord = async (
  programId: PublicKey,
  rewardPool: PublicKey,
  rewardEpoch: number
) => {
  return await PublicKey.findProgramAddress(
    [
      rewardPool.toBuffer(),
      new BN(rewardEpoch).toArrayLike(Buffer, "le", 2),
      Buffer.from("RewardRecord"),
    ],
    programId
  );
};

export const deriveRewardTokenAccount = async (
  programId: PublicKey,
  rewardPool: PublicKey
) => {
  return await PublicKey.findProgramAddress(
    [rewardPool.toBuffer(), Buffer.from("RewardTokenAccount")],
    programId
  );
};

export const deriveAllocatedTokenAccount = async (
  programId: PublicKey,
  rewardPool: PublicKey
) => {
  return await PublicKey.findProgramAddress(
    [rewardPool.toBuffer(), Buffer.from("AllocatedTokenAccount")],
    programId
  );
};

export const deriveStakingRecord = async (
  programId: PublicKey,
  recordOwner: PublicKey,
  stakePool: PublicKey
) => {
  return await PublicKey.findProgramAddress(
    [
      recordOwner.toBuffer(),
      stakePool.toBuffer(),
      Buffer.from("StakingRecord"),
    ],
    programId
  );
};

export const deriveRewardPoolPdas = async (
  programId: PublicKey,
  stakePool: PublicKey,
  poolId: number,
  rewardEpoch?: number
) => {
  let rewardRecord, rewardRecordBump;
  let nextRewardRecord, nextRewardRecordBump;

  const [rewardPool, rewardPoolBump] = await deriveRewardPool(
    programId,
    stakePool,
    poolId
  );

  if (rewardEpoch != undefined) {
    [rewardRecord, rewardRecordBump] = await deriveRewardRecord(
      programId,
      rewardPool,
      rewardEpoch
    );

    [nextRewardRecord, nextRewardRecordBump] = await deriveRewardRecord(
      programId,
      rewardPool,
      rewardEpoch + 1
    );
  }
  const [rewardTokenAccount, rewardTokenAccountBump] =
    await deriveRewardTokenAccount(programId, rewardPool);
  const [allocatedTokenAccount, allocatedTokenAccountBump] =
    await deriveAllocatedTokenAccount(programId, rewardPool);

  return {
    rewardPool,
    rewardPoolBump,
    rewardRecord,
    rewardRecordBump,
    nextRewardRecord,
    nextRewardRecordBump,
    rewardTokenAccount,
    rewardTokenAccountBump,
    allocatedTokenAccount,
    allocatedTokenAccountBump,
  };
};