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";
import { sleep } from "@project-serum/common";
const { setupTestAccounts, MIN_EPOCH_DURATION } = require("./setup");
const {
  getAccountData,
  waitForStakePoolEpoch,
  assertParsedError,
} = require("./utils");
import { updateRewardRecordRpc } from "../packages/psystake/src/instructions";

// Warning: this test case may be flaky due to the use of fixed time durations to
// test time-based constraints on the RewardRecord update and create instructions.
describe("Test RewardRecord update 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,
    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,
    rewardRecord5: PublicKey,
    rewardRecord6: PublicKey,
    stakingRecord;

  const startingEpoch = 1;
  let poolId = 1;
  // Make epoch duration longer to test multiple update calls in the same epoch.
  const epochDuration = new BN(MIN_EPOCH_DURATION + 7);
  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 claimRewardHelper = async (startRewardRecord, endRewardRecord) => {
    await program.methods
      .claimReward()
      .accounts({
        recordOwner: userWalletKeypair.publicKey,
        stakePool,
        stakingRecord,
        rewardPool,
        startRewardRecord,
        endRewardRecord,
        ownerTokenAccount: userBtcTokenAccount,
        allocatedTokenAccount,
      })
      .signers([userWalletKeypair])
      .rpc();
  };

  const updateRewardRecordHelper = async (
    epochToUpdate: number,
    isStartingEpoch = false
  ) => {
    await updateRewardRecordRpc(
      program,
      stakePool,
      poolId,
      epochToUpdate,
      isStartingEpoch,
      []
    );
  };

  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 distributionType = { constant: {} };
    const rewardPerEpoch = new BN(1000);

    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);
    [rewardRecord5] = await deriveRewardRecord(program, rewardPool, 5);
    [rewardRecord6] = await deriveRewardRecord(program, rewardPool, 6);

    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();
  });

  it("Update RewardRecord in epoch 1 before RewardPool starts with no tokens staked successfully.", async () => {
    const accountDataBefore = await getAccountData(program, {
      rewardRecord: rewardRecord1,
      rewardTokenAccount,
      allocatedTokenAccount,
    });
    const rewardRecordBefore = accountDataBefore.rewardRecord;

    const success = await sleep(2000).then(async () => {
      updateRewardRecordHelper(1, true);
      return true;
    });

    assert(success);
    if (success) {
      const accountDataAfter = await getAccountData(program, {
        rewardRecord: rewardRecord1,
        rewardTokenAccount,
        allocatedTokenAccount,
      });
      const rewardRecordAfter = accountDataAfter.rewardRecord;
      assert(rewardRecordAfter.stakePool.equals(rewardRecordBefore.stakePool));
      assert(
        rewardRecordAfter.rewardPool.equals(rewardRecordBefore.rewardPool)
      );
      assert.equal(rewardRecordAfter.epoch, rewardRecordBefore.epoch);
      assert.equal(rewardRecordAfter.recordBump, rewardRecordBefore.recordBump);
      assert(
        rewardRecordAfter.epochEndTimestamp.eq(
          rewardRecordBefore.epochEndTimestamp
        )
      );
      assert(rewardRecordAfter.cummulativeReward.isZero());
      assert(rewardRecordAfter.epochReward.isZero());
      assert(rewardRecordAfter.totalRewardAllocated.isZero());
      assert(rewardRecordAfter.isUpdatable);
    }
  });

  it("Update RewardRecord in epoch 2 with no tokens staked successfully.", async () => {
    const 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, {
      rewardRecord: rewardRecord2,
      rewardTokenAccount,
      allocatedTokenAccount,
    });
    const rewardRecordBefore = accountDataBefore.rewardRecord;

    const isRewardRecordUpdated = await sleep(2000).then(async () => {
      await updateRewardRecordHelper(2);
      return true;
    });

    assert(isRewardRecordUpdated);
    const accountDataAfter = await getAccountData(program, {
      rewardRecord: rewardRecord2,
      rewardTokenAccount,
      allocatedTokenAccount,
    });
    const rewardRecordAfter = accountDataAfter.rewardRecord;
    assert(rewardRecordAfter.stakePool.equals(rewardRecordBefore.stakePool));
    assert(rewardRecordAfter.rewardPool.equals(rewardRecordBefore.rewardPool));
    assert.equal(rewardRecordAfter.epoch, rewardRecordBefore.epoch);
    assert.equal(rewardRecordAfter.recordBump, rewardRecordBefore.recordBump);
    assert(
      rewardRecordAfter.epochEndTimestamp.eq(
        rewardRecordBefore.epochEndTimestamp
      )
    );
    assert(rewardRecordAfter.cummulativeReward.isZero());
    assert(rewardRecordAfter.epochReward.isZero());
    assert(rewardRecordAfter.totalRewardAllocated.isZero());
    assert(rewardRecordAfter.isUpdatable);
  });

  it("Update RewardRecord in epoch 2 does not allocate reward when insufficient reward tokens available.", async () => {
    // Claim rewards for only epoch 1, since 2 has not ended.
    await claimRewardHelper(rewardRecord1, rewardRecord1);

    // Stake token into a pool, to increment reward units.
    const stakeAmount = new BN(100);
    const lockUpPeriod = 0;
    await stakeTokenHelper(stakeAmount, lockUpPeriod);

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

    await updateRewardRecordHelper(2);

    const accountDataAfter = await getAccountData(program, {
      rewardRecord: rewardRecord2,
      rewardTokenAccount,
      allocatedTokenAccount,
    });
    const rewardRecordAfter = accountDataAfter.rewardRecord;
    assert(rewardRecordAfter.stakePool.equals(rewardRecordBefore.stakePool));
    assert(rewardRecordAfter.rewardPool.equals(rewardRecordBefore.rewardPool));
    assert.equal(rewardRecordAfter.epoch, rewardRecordBefore.epoch);
    assert.equal(rewardRecordAfter.recordBump, rewardRecordBefore.recordBump);
    assert(
      rewardRecordAfter.epochEndTimestamp.eq(
        rewardRecordBefore.epochEndTimestamp
      )
    );
    assert(rewardRecordAfter.cummulativeReward.isZero());
    assert(rewardRecordAfter.epochReward.isZero());
    assert(rewardRecordAfter.totalRewardAllocated.isZero());
    assert(rewardRecordAfter.isUpdatable);
  });

  it("Update RewardRecord succeeds even when called after epoch ends.", async () => {
    // Mint tokens to reward account.
    await network.btcToken.mintTo(
      rewardTokenAccount,
      mintAuthority,
      [],
      10 ** 6
    );

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

    // Try to initialize a RewardRecord for epoch 3.
    await createRewardRecordHelper(rewardRecord3, rewardRecord2);

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

    // Get time difference to epoch end for Epoch 2.
    const lastUpdatedTimeDIff =
      rewardRecordData2.epochEndTimestamp.toNumber() -
      rewardRecordData2.lastUpdated.toNumber();
    const currentTimeDiff =
      new Date().getTime() / 1000 -
      rewardRecordData2.epochEndTimestamp.toNumber();

    // Verify that current time diff from epoch end is less than last updated.
    assert(currentTimeDiff < lastUpdatedTimeDIff);

    // Update RewardRecord for epoch 2, during epoch 3.
    await updateRewardRecordHelper(2);

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

    const expectedRewardAllocated =
      accountDataAfter.rewardPool.constantRewardPerEpoch;
    const expectedEpochReward = expectedRewardAllocated
      .mul(new BN(10 ** epochRewardDecimals))
      .div(stakePoolAfter.totalRewardUnits);

    assert(rewardRecordAfter.stakePool.equals(rewardRecordData2.stakePool));
    assert(rewardRecordAfter.rewardPool.equals(rewardRecordData2.rewardPool));
    assert.equal(rewardRecordAfter.epoch, rewardRecordData2.epoch);
    assert.equal(rewardRecordAfter.recordBump, rewardRecordData2.recordBump);
    assert.equal(
      rewardRecordAfter.epochEndTimestamp.toNumber(),
      rewardRecordData2.epochEndTimestamp.toNumber()
    );
    assert(rewardRecordAfter.isUpdatable);

    // Verify the tokens are allocated now that there are sufficient reward tokens available.
    assert(rewardRecordAfter.totalRewardAllocated.eq(expectedRewardAllocated));
    assert(rewardRecordAfter.cummulativeReward.eq(expectedEpochReward));
    assert(rewardRecordAfter.epochReward.eq(expectedEpochReward));

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

  it("Update RewardRecord recalculates epoch reward when total reward units changes.", async () => {
    // Stake token into a pool, to increment reward units.
    const stakeAmount = new BN(100);
    const lockUpPeriod = 0;
    await stakeTokenHelper(stakeAmount, lockUpPeriod);

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

    const isRewardRecordUpdated = await sleep(0).then(async () => {
      await updateRewardRecordHelper(3);
      return true;
    });

    assert(isRewardRecordUpdated);

    const accountDataAfter = await getAccountData(program, {
      rewardRecord: rewardRecord3,
      rewardTokenAccount,
      allocatedTokenAccount,
      rewardPool,
      stakePool,
    });
    const rewardRecordAfter = accountDataAfter.rewardRecord;
    const stakePoolAfter = accountDataAfter.stakePool;

    const expectedRewardAllocated =
      accountDataAfter.rewardPool.constantRewardPerEpoch;
    const expectedEpochReward = expectedRewardAllocated
      .mul(new BN(10 ** epochRewardDecimals))
      .div(stakePoolAfter.totalRewardUnits);

    const cummulatativeRewardForEpoch2 = await (
      await program.account.rewardRecord.fetch(rewardRecord2)
    ).cummulativeReward;
    const expectedCummulativeRewardForEpoch3 =
      cummulatativeRewardForEpoch2.add(expectedEpochReward);

    // Verify that only cummulatativeReward, epochReward and lastUpdated should change.
    assert(rewardRecordAfter.stakePool.equals(rewardRecordBefore.stakePool));
    assert(rewardRecordAfter.rewardPool.equals(rewardRecordBefore.rewardPool));
    assert.equal(rewardRecordAfter.epoch, rewardRecordBefore.epoch);
    assert.equal(rewardRecordAfter.recordBump, rewardRecordBefore.recordBump);
    assert.equal(
      rewardRecordAfter.epochEndTimestamp.toNumber(),
      rewardRecordBefore.epochEndTimestamp.toNumber()
    );
    assert(
      rewardRecordAfter.totalRewardAllocated.eq(
        rewardRecordBefore.totalRewardAllocated
      )
    );
    assert(
      rewardRecordAfter.cummulativeReward.eq(expectedCummulativeRewardForEpoch3)
    );
    assert(rewardRecordAfter.epochReward.eq(expectedEpochReward));
    assert(rewardRecordAfter.isUpdatable);

    // Verify amounts in token accounts unchanged.
    assert(
      accountDataBefore.rewardTokenAccountAmount
        .sub(accountDataAfter.rewardTokenAccountAmount)
        .isZero()
    );
    assert(
      accountDataAfter.allocatedTokenAccountAmount
        .sub(accountDataBefore.allocatedTokenAccountAmount)
        .isZero()
    );
  });

  it("Update RewardRecord fails when called after a longer period from epoch end than last updated.", async () => {
    const accountDataBefore = await getAccountData(program, {
      rewardRecord: rewardRecord2,
    });
    const rewardRecordBefore = accountDataBefore.rewardRecord;
    const lastUpdatedTimeDIff =
      rewardRecordBefore.epochEndTimestamp.toNumber() -
      rewardRecordBefore.lastUpdated.toNumber();
    const currentTimeDiff =
      new Date().getTime() / 1000 -
      rewardRecordBefore.epochEndTimestamp.toNumber();

    // Verify that current time diff from epoch end is more than last updated,
    // since this test is invoked after last update was called after epoch ended.
    assert(Math.abs(currentTimeDiff) > Math.abs(lastUpdatedTimeDIff));

    try {
      // Update RewardRecord for epoch 2, during epoch 3.
      await updateRewardRecordHelper(2);
      assert(false);
    } catch (err) {
      assertParsedError(
        err,
        "RewardRecord cannot be updated as last update was more recent to epoch end."
      );
    }
  });

  it("Update RewardRecord fails when the prev record is not from previous epoch.", async () => {
    // Wait for Epoch 4.
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      4
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Try to initialize a RewardRecord for epoch 4.
    await createRewardRecordHelper(rewardRecord4, rewardRecord3);

    try {
      // Update RewardRecord for epoch 4, setting previous reward record as epoch 2.
      await program.methods
        .updateRewardRecord()
        .accounts({
          stakePool,
          rewardPool,
          curRewardRecord: rewardRecord4,
          prevRewardRecord: rewardRecord2,
          rewardTokenAccount,
          allocatedTokenAccount,
        })
        .rpc();
      assert(false);
    } catch (err) {
      assertParsedError(
        err,
        "Previous record must be from one epoch before current record."
      );
    }
  });

  it("Update RewardRecord fails when fails when more than an epoch has passed, even if unclaimed.", async () => {
    // Wait for Epoch 5.
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      5
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Try to initialize a RewardRecord for epoch 5.
    await createRewardRecordHelper(rewardRecord5, rewardRecord4);

    const rewardRecordData3 = await program.account.rewardRecord.fetch(
      rewardRecord3
    );
    assert(rewardRecordData3.isUpdatable);

    // Update RewardRecord for epoch 3. This should fail as more than an epoch has passed,
    // even though reward was not claimed (isUpdatable is true).
    try {
      await updateRewardRecordHelper(3);
      assert(false);
    } catch (err) {
      assertParsedError(err, "RewardRecord is not updatable.");
    }
  });

  it("Update RewardRecord should fail if reward was claimed for the epoch.", async () => {
    // Claim rewards for epoch 4.
    await claimRewardHelper(rewardRecord3, rewardRecord4);

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

    // Get time difference to epoch end for Epoch 4.
    const lastUpdatedTimeDIff =
      rewardRecordData4.epochEndTimestamp.toNumber() -
      rewardRecordData4.lastUpdated.toNumber();
    const currentTimeDiff =
      new Date().getTime() / 1000 -
      rewardRecordData4.epochEndTimestamp.toNumber();
    console.log(currentTimeDiff, lastUpdatedTimeDIff);

    // Verify that current time diff from epoch end is less than last updated.
    assert(currentTimeDiff < lastUpdatedTimeDIff);

    // Check that RewardRecord should not be updatable, after reward was claimed.
    assert(!rewardRecordData4.isUpdatable);

    // Update RewardRecord for epoch 4.

    try {
      // Update RewardRecord for epoch 4 should fail since it is not updatable, even if
      // time from last updated is longer than current time diff.
      await updateRewardRecordHelper(4);
      assert(false);
    } catch (err) {
      assertParsedError(err, "RewardRecord is not updatable.");
    }
  });
});