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,
  deriveStakingRecord,
  deriveRewardRecord,
} from "./pdas";

import { assert } from "chai";
import { Keypair, PublicKey } from "@solana/web3.js";
import {
  createStakingRecordRpc,
  stakeTokenRpc,
  updateRewardRecordRpc,
  updateStakePoolEpochRpc,
} from "../packages/psystake/src/instructions";
import { sleep } from "@project-serum/common";
const {
  setupTestAccounts,
  MAX_REWARD_POOLS,
  MIN_EPOCH_DURATION,
} = require("./setup");
const {
  getAccountData,
  defaultVec,
  waitForStakePoolEpoch,
  assertParsedError,
} = require("./utils");

describe("Test setting up RewardPool after Staking.", () => {
  const program = anchor.workspace.PsyStake as Program<PsyStake>;
  const provider = anchor.Provider.env();
  anchor.setProvider(provider);

  let network, stakePoolAuthority: Keypair, rewardPoolAuthority: Keypair;

  const uniqueSeed = 1;

  let stakePool: PublicKey,
    stakingTokenAccount: PublicKey,
    stakingTokenMint: PublicKey,
    userUsdcTokenAccount: PublicKey,
    userBtcTokenAccount: PublicKey,
    userWalletKeypair: Keypair,
    userWalletKeypair2: Keypair,
    userUsdcTokenAccount2: PublicKey;

  let rewardTokenAccount: PublicKey,
    mintAuthority: PublicKey,
    stakingRecord: PublicKey,
    stakingRecord2: PublicKey,
    rewardRecord2: PublicKey,
    rewardRecord3: PublicKey,
    rewardRecord4: PublicKey,
    rewardPool: PublicKey,
    allocatedTokenAccount: PublicKey;

  let poolId = 1;
  const epochDuration = new BN(MIN_EPOCH_DURATION);
  const expectedStartingEpochVec = defaultVec();
  const epochRewardDecimals = 8;

  const createRewardPoolHelper = async (
    startingEpoch,
    poolId,
    distributionType,
    rewardPerEpoch,
    decimals = epochRewardDecimals
  ) => {
    // Create a RewardPool.
    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      poolId,
      startingEpoch
    );
    await program.methods
      .createRewardPool(
        startingEpoch,
        distributionType,
        rewardPerEpoch,
        decimals
      )
      .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 createRewardRecordHelper = async (
    newRewardRecord,
    prevRewardRecord
  ) => {
    await program.methods
      .createRewardRecord()
      .accounts({
        payer: userWalletKeypair.publicKey,
        stakePool,
        rewardPool,
        newRewardRecord,
        prevRewardRecord,
        rewardTokenAccount,
        allocatedTokenAccount,
      })
      .signers([userWalletKeypair])
      .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 createStakingRecordHelper = async (
    recordOwner = userWalletKeypair.publicKey,
    signer = userWalletKeypair
  ) => {
    await createStakingRecordRpc(program, recordOwner, stakePool, [signer]);
  };

  const stakeTokenHelper = async (
    stakeAmount,
    lockUpPeriod,
    ownerTokenAccount,
    recordOwner = userWalletKeypair.publicKey,
    signer = userWalletKeypair
  ) => {
    await stakeTokenRpc(
      program,
      recordOwner,
      stakePool,
      ownerTokenAccount,
      stakeAmount,
      lockUpPeriod,
      [signer]
    );
  };

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

  before(async () => {
    ({
      network,
      stakePoolAuthority,
      userUsdcTokenAccount,
      rewardPoolAuthority,
      userWalletKeypair,
      userWalletKeypair2,
      userBtcTokenAccount,
      userUsdcTokenAccount2,
      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 + 1);

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

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

    [stakingRecord2] = await deriveStakingRecord(
      program,
      userWalletKeypair2.publicKey,
      stakePool
    );
    await createStakingRecordHelper(
      userWalletKeypair2.publicKey,
      userWalletKeypair2
    );

    // Stake Tokens
    const stakeAmount = new BN(1000);
    const lockUpPeriod = 0;
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);
  });

  it("Verify StakePool and StakingRecord is set up correctly.", async () => {
    const accountData = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });

    // Verify StakePool account state
    const stakePoolData = accountData.stakePool;

    // Starting epoch should be all zero.
    assert.deepEqual(stakePoolData.startingEpochVec, expectedStartingEpochVec);
    assert.equal(stakePoolData.rewardPoolCount, 0);
    assert.equal(stakePoolData.currentEpoch, 0);
    assert.equal(stakePoolData.totalRewardUnits, 1000);

    // Verify staked amount and reward units in StakingRecord changed.
    const stakingRecordData = accountData.stakingRecord;
    assert.equal(stakingRecordData.rewardUnits, 1000);
    assert.equal(stakingRecordData.stakedAmount, 1000);
    assert.equal(accountData.stakingTokenAccountAmount, 1000);
  });

  it("Advances StakePool epoch as expected.", async () => {
    let stakePoolData = await program.account.stakePool.fetch(stakePool);
    let sleepDurationSeconds = stakePoolData.nextEpochStartTime.sub(
      new BN(new Date().getTime() / 1000)
    );
    let sleepDurationMs = sleepDurationSeconds.toNumber() * 1000 + 2000;
    console.log(`Sleep for ${sleepDurationMs}ms`);
    const success = await sleep(sleepDurationMs).then(async () => {
      await updateStakePoolEpochRpc(program, stakePool, []);
      const stakePoolData2 = await program.account.stakePool.fetch(stakePool);
      assert.equal(stakePoolData2.currentEpoch, stakePoolData.currentEpoch + 1);
      return true;
    });

    success &&
      (await sleep(MIN_EPOCH_DURATION * 1000 + 2000).then(async () => {
        await updateStakePoolEpochRpc(program, stakePool, []);
        const stakePoolData2 = await program.account.stakePool.fetch(stakePool);
        assert.equal(
          stakePoolData2.currentEpoch,
          stakePoolData.currentEpoch + 2
        );
      }));

    // Stake Tokens for user 1 and 2.
    const stakeAmount = new BN(1000);
    const lockUpPeriod = 0;
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);
    await stakeTokenHelper(
      stakeAmount,
      lockUpPeriod,
      userUsdcTokenAccount2,
      userWalletKeypair2.publicKey,
      userWalletKeypair2
    );

    // Verify StakePool account state
    const accountData = await getAccountData(program, {
      stakingRecord,
    });
    stakePoolData = accountData.stakePool;

    // Verify staked amount and reward units in StakingRecord changed for user 1.
    const stakingRecordData = accountData.stakingRecord;
    assert.equal(stakingRecordData.rewardUnits, 2000);
    assert.equal(stakingRecordData.stakedAmount, 2000);
    assert.deepEqual(stakingRecordData.lastEpochClaimedVec, defaultVec());
  });

  it("Initializes a RewardPool successfully with constant reward.", async () => {
    const distributionType = { constant: {} };
    const startingEpoch = 2;
    const rewardPerEpoch = new BN(1000000);
    const rewardToMint = 10 ** 6;

    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      poolId,
      startingEpoch
    );

    ({ rewardPool, rewardTokenAccount, allocatedTokenAccount } = poolAccounts);

    await createRewardPoolHelper(
      startingEpoch,
      poolId,
      distributionType,
      rewardPerEpoch
    );

    // Mint tokens to reward account.
    await network.btcToken.mintTo(
      rewardTokenAccount,
      mintAuthority,
      [],
      rewardToMint
    );

    const accountData = await getAccountData(program, {
      stakePool,
      ...poolAccounts,
    });

    // Verify StakePool account state
    const stakePoolData = accountData.stakePool;
    expectedStartingEpochVec[poolId - 1] = startingEpoch;
    assert.deepEqual(stakePoolData.startingEpochVec, expectedStartingEpochVec);

    // Verify RewardPool account state.
    const rewardPoolData = accountData.rewardPool;
    assert(rewardPoolData.stakePool.equals(stakePool));
    assert(
      rewardPoolData.rewardPoolAuthority.equals(rewardPoolAuthority.publicKey)
    );
    assert.equal(rewardPoolData.poolId, poolId);
    assert(rewardPoolData.rewardTokenMint.equals(network.btcToken.publicKey));
    assert(
      rewardPoolData.rewardTokenAccount.equals(poolAccounts.rewardTokenAccount)
    );
    assert(
      rewardPoolData.allocatedTokenAccount.equals(
        poolAccounts.allocatedTokenAccount
      )
    );
    assert.deepEqual(rewardPoolData.distributionType, distributionType);
    assert(rewardPoolData.constantRewardPerEpoch.eq(rewardPerEpoch));
    assert(rewardPoolData.percentageRewardMbpsPerEpoch.isZero());
    assert.equal(rewardPoolData.startingEpoch, startingEpoch);
    assert(rewardPoolData.isActive);
    assert.equal(rewardPoolData.epochRewardDecimals, epochRewardDecimals);
    assert.equal(rewardPoolData.rewardPoolBump, poolAccounts.rewardPoolBump);
    assert.equal(
      rewardPoolData.allocatedTokenAccountBump,
      poolAccounts.allocatedTokenAccountBump
    );

    // Verify RewardRecord account state
    const rewardRecordData = accountData.rewardRecord;
    assert(rewardRecordData.stakePool.equals(stakePool));
    assert(rewardRecordData.rewardPool.equals(poolAccounts.rewardPool));
    assert.equal(rewardRecordData.epoch, startingEpoch);
    assert(
      rewardRecordData.epochEndTimestamp.eq(stakePoolData.nextEpochStartTime)
    );
    assert(rewardRecordData.cummulativeReward.isZero());
    assert(rewardRecordData.epochReward.isZero());
    assert(rewardRecordData.totalRewardAllocated.isZero());
    assert.equal(rewardRecordData.recordBump, poolAccounts.rewardRecordBump);
    assert(rewardRecordData.isUpdatable);

    // Verify Token account states.
    assert(accountData.allocatedTokenAccountAmount.isZero());
    assert.equal(accountData.rewardTokenAccountAmount.toNumber(), rewardToMint);
  });

  it("Try to stake token after RewardPool is created.", async () => {
    // Stake Tokens
    const stakeAmount = new BN(1000);
    const lockUpPeriod = 0;
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);
    // Verify StakePool account state
    const accountData = await getAccountData(program, {
      stakingRecord,
    });
    // Verify staked amount and reward units in StakingRecord changed.
    const stakingRecordData = accountData.stakingRecord;
    assert.equal(stakingRecordData.rewardUnits, 3000);
    assert.equal(stakingRecordData.stakedAmount, 3000);

    // Last epoch claimed should include current epoch for RewardPool.
    assert.deepEqual(stakingRecordData.lastEpochClaimedVec, [2, 0, 0]);
  });

  it("Update StakePoolEpoch to epoch 4.", async () => {
    let stakePoolData = await program.account.stakePool.fetch(stakePool);
    let sleepDurationSeconds = stakePoolData.nextEpochStartTime.sub(
      new BN(new Date().getTime() / 1000)
    );
    let sleepDurationMs = sleepDurationSeconds.toNumber() * 1000 + 2000;
    console.log(`Sleep for ${sleepDurationMs}ms`);
    let success = await sleep(sleepDurationMs).then(async () => {
      await updateStakePoolEpochRpc(program, stakePool, []);
      const stakePoolData2 = await program.account.stakePool.fetch(stakePool);
      assert.equal(stakePoolData2.currentEpoch, stakePoolData.currentEpoch + 1);
      return true;
    });

    [rewardRecord3] = await deriveRewardRecord(program, rewardPool, 3);
    [rewardRecord2] = await deriveRewardRecord(program, rewardPool, 2);
    [rewardRecord4] = await deriveRewardRecord(program, rewardPool, 4);
    await createRewardRecordHelper(rewardRecord3, rewardRecord2);

    stakePoolData = await program.account.stakePool.fetch(stakePool);
    sleepDurationSeconds = stakePoolData.nextEpochStartTime.sub(
      new BN(new Date().getTime() / 1000)
    );
    sleepDurationMs = sleepDurationSeconds.toNumber() * 1000 + 2000;
    console.log(`Sleep for ${sleepDurationMs}ms`);
    success = await sleep(sleepDurationMs).then(async () => {
      await updateStakePoolEpochRpc(program, stakePool, []);
      const stakePoolData2 = await program.account.stakePool.fetch(stakePool);
      assert.equal(stakePoolData2.currentEpoch, stakePoolData.currentEpoch + 1);
      return true;
    });

    await createRewardRecordHelper(rewardRecord4, rewardRecord3);
  });

  it("Claim reward successfully for epoch 3", async () => {
    const accountDataBefore = await getAccountData(program, {
      rewardRecord: rewardRecord3,
      rewardPool,
      stakingRecord,
      allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount,
    });

    // Claim reward for this RewardPool, for user 1.
    await claimRewardHelper(
      rewardRecord2,
      rewardRecord3,
      stakingRecord,
      userBtcTokenAccount,
      userWalletKeypair
    );

    // Verify that rewards for epoch 3 is claimed.
    const accountDataAfter = await getAccountData(program, {
      stakingRecord,
      rewardRecord: rewardRecord3,
      allocatedTokenAccount,
      ownerTokenAccount: userBtcTokenAccount,
    });

    const expectedReward = accountDataBefore.rewardRecord.cummulativeReward
      .mul(accountDataBefore.stakingRecord.rewardUnits)
      .div(new BN(10 ** accountDataBefore.rewardPool.epochRewardDecimals));
    console.log("Expected Reward Claimed:", expectedReward.toNumber());

    // 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, for first RewardPool.
    expectedLastClaimedVec[0] = 3;
    assert.deepEqual(
      accountDataAfter.stakingRecord.lastEpochClaimedVec,
      expectedLastClaimedVec
    );
  });
});