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,
  defaultVec,
  assertParsedError,
} = require("./utils");

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

  let userUsdcTokenAccount: PublicKey,
    userBtcTokenAccount: PublicKey,
    network,
    userWalletKeypair: Keypair,
    userWalletKeypair2: Keypair,
    userUsdcTokenAccount2: PublicKey,
    userBtcTokenAccount2: PublicKey,
    stakePoolAuthority: Keypair,
    rewardPoolAuthority: Keypair,
    mintAuthority: 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,
    rewardRecord5: PublicKey,
    rewardRecord6: PublicKey,
    rewardRecord7: PublicKey,
    rewardRecord8: PublicKey,
    stakingRecord: PublicKey,
    stakingRecord2: PublicKey;

  const epochDuration = new BN(MIN_EPOCH_DURATION);
  const distributionType = { constant: {} };
  const rewardPerEpoch = new BN(1000);
  const epochRewardDecimals = 8;

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

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

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

  const claimRewardHelper = async (
    startRewardRecord: PublicKey,
    endRewardRecord: PublicKey,
    stakingRecord: PublicKey,
    ownerTokenAccount: PublicKey,
    keypair: Keypair,
    selectedRewardPool = rewardPool,
    selectedAllocatedTokenAccount = allocatedTokenAccount
  ) => {
    await program.methods
      .claimReward()
      .accounts({
        recordOwner: keypair.publicKey,
        stakePool,
        stakingRecord,
        rewardPool: selectedRewardPool,
        startRewardRecord,
        endRewardRecord,
        ownerTokenAccount,
        allocatedTokenAccount: selectedAllocatedTokenAccount,
      })
      .signers(keypair)
      .rpc();
  };

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

  const createRewardPoolHelper = async (startingEpoch, poolId) => {
    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();
  };

  before(async () => {
    ({
      userWalletKeypair,
      userUsdcTokenAccount,
      userBtcTokenAccount,
      userWalletKeypair2,
      userUsdcTokenAccount2,
      userBtcTokenAccount2,
      network,
      stakePoolAuthority,
      rewardPoolAuthority,
      mintAuthority,
    } = 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=*/ 1,
      /*startingEpoch=*/ 1
    );
    ({ 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);
    [rewardRecord5] = await deriveRewardRecord(program, rewardPool, 5);
    [rewardRecord6] = await deriveRewardRecord(program, rewardPool, 6);
    [rewardRecord7] = await deriveRewardRecord(program, rewardPool, 7);
    [rewardRecord8] = await deriveRewardRecord(program, rewardPool, 8);

    await createRewardPoolHelper(/*startingEpoch=*/ 1, /*poolId=*/ 1);
    await createRewardPoolHelper(/*startingEpoch=*/ 3, /*poolId=*/ 2);

    // Create a StakingRecord for user 1 and 2.
    [stakingRecord] = await deriveStakingRecord(
      program,
      userWalletKeypair.publicKey,
      stakePool
    );
    [stakingRecord2] = await deriveStakingRecord(
      program,
      userWalletKeypair2.publicKey,
      stakePool
    );
    await createStakingRecordHelper(userWalletKeypair, stakingRecord);
    await createStakingRecordHelper(userWalletKeypair2, stakingRecord2);

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

  it("Claim reward fails if claiming reward from a RewardPool that has not started.", async () => {
    // Stake tokens in epoch 0 for user 1
    const lockUpPeriod = 0;
    await stakeTokenHelper(
      new BN(2000),
      lockUpPeriod,
      userUsdcTokenAccount,
      stakingRecord,
      userWalletKeypair
    );

    // Stake tokens in epoch 0 for user 2
    await stakeTokenHelper(
      new BN(3000),
      lockUpPeriod,
      userUsdcTokenAccount2,
      stakingRecord2,
      userWalletKeypair2
    );

    try {
      // Claim reward for this RewardPool, for user 1.
      await claimRewardHelper(
        rewardRecord1,
        rewardRecord1,
        stakingRecord,
        userBtcTokenAccount,
        userWalletKeypair
      );
      assert(false);
    } catch (err) {
      // Expect failure as the RewardPool has not started.
      assertParsedError(err, "No rewards available for claiming.");
    }
  });

  it("Claim reward for an epoch fails if epoch has not ended.", async () => {
    // Wait for epoch 1.
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      1
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Update RewardRecord for epoch 1.
    await updateRewardRecordHelper(rewardRecord1, rewardRecord1);

    try {
      // Claim reward for the first epoch, for user 1.
      await claimRewardHelper(
        rewardRecord1,
        rewardRecord1,
        stakingRecord,
        userBtcTokenAccount,
        userWalletKeypair
      );
      assert(false);
    } catch (err) {
      // Expect failure as the epoch has not ended.
      assertParsedError(err, "The current epoch has not ended.");
    }
  });

  it("Claim reward succeeds after epoch 1 has ended.", async () => {
    // Wait for epoch 2
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      2
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

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

    const accountDataBefore = await getAccountData(program, {
      stakingRecord,
      rewardRecord: rewardRecord1,
      rewardPool,
      allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount,
    });

    // Claim reward for the first epoch, for user 1.
    await claimRewardHelper(
      rewardRecord1,
      rewardRecord1,
      stakingRecord,
      userBtcTokenAccount,
      userWalletKeypair
    );

    const accountDataAfter = await getAccountData(program, {
      stakingRecord,
      rewardRecord: rewardRecord1,
      allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount,
    });

    const expectedReward = accountDataBefore.rewardRecord.cummulativeReward
      .mul(accountDataBefore.stakingRecord.rewardUnits)
      .div(new BN(10 ** accountDataBefore.rewardPool.epochRewardDecimals));

    // Check transfer of expected reward.
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(expectedReward)
    );
    assert(
      accountDataBefore.allocatedTokenAccountAmount
        .sub(accountDataAfter.allocatedTokenAccountAmount)
        .eq(expectedReward)
    );

    // Check that RewardRecord is now not updatable.
    assert(!accountDataAfter.rewardRecord.isUpdatable);

    // Verify last epoch claimed has been updated to include claimed epoch.
    const expectedLastClaimedVec = defaultVec();
    // First epoch should be claimed, for first RewardPool.
    expectedLastClaimedVec[0] = 1;
    assert.deepEqual(
      accountDataAfter.stakingRecord.lastEpochClaimedVec,
      expectedLastClaimedVec
    );
  });

  it("Claim reward fails when epoch has already been claimed.", async () => {
    try {
      // Claim reward for the first epoch, for user 1, again.
      await claimRewardHelper(
        rewardRecord1,
        rewardRecord1,
        stakingRecord,
        userBtcTokenAccount,
        userWalletKeypair
      );
      assert(false);
    } catch (err) {
      assertParsedError(err, "Reward has already been claimed.");
    }
  });

  it("Claim reward fails when starting RewardRecord is not from last epoch claimed.", async () => {
    // Wait for epoch 3
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      3
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Initialize a RewardRecord for epoch 3.
    await createRewardRecordHelper(rewardRecord3, rewardRecord2);

    try {
      await claimRewardHelper(
        rewardRecord2,
        rewardRecord2,
        stakingRecord,
        userBtcTokenAccount,
        userWalletKeypair
      );
      assert(false);
    } catch (err) {
      assertParsedError(
        err,
        "Start RewardRecord should be from the last epoch claimed."
      );
    }
  });

  it("Claim reward fails if starting RewardRecord is newer than end RewardRecord claimed.", async () => {
    // Claim Reward for epoch 2, for user 1.
    await claimRewardHelper(
      rewardRecord1,
      rewardRecord2,
      stakingRecord,
      userBtcTokenAccount,
      userWalletKeypair
    );

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

    // Initialize a RewardRecord for epoch 4.
    await createRewardRecordHelper(rewardRecord4, rewardRecord3);

    try {
      await claimRewardHelper(
        rewardRecord2,
        rewardRecord1,
        stakingRecord,
        userBtcTokenAccount,
        userWalletKeypair
      );
      assert(false);
    } catch (err) {
      assertParsedError(
        err,
        "End RewardRecord must be from the same or later epoch than start RewardRecord."
      );
    }
  });

  it("Claim reward succeeds for first claim on RewardPool after more than one epoch.", async () => {
    const accountDataBefore = await getAccountData(program, {
      stakingRecord: stakingRecord2,
      rewardRecord: rewardRecord3,
      rewardPool,
      allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount2,
    });

    // Claim Reward from epoch 1 to 3, for user 2.
    await claimRewardHelper(
      rewardRecord2, // Starting reward record is not read, for first claims.
      rewardRecord3,
      stakingRecord2,
      userBtcTokenAccount2,
      userWalletKeypair2
    );
    const accountDataAfter = await getAccountData(program, {
      stakingRecord: stakingRecord2,
      rewardRecord: rewardRecord3,
      allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount2,
    });

    // Expect all cummulative reward uptil epoch 3 to be claimed.
    const expectedReward = accountDataBefore.rewardRecord.cummulativeReward
      .mul(accountDataBefore.stakingRecord.rewardUnits)
      .div(new BN(10 ** accountDataBefore.rewardPool.epochRewardDecimals));

    // Check transfer of expected reward.
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(expectedReward)
    );
    assert(
      accountDataBefore.allocatedTokenAccountAmount
        .sub(accountDataAfter.allocatedTokenAccountAmount)
        .eq(expectedReward)
    );

    // Check that RewardRecord is now not updatable.
    assert(!accountDataAfter.rewardRecord.isUpdatable);

    // Verify last epoch claimed has been updated to include claimed epoch.
    const expectedLastClaimedVec = defaultVec();
    // Third epoch should be claimed.
    expectedLastClaimedVec[0] = 3;
    assert.deepEqual(
      accountDataAfter.stakingRecord.lastEpochClaimedVec,
      expectedLastClaimedVec
    );
  });

  it("Claim reward succeeds for subset of epochs, not including the most recent completed epoch.", async () => {
    // Wait for epoch 6
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      6
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Initialize a RewardRecord for epoch 5 and 6.
    await createRewardRecordHelper(rewardRecord5, rewardRecord4);
    await createRewardRecordHelper(rewardRecord6, rewardRecord5);

    const accountDataBefore = await getAccountData(program, {
      stakingRecord,
      rewardRecord: rewardRecord4,
      rewardPool,
      allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount,
    });
    const rewardRecordData2 = await program.account.rewardRecord.fetch(
      rewardRecord2
    );

    // Claim Reward from epoch 3 to 4, for user 1.
    await claimRewardHelper(
      rewardRecord2,
      rewardRecord4,
      stakingRecord,
      userBtcTokenAccount,
      userWalletKeypair
    );
    const accountDataAfter = await getAccountData(program, {
      stakingRecord,
      rewardRecord: rewardRecord4,
      allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount,
    });

    // Expect cummulative reward inc. epoch 3 and 4 to be claimed.
    const expectedReward = accountDataBefore.rewardRecord.cummulativeReward
      .sub(rewardRecordData2.cummulativeReward)
      .mul(accountDataBefore.stakingRecord.rewardUnits)
      .div(new BN(10 ** accountDataBefore.rewardPool.epochRewardDecimals));

    // Check transfer of expected reward.
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(expectedReward)
    );
    assert(
      accountDataBefore.allocatedTokenAccountAmount
        .sub(accountDataAfter.allocatedTokenAccountAmount)
        .eq(expectedReward)
    );

    // Check that RewardRecord is now not updatable.
    assert(!accountDataAfter.rewardRecord.isUpdatable);

    // Verify last epoch claimed has been updated to include claimed epoch.
    const expectedLastClaimedVec = defaultVec();
    // Fourth epoch should be claimed for first RewardPool.
    expectedLastClaimedVec[0] = 4;
    assert.deepEqual(
      accountDataAfter.stakingRecord.lastEpochClaimedVec,
      expectedLastClaimedVec
    );
  });

  it("Claim reward succeeds for second RewardPool.", async () => {
    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      /*poolId=*/ 2,
      /*startingEpoch=*/ 3
    );

    const rewardPool2 = poolAccounts.rewardPool;

    const [pool2Reward3] = await deriveRewardRecord(program, rewardPool2, 3);
    const [pool2Reward4] = await deriveRewardRecord(program, rewardPool2, 4);
    const [pool2Reward5] = await deriveRewardRecord(program, rewardPool2, 5);

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

    // Initialize RewardRecord epoch 4 and 5 for second RewardPool.
    await createRewardRecordHelper(
      pool2Reward4,
      pool2Reward3,
      rewardPool2,
      poolAccounts.rewardTokenAccount,
      poolAccounts.allocatedTokenAccount,
      userWalletKeypair
    );
    await createRewardRecordHelper(
      pool2Reward5,
      pool2Reward4,
      rewardPool2,
      poolAccounts.rewardTokenAccount,
      poolAccounts.allocatedTokenAccount,
      userWalletKeypair
    );

    const accountDataBefore = await getAccountData(program, {
      stakingRecord,
      rewardRecord: pool2Reward5,
      rewardPool: rewardPool2,
      allocatedTokenAccount: poolAccounts.allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount,
    });

    // Claim Reward for RewardPool 2 from epoch 3 to 5, for user 1.
    await claimRewardHelper(
      pool2Reward3, // Starting reward record is not read, for first claims.
      pool2Reward5,
      stakingRecord,
      userBtcTokenAccount,
      userWalletKeypair,
      rewardPool2,
      poolAccounts.allocatedTokenAccount
    );
    const accountDataAfter = await getAccountData(program, {
      stakingRecord,
      rewardRecord: pool2Reward5,
      allocatedTokenAccount: poolAccounts.allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount,
    });

    // Expect cummulative reward up to epoch 5 to be claimed.
    const expectedReward = accountDataBefore.rewardRecord.cummulativeReward
      .mul(accountDataBefore.stakingRecord.rewardUnits)
      .div(new BN(10 ** accountDataBefore.rewardPool.epochRewardDecimals));

    // Check transfer of expected reward.
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(expectedReward)
    );
    assert(
      accountDataBefore.allocatedTokenAccountAmount
        .sub(accountDataAfter.allocatedTokenAccountAmount)
        .eq(expectedReward)
    );

    // Check that RewardRecord is now not updatable.
    assert(!accountDataAfter.rewardRecord.isUpdatable);

    // Verify last epoch claimed has been updated to include claimed epoch.
    const expectedLastClaimedVec = defaultVec();
    // Fourth epoch should be claimed for first RewardPool.
    expectedLastClaimedVec[0] = 4;
    // Fifth epoch should be claimed for second RewardPool.
    expectedLastClaimedVec[1] = 5;
    assert.deepEqual(
      accountDataAfter.stakingRecord.lastEpochClaimedVec,
      expectedLastClaimedVec
    );
  });
});