import * as anchor from "@project-serum/anchor";
import { Program, BN } from "@project-serum/anchor";
import { assert } from "chai";
import { Keypair, PublicKey } from "@solana/web3.js";
import { PsyStake } from "../target/types/psy_stake";
import {
  deriveStakePool,
  deriveStakingTokenAccount,
  deriveRewardPoolPdas,
  deriveStakingRecord,
} from "./pdas";
const { setupTestAccounts, MIN_EPOCH_DURATION } = require("./setup");
const { defaultVec, waitForStakePoolEpoch } = require("./utils");

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

  let network,
    userWalletKeypair: Keypair,
    stakePoolAuthority: Keypair,
    rewardPoolAuthority: Keypair;

  const uniqueSeed = 1;

  let stakePool: PublicKey,
    stakingTokenAccount: PublicKey,
    stakingTokenMint: PublicKey,
    stakingRecord: PublicKey,
    stakingRecordBump: number;
  const distributionType = { constant: {} };
  const rewardPerEpoch = new BN(1000000);
  const epochRewardDecimals = 9;

  const createStakingRecordHelper = async () => {
    await program.methods
      .createStakingRecord()
      .accounts({
        recordOwner: userWalletKeypair.publicKey,
        stakePool,
        stakingRecord,
      })
      .signers([userWalletKeypair])
      .rpc();
  };

  const createRewardPoolHelper = async (startingEpoch, poolId) => {
    // Create a RewardPool.
    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      poolId,
      startingEpoch
    );
    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();
  };

  const updateRewadPoolHelper = async (isActive, poolId) => {
    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      poolId,
      /*startingEpoch=*/ 1 // Not used here.
    );
    await program.methods
      .updateRewardPool(distributionType, rewardPerEpoch, isActive)
      .accounts({
        rewardPoolAuthority: rewardPoolAuthority.publicKey,
        stakePool,
        rewardPool: poolAccounts.rewardPool,
      })
      .signers([rewardPoolAuthority])
      .rpc();
  };

  const epochDuration = new BN(MIN_EPOCH_DURATION);

  beforeEach(async () => {
    ({ network, userWalletKeypair, 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.
    await createRewardPoolHelper(1, 1);
  });

  it("Creates a StakingRecord before first RewardPool starts.", async () => {
    // Create a StakingRecord.
    [stakingRecord, stakingRecordBump] = await deriveStakingRecord(
      program,
      userWalletKeypair.publicKey,
      stakePool
    );
    await createStakingRecordHelper();

    const stakingRecordData = await program.account.stakingRecord.fetch(
      stakingRecord
    );
    assert(stakingRecordData.stakePool.equals(stakePool));
    assert(stakingRecordData.recordOwner.equals(userWalletKeypair.publicKey));
    assert(stakingRecordData.stakedAmount.isZero());
    assert(stakingRecordData.rewardUnits.isZero());

    // Verify that lock up should not be in the future.
    const currentTime = new Date().getTime() / 1000;
    assert(stakingRecordData.lockUpExpiry.toNumber() <= currentTime);

    assert.deepEqual(stakingRecordData.lastEpochClaimedVec, defaultVec());
    assert.equal(stakingRecordData.recordBump, stakingRecordBump);
  });

  it("Creates a StakingRecord after first RewardPool starts.", async () => {
    const isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      2
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Create a StakingRecord.
    [stakingRecord, stakingRecordBump] = await deriveStakingRecord(
      program,
      userWalletKeypair.publicKey,
      stakePool
    );
    await createStakingRecordHelper();

    const stakingRecordData = await program.account.stakingRecord.fetch(
      stakingRecord
    );
    assert(stakingRecordData.stakePool.equals(stakePool));
    assert(stakingRecordData.recordOwner.equals(userWalletKeypair.publicKey));
    assert(stakingRecordData.stakedAmount.isZero());
    assert(stakingRecordData.rewardUnits.isZero());

    // Verify that lock up should not be in the future.
    const currentTime = new Date().getTime() / 1000;
    assert(stakingRecordData.lockUpExpiry.toNumber() <= currentTime);

    // Expect that first RewardPool should be claimed for epoch 2.
    const expectedlastEpochClaimedVec = defaultVec();
    expectedlastEpochClaimedVec[0] = 2;

    assert.deepEqual(
      stakingRecordData.lastEpochClaimedVec,
      expectedlastEpochClaimedVec
    );
    assert.equal(stakingRecordData.recordBump, stakingRecordBump);
  });

  it("Creates a StakingRecord when there is an inactive RewardPool.", async () => {
    const isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      3
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Set first RewardPool as inactive.
    await updateRewadPoolHelper(false, 1);

    // Add another RewardPool, starting in epoch 5.
    await createRewardPoolHelper(5, 2);

    // Create a StakingRecord.
    [stakingRecord, stakingRecordBump] = await deriveStakingRecord(
      program,
      userWalletKeypair.publicKey,
      stakePool
    );
    await createStakingRecordHelper();

    const stakingRecordData = await program.account.stakingRecord.fetch(
      stakingRecord
    );
    assert(stakingRecordData.stakePool.equals(stakePool));
    assert(stakingRecordData.recordOwner.equals(userWalletKeypair.publicKey));
    assert(stakingRecordData.stakedAmount.isZero());
    assert(stakingRecordData.rewardUnits.isZero());

    // Verify that lock up should not be in the future.
    const currentTime = new Date().getTime() / 1000;
    assert(stakingRecordData.lockUpExpiry.toNumber() <= currentTime);

    const expectedlastEpochClaimedVec = defaultVec();

    // Expect that first RewardPool should be claimed for epoch 3, the last active epoch.
    expectedlastEpochClaimedVec[0] = 3;
    // Expect that second RewardPool should be claimed for epoch 3, the current epoch.
    expectedlastEpochClaimedVec[1] = 3;

    assert.deepEqual(
      stakingRecordData.lastEpochClaimedVec,
      expectedlastEpochClaimedVec
    );
    assert.equal(stakingRecordData.recordBump, stakingRecordBump);
  });

  it("Creates a StakingRecord when there are three RewardPools.", async () => {
    // Wait till epoch 3.
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      3
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Set first RewardPool as inactive.
    await updateRewadPoolHelper(false, 1);

    // Add another RewardPool, starting in epoch 5.
    await createRewardPoolHelper(5, 2);

    // Wait till epoch 5.
    isTargetEpochReached = await waitForStakePoolEpoch(program, stakePool, 5);
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Add another RewardPool, starting in epoch 7.
    await createRewardPoolHelper(7, 3);

    // Create a StakingRecord.
    [stakingRecord, stakingRecordBump] = await deriveStakingRecord(
      program,
      userWalletKeypair.publicKey,
      stakePool
    );
    await createStakingRecordHelper();

    const stakingRecordData = await program.account.stakingRecord.fetch(
      stakingRecord
    );
    assert(stakingRecordData.stakePool.equals(stakePool));
    assert(stakingRecordData.recordOwner.equals(userWalletKeypair.publicKey));
    assert(stakingRecordData.stakedAmount.isZero());
    assert(stakingRecordData.rewardUnits.isZero());

    // Verify that lock up should not be in the future.
    const currentTime = new Date().getTime() / 1000;
    assert(stakingRecordData.lockUpExpiry.toNumber() <= currentTime);

    const expectedlastEpochClaimedVec = defaultVec();

    // Expect that first RewardPool should be claimed for epoch 3, the last active epoch.
    expectedlastEpochClaimedVec[0] = 3;
    // Expect that second RewardPool should be claimed for epoch 5, the current epoch.
    expectedlastEpochClaimedVec[1] = 5;
    // Expect that third RewardPool should be claimed for epoch 5, the current epoch.
    expectedlastEpochClaimedVec[2] = 5;

    assert.deepEqual(
      stakingRecordData.lastEpochClaimedVec,
      expectedlastEpochClaimedVec
    );
    assert.equal(stakingRecordData.recordBump, stakingRecordBump);
  });
});