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

describe("Test RewardRecord create 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,
    rewardRecord7: PublicKey,
    rewardRecord8: PublicKey,
    stakingRecord;

  const startingEpoch = 1;
  let poolId = 1;
  const epochDuration = new BN(MIN_EPOCH_DURATION);
  const stakeAmount = new BN(100);
  const distributionType = { constant: {} };
  const rewardPerEpoch = new BN(1000);
  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 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);
    [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 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("Create RewardRecord for a future epoch fails.", async () => {
    // Try to initialize a RewardRecord for epoch 2.
    try {
      await createRewardRecordHelper(rewardRecord2, rewardRecord1);
      assert(false);
    } catch (err) {
      assertParsedError(
        err,
        "RewardRecord cannot be created for future epochs."
      );
    }
  });

  it("Create RewardRecord in epoch 2 succeeds with no tokens staked.", async () => {
    const isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      2
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

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

    const rewardRecordData1 = await program.account.rewardRecord.fetch(
      rewardRecord1
    );
    const rewardRecordData2 = await program.account.rewardRecord.fetch(
      rewardRecord2
    );
    const [_, expectedBump] = await deriveRewardRecord(program, rewardPool, 2);
    assert(rewardRecordData2.stakePool.equals(stakePool));
    assert(rewardRecordData2.rewardPool.equals(rewardPool));
    assert.equal(rewardRecordData2.epoch, 2);
    assert.equal(rewardRecordData2.recordBump, expectedBump);
    assert(
      rewardRecordData2.epochEndTimestamp.eq(
        rewardRecordData1.epochEndTimestamp.add(epochDuration)
      )
    );
    assert(rewardRecordData2.cummulativeReward.isZero());
    assert(rewardRecordData2.epochReward.isZero());
    assert(rewardRecordData2.totalRewardAllocated.isZero());
    assert(rewardRecordData2.isUpdatable);
  });

  it("Create RewardRecord succeeds with tokens staked, but empty reward token account.", async () => {
    // Stake token into a pool, to increment reward units.
    const lockUpPeriod = 0;
    // Claim rewards for only epoch 1, since 2 has not ended.
    await claimRewardHelper(rewardRecord1, rewardRecord1);
    await stakeTokenHelper(stakeAmount, lockUpPeriod);

    const 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 rewardRecordData3 = await program.account.rewardRecord.fetch(
      rewardRecord3
    );
    const rewardRecordData2 = await program.account.rewardRecord.fetch(
      rewardRecord2
    );
    const [_, expectedBump] = await deriveRewardRecord(program, rewardPool, 3);
    assert(rewardRecordData3.stakePool.equals(stakePool));
    assert(rewardRecordData3.rewardPool.equals(rewardPool));
    assert.equal(rewardRecordData3.epoch, 3);
    assert.equal(rewardRecordData3.recordBump, expectedBump);
    assert(
      rewardRecordData3.epochEndTimestamp.eq(
        rewardRecordData2.epochEndTimestamp.add(epochDuration)
      )
    );
    assert(rewardRecordData3.isUpdatable);

    // Still 0 reward since no rewards in reward token pool.
    assert(rewardRecordData3.cummulativeReward.isZero());
    assert(rewardRecordData3.epochReward.isZero());
    assert(rewardRecordData3.totalRewardAllocated.isZero());
  });

  it("Create RewardRecord succeeds when tokens are available for allocation.", async () => {
    // Mint tokens to reward account.
    await network.btcToken.mintTo(
      rewardTokenAccount,
      mintAuthority,
      [],
      10 ** 6
    );

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

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

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

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

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

    // Verify that RewardRecord for epoch 4 allocates reward tokens.
    assert(rewardRecordData4.totalRewardAllocated.eq(expectedRewardAllocated));
    assert(rewardRecordData4.epochReward.eq(expectedEpochReward));
    assert(rewardRecordData4.cummulativeReward.eq(expectedEpochReward));

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

    // Verify remaining fields.
    assert(rewardRecordData4.stakePool.equals(stakePool));
    assert(rewardRecordData4.rewardPool.equals(rewardPool));
    assert.equal(rewardRecordData4.epoch, 4);
    assert(
      rewardRecordData4.epochEndTimestamp.eq(
        rewardRecordData3.epochEndTimestamp.add(epochDuration)
      )
    );
    assert(rewardRecordData4.isUpdatable);
  });

  it("Create RewardRecord adds to cummulative reward when tokens are allocated.", async () => {
    // Wait for Epoch 5.
    const isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      5
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

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

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

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

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

    // Verify that RewardRecord for epoch 5 allocates reward tokens.
    assert(rewardRecordData5.totalRewardAllocated.eq(expectedRewardAllocated));
    assert(rewardRecordData5.epochReward.eq(expectedEpochReward));

    // Verify that epoch's reward is added to the cummulative reward for epoch 4.
    assert(
      rewardRecordData5.cummulativeReward
        .sub(rewardRecordData4.cummulativeReward)
        .eq(expectedEpochReward)
    );

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

    // Verify remaining fields.
    assert(rewardRecordData5.stakePool.equals(stakePool));
    assert(rewardRecordData5.rewardPool.equals(rewardPool));
    assert.equal(rewardRecordData5.epoch, 5);
    assert(
      rewardRecordData5.epochEndTimestamp.eq(
        rewardRecordData4.epochEndTimestamp.add(epochDuration)
      )
    );
    assert(rewardRecordData5.isUpdatable);
  });

  it("CreateRewardRecord allocates rewards even when no tokens are staked.", async () => {
    // Claim rewards for only epoch 3 to 4, since 5 has not ended,
    // and last stake happened in epoch 2.
    await claimRewardHelper(rewardRecord2, rewardRecord4);

    // Unstake all tokens.
    await unstakeTokenHelper(stakeAmount);

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

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

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

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

    const expectedRewardAllocated =
      accountDataAfter.rewardPool.constantRewardPerEpoch;

    // Expect 0 epoch reward since there are no reward units.
    const expectedEpochReward = new BN(0);

    // Verify that RewardRecord for epoch 6 allocates reward tokens.
    assert(rewardRecordData6.totalRewardAllocated.eq(expectedRewardAllocated));
    assert(rewardRecordData6.epochReward.eq(expectedEpochReward));

    // Verify that epoch's reward is added to the cummulative reward for epoch 5.
    assert(
      rewardRecordData6.cummulativeReward
        .sub(rewardRecordData5.cummulativeReward)
        .eq(expectedEpochReward)
    );

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

    // Verify remaining fields.
    assert(rewardRecordData6.stakePool.equals(stakePool));
    assert(rewardRecordData6.rewardPool.equals(rewardPool));
    assert.equal(rewardRecordData6.epoch, 6);
    assert(
      rewardRecordData6.epochEndTimestamp.eq(
        rewardRecordData5.epochEndTimestamp.add(epochDuration)
      )
    );
    assert(rewardRecordData6.isUpdatable);
  });

  it("CreateRewardRecord succeeds when pool is inactive, during the last active epoch.", async () => {
    // Claim reward for epoch 5.
    await claimRewardHelper(rewardRecord4, rewardRecord5);

    // Stake tokens back into the pool, to increment reward units.
    const lockUpPeriod = 0;
    await stakeTokenHelper(stakeAmount, lockUpPeriod);

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

    // Set the RewardPool to inactive.
    await program.methods
      .updateRewardPool(distributionType, rewardPerEpoch, false)
      .accounts({
        rewardPoolAuthority: rewardPoolAuthority.publicKey,
        stakePool,
        rewardPool,
      })
      .signers([rewardPoolAuthority])
      .rpc();

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

    // Try to initialize a RewardRecord for epoch 7.
    await createRewardRecordHelper(rewardRecord7, rewardRecord6);

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

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

    // Verify that RewardRecord for epoch 7 allocates reward tokens.
    assert(rewardRecordData7.totalRewardAllocated.eq(expectedRewardAllocated));
    assert(rewardRecordData7.epochReward.eq(expectedEpochReward));

    // Verify that epoch's reward is added to the cummulative reward for epoch 7.
    assert(
      rewardRecordData7.cummulativeReward
        .sub(rewardRecordData6.cummulativeReward)
        .eq(expectedEpochReward)
    );

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

    // Verify remaining fields.
    assert(rewardRecordData7.stakePool.equals(stakePool));
    assert(rewardRecordData7.rewardPool.equals(rewardPool));
    assert.equal(rewardRecordData7.epoch, 7);
    assert(
      rewardRecordData7.epochEndTimestamp.eq(
        rewardRecordData6.epochEndTimestamp.add(epochDuration)
      )
    );
    assert(rewardRecordData7.isUpdatable);
  });

  it("CreateRewardRecord fails when pool is inactive, after the last active epoch.", async () => {
    // Wait for Epoch 8.
    const isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      8
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    try {
      // Try to initialize a RewardRecord for epoch 8.
      await createRewardRecordHelper(rewardRecord8, rewardRecord7);
      assert(false);
    } catch (err) {
      assertParsedError(err, "RewardPool is no longer active.");
    }
  });
});