import * as anchor from "@project-serum/anchor";
import { Program, BN } from "@project-serum/anchor";
import { PsyStake } from "../target/types/psy_stake";
import {
  deriveStakePool,
  deriveStakingTokenAccount,
  deriveRewardPoolPdas,
  deriveRewardRecord,
  deriveStakingRecord,
} from "./pdas";

import { assert } from "chai";
import { Keypair, PublicKey } from "@solana/web3.js";
const { setupTestAccounts, MIN_EPOCH_DURATION } = require("./setup");
const { getAccountData, waitForStakePoolEpoch } = require("./utils");

describe("Test distribution strategies.", () => {
  const program = anchor.workspace.PsyStake as Program<PsyStake>;
  const provider = anchor.Provider.env();
  anchor.setProvider(provider);

  let userUsdcTokenAccount: PublicKey,
    userBtcTokenAccount: PublicKey,
    network,
    userWalletKeypair: Keypair,
    mintAuthority: PublicKey,
    stakePoolAuthority: Keypair,
    rewardPoolAuthority: Keypair;

  const uniqueSeed = 1;

  let stakePool: PublicKey,
    stakingTokenAccount: PublicKey,
    stakingTokenMint: PublicKey,
    rewardPool: PublicKey,
    rewardTokenAccount: PublicKey,
    allocatedTokenAccount: PublicKey,
    rewardRecord1: PublicKey,
    rewardRecord2: PublicKey,
    rewardRecord3: PublicKey,
    rewardRecord4: PublicKey,
    stakingRecord;

  const startingEpoch = 1;
  let poolId = 1;
  const epochDuration = new BN(MIN_EPOCH_DURATION);
  const stakeAmount = new BN(100);
  const distributionType = { percentage: {} };
  const rewardPerEpoch = new BN(10);
  const epochRewardDecimals = 8;

  const stakeTokenHelper = async (stakeAmount, lockUpPeriod) => {
    await program.methods
      .stakeToken(stakeAmount, lockUpPeriod)
      .accounts({
        recordOwner: userWalletKeypair.publicKey,
        stakePool,
        stakingRecord,
        ownerTokenAccount: userUsdcTokenAccount,
        stakingTokenAccount,
      })
      .signers([userWalletKeypair])
      .rpc();
  };

  const updateRewardRecordHelper = async (
    curRewardRecord,
    prevRewardRecord
  ) => {
    await program.methods
      .updateRewardRecord()
      .accounts({
        stakePool,
        rewardPool,
        curRewardRecord,
        prevRewardRecord,
        rewardTokenAccount,
        allocatedTokenAccount,
      })
      .rpc();
  };

  const unstakeTokenHelper = async (unstakeAmount) => {
    await program.methods
      .unstakeToken(unstakeAmount)
      .accounts({
        recordOwner: userWalletKeypair.publicKey,
        stakePool,
        stakingRecord,
        ownerTokenAccount: userUsdcTokenAccount,
        stakingTokenAccount,
      })
      .signers([userWalletKeypair])
      .rpc();
  };

  const claimRewardHelper = async (startRewardRecord, endRewardRecord) => {
    await program.methods
      .claimReward()
      .accounts({
        recordOwner: userWalletKeypair.publicKey,
        stakePool,
        stakingRecord,
        rewardPool,
        startRewardRecord,
        endRewardRecord,
        ownerTokenAccount: userBtcTokenAccount,
        allocatedTokenAccount,
      })
      .signers([userWalletKeypair])
      .rpc();
  };

  const createRewardRecordHelper = async (
    newRewardRecord,
    prevRewardRecord
  ) => {
    await program.methods
      .createRewardRecord()
      .accounts({
        payer: userWalletKeypair.publicKey,
        stakePool,
        rewardPool,
        newRewardRecord,
        prevRewardRecord,
        rewardTokenAccount,
        allocatedTokenAccount,
      })
      .signers([userWalletKeypair])
      .rpc();
  };

  before(async () => {
    ({
      userUsdcTokenAccount,
      userBtcTokenAccount,
      network,
      userWalletKeypair,
      mintAuthority,
      stakePoolAuthority,
      rewardPoolAuthority,
    } = await setupTestAccounts(program, provider));
    stakingTokenMint = network.usdcToken.publicKey;

    [stakePool] = await deriveStakePool(
      program,
      network.usdcToken.publicKey,
      new BN(uniqueSeed)
    );
    [stakingTokenAccount] = await deriveStakingTokenAccount(program, stakePool);

    const nextEpochStartTime = new BN(new Date().getTime() / 1000 + 3);

    // Create a StakePool.
    await program.methods
      .createStakePool(uniqueSeed, nextEpochStartTime, epochDuration)
      .accounts({
        stakePoolAuthority: stakePoolAuthority.publicKey,
        stakePool,
        stakingTokenMint,
        stakingTokenAccount,
      })
      .signers([stakePoolAuthority])
      .rpc();

    // Create a RewardPool.
    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      poolId,
      startingEpoch
    );
    ({ rewardPool, rewardTokenAccount, allocatedTokenAccount } = poolAccounts);
    rewardRecord1 = poolAccounts.rewardRecord;
    [rewardRecord2] = await deriveRewardRecord(program, rewardPool, 2);
    [rewardRecord3] = await deriveRewardRecord(program, rewardPool, 3);
    [rewardRecord4] = await deriveRewardRecord(program, rewardPool, 4);

    await program.methods
      .createRewardPool(
        startingEpoch,
        distributionType,
        rewardPerEpoch,
        epochRewardDecimals
      )
      .accounts({
        stakePoolAuthority: stakePoolAuthority.publicKey,
        rewardPoolAuthority: rewardPoolAuthority.publicKey,
        stakePool,
        rewardPool: poolAccounts.rewardPool,
        rewardRecord: poolAccounts.rewardRecord,
        rewardTokenMint: network.btcToken.publicKey,
        rewardTokenAccount: poolAccounts.rewardTokenAccount,
        allocatedTokenAccount: poolAccounts.allocatedTokenAccount,
      })
      .signers([stakePoolAuthority])
      .rpc();

    // Create a StakingRecord.
    [stakingRecord] = await deriveStakingRecord(
      program,
      userWalletKeypair.publicKey,
      stakePool
    );
    await program.methods
      .createStakingRecord()
      .accounts({
        recordOwner: userWalletKeypair.publicKey,
        stakePool,
        stakingRecord,
      })
      .signers([userWalletKeypair])
      .rpc();

    // Mint tokens to reward account.
    await network.btcToken.mintTo(
      rewardTokenAccount,
      mintAuthority,
      [],
      10 ** 6
    );
  });

  it("Update RewardRecord with percentage distribution strategy.", async () => {
    // Stake token into a pool, to increment reward units.
    const lockUpPeriod = 0;
    await stakeTokenHelper(stakeAmount, lockUpPeriod);

    const isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      1
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    const accountDataBefore = await getAccountData(program, {
      rewardRecord: rewardRecord1,
      rewardTokenAccount,
      allocatedTokenAccount,
    });

    await updateRewardRecordHelper(rewardRecord1, rewardRecord1);

    const accountDataAfter = await getAccountData(program, {
      rewardRecord: rewardRecord1,
      rewardTokenAccount,
      allocatedTokenAccount,
      rewardPool,
      stakePool,
    });

    const expectedRewardAllocated = accountDataBefore.rewardTokenAccountAmount
      .mul(accountDataAfter.rewardPool.percentageRewardMbpsPerEpoch)
      .div(new BN(10000 * 1000));
    const expectedEpochReward = expectedRewardAllocated
      .mul(new BN(10 ** epochRewardDecimals))
      .div(accountDataAfter.stakePool.totalRewardUnits);

    // Verify that RewardRecord allocates reward tokens.
    assert(
      accountDataAfter.rewardRecord.totalRewardAllocated.eq(
        expectedRewardAllocated
      )
    );
    assert(accountDataAfter.rewardRecord.epochReward.eq(expectedEpochReward));

    // Verify that epoch's reward is added to the cummulative reward.
    assert(
      accountDataAfter.rewardRecord.cummulativeReward
        .sub(accountDataBefore.rewardRecord.cummulativeReward)
        .eq(expectedEpochReward)
    );

    // Verify amounts in token account.
    assert(
      accountDataBefore.rewardTokenAccountAmount
        .sub(accountDataAfter.rewardTokenAccountAmount)
        .eq(expectedRewardAllocated)
    );
    assert(
      accountDataAfter.allocatedTokenAccountAmount
        .sub(accountDataBefore.allocatedTokenAccountAmount)
        .eq(expectedRewardAllocated)
    );
  });

  it("Create RewardRecord with percentage distribution strategy.", async () => {
    // Change to a different percentage distribution per epoch.
    await program.methods
      .updateRewardPool(distributionType, new BN(75), true)
      .accounts({
        rewardPoolAuthority: rewardPoolAuthority.publicKey,
        stakePool,
        rewardPool,
      })
      .signers([rewardPoolAuthority])
      .rpc();

    const isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      2
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    const accountDataBefore = await getAccountData(program, {
      rewardRecord: rewardRecord1,
      rewardTokenAccount,
      allocatedTokenAccount,
    });

    // Initialize a RewardRecord for epoch 2.
    await createRewardRecordHelper(rewardRecord2, rewardRecord1);

    const accountDataAfter = await getAccountData(program, {
      rewardRecord: rewardRecord2,
      rewardTokenAccount,
      allocatedTokenAccount,
      rewardPool,
      stakePool,
    });

    console.log(accountDataBefore.rewardTokenAccountAmount.toNumber());
    const expectedRewardAllocated = accountDataBefore.rewardTokenAccountAmount
      .mul(accountDataAfter.rewardPool.percentageRewardMbpsPerEpoch)
      .div(new BN(10000 * 1000));
    const expectedEpochReward = expectedRewardAllocated
      .mul(new BN(10 ** epochRewardDecimals))
      .div(accountDataAfter.stakePool.totalRewardUnits);
    console.log(expectedEpochReward.toNumber());
    console.log(expectedRewardAllocated.toNumber());

    // Verify that RewardRecord allocates reward tokens.
    assert(
      accountDataAfter.rewardRecord.totalRewardAllocated.eq(
        expectedRewardAllocated
      )
    );
    assert(accountDataAfter.rewardRecord.epochReward.eq(expectedEpochReward));

    // Verify that epoch's reward is added to the cummulative reward.
    assert(
      accountDataAfter.rewardRecord.cummulativeReward
        .sub(accountDataBefore.rewardRecord.cummulativeReward)
        .eq(expectedEpochReward)
    );

    // Verify amounts in token account.
    assert(
      accountDataBefore.rewardTokenAccountAmount
        .sub(accountDataAfter.rewardTokenAccountAmount)
        .eq(expectedRewardAllocated)
    );
    assert(
      accountDataAfter.allocatedTokenAccountAmount
        .sub(accountDataBefore.allocatedTokenAccountAmount)
        .eq(expectedRewardAllocated)
    );
  });
});