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

import { assert } from "chai";
import { Keypair, PublicKey } from "@solana/web3.js";
const {
  setupTestAccounts,
  MAX_REWARD_POOLS,
  MIN_EPOCH_DURATION,
} = require("./setup");
const {
  getAccountData,
  defaultVec,
  waitForStakePoolEpoch,
  assertParsedError,
} = require("./utils");

describe("Test RewardPool instructions.", () => {
  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;

  const startingEpoch = 1;
  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();
  };

  before(async () => {
    ({ network, 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 + 1);
    await program.methods
      .createStakePool(uniqueSeed, nextEpochStartTime, epochDuration)
      .accounts({
        stakePoolAuthority: stakePoolAuthority.publicKey,
        stakePool,
        stakingTokenMint,
        stakingTokenAccount,
      })
      .signers([stakePoolAuthority])
      .rpc();
  });

  it("Initialization of RewardPool fails when starting epoch is zero.", async () => {
    const distributionType = { constant: {} };
    const rewardPerEpoch = new BN(0);
    const zeroEpoch = 0;

    try {
      await createRewardPoolHelper(
        zeroEpoch,
        poolId,
        distributionType,
        rewardPerEpoch
      );
      assert(false);
    } catch (err) {
      assertParsedError(
        err,
        "Starting epoch of RewardPool cannot be in the past."
      );
    }
  });

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

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

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

    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.add(epochDuration)
      )
    );
    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.rewardTokenAccountAmount.isZero());
    assert(accountData.allocatedTokenAccountAmount.isZero());
  });

  it("Initialization of RewardPool fails when max percentage is exceeded.", async () => {
    const distributionType = { percentage: {} };
    const rewardPerEpoch = new BN(10000 * 1000 + 1); // More than 10000 Mbps.

    // Increment pool id.
    poolId++;

    try {
      await createRewardPoolHelper(
        startingEpoch,
        poolId,
        distributionType,
        rewardPerEpoch
      );
      assert(false);
    } catch (err) {
      assertParsedError(err, "Percentage chosen is invalid.");
    }
  });

  it("Initializes a RewardPool succeeds with percentage reward.", async () => {
    const distributionType = { percentage: {} };
    const rewardPerEpoch = new BN(10000); // 10 bps.

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

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

    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.isZero());
    assert(rewardPoolData.percentageRewardMbpsPerEpoch.eq(rewardPerEpoch));
    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.add(epochDuration)
      )
    );
    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.rewardTokenAccountAmount.isZero());
    assert(accountData.allocatedTokenAccountAmount.isZero());
  });

  it("Initialization of RewardPool fails when pool id already exists.", async () => {
    const distributionType = { percentage: {} };
    const rewardPerEpoch = new BN(10000); // 10 bps.

    try {
      await createRewardPoolHelper(
        startingEpoch,
        poolId,
        distributionType,
        rewardPerEpoch
      );
      assert(false);
    } catch (err) {
      assert.match(
        err.message,
        /Cross-program invocation with unauthorized signer or writable account/
      );
    }
  });

  it("Updating RewardPool fails when setting inactive before pool starts.", async () => {
    const distributionType = { percentage: {} };
    const rewardPerEpoch = new BN(100);
    const isActive = false;

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

    try {
      await program.methods
        .updateRewardPool(distributionType, rewardPerEpoch, isActive)
        .accounts({
          rewardPoolAuthority: rewardPoolAuthority.publicKey,
          stakePool,
          rewardPool: poolAccounts.rewardPool,
        })
        .signers([rewardPoolAuthority])
        .rpc();
      assert(false);
    } catch (err) {
      assertParsedError(
        err,
        "Cannot change a RewardPool to inactive before it starts."
      );
    }
  });

  it("Initialization of RewardPool fails when starting epoch is older than current epoch.", async () => {
    const distributionType = { percentage: {} };
    const rewardPerEpoch = new BN(10000); // 10 bps.
    // Increment pool id.
    poolId++;

    const testPassed = await waitForStakePoolEpoch(program, stakePool, 2).then(
      async () => {
        console.log("Target Epoch Reached");
        // Verify that instruction fails when current epoch is 2 and starting epoch for RewardPool is 1.
        try {
          await createRewardPoolHelper(
            startingEpoch,
            poolId,
            distributionType,
            rewardPerEpoch
          );
          assert(false);
        } catch (err) {
          assertParsedError(
            err,
            "Starting epoch of RewardPool cannot be in the past."
          );
          return true;
        }
      }
    );

    assert(testPassed);
  });

  it("Initialization of RewardPool fails when epoch reward decimals exceeded.", async () => {
    const distributionType = { percentage: {} };
    const rewardPerEpoch = new BN(10000); // 10 bps.

    try {
      await createRewardPoolHelper(
        2,
        poolId,
        distributionType,
        rewardPerEpoch,
        /*decimals=*/ 16
      );
      assert(false);
    } catch (err) {
      assertParsedError(err, "Limit of decimals allowed has been exceeded.");
      return true;
    }
  });

  it("Initialization of RewardPool fails when max pool limit exceeded.", async () => {
    const distributionType = { percentage: {} };
    const rewardPerEpoch = new BN(10000); // 10 bps.

    let stakePoolData = await program.account.stakePool.fetch(stakePool);

    // Initialize the third RewardPool which should succeed.
    await createRewardPoolHelper(
      stakePoolData.currentEpoch,
      poolId,
      distributionType,
      rewardPerEpoch
    );

    // Get updated StakePool.
    stakePoolData = await program.account.stakePool.fetch(stakePool);

    // Check that max pool count is reached.
    assert(stakePoolData.rewardPoolCount, MAX_REWARD_POOLS);

    // Verify that RewardPool creation fails when max pool count is reached.
    try {
      poolId++;
      await createRewardPoolHelper(
        stakePoolData.currentEpoch,
        poolId,
        distributionType,
        rewardPerEpoch
      );
    } catch (err) {
      assertParsedError(
        err,
        "Max. no. of RewardPool per stake pool has been reached."
      );
      return true;
    }
  });

  it("Updating RewardPool succeeds in changing distribution strategy.", async () => {
    const distributionType = { percentage: {} };
    const rewardPerEpoch = new BN(222);

    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      1,
      startingEpoch
    );
    await program.methods
      .updateRewardPool(distributionType, rewardPerEpoch, true)
      .accounts({
        rewardPoolAuthority: rewardPoolAuthority.publicKey,
        stakePool,
        rewardPool: poolAccounts.rewardPool,
      })
      .signers([rewardPoolAuthority])
      .rpc();

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

    // Verify RewardPool account state.
    const rewardPoolData = accountData.rewardPool;
    assert(rewardPoolData.stakePool.equals(stakePool));
    assert(
      rewardPoolData.rewardPoolAuthority.equals(rewardPoolAuthority.publicKey)
    );
    assert.equal(rewardPoolData.poolId, 1);
    assert(rewardPoolData.rewardTokenMint.equals(network.btcToken.publicKey));
    assert(
      rewardPoolData.rewardTokenAccount.equals(poolAccounts.rewardTokenAccount)
    );
    assert(
      rewardPoolData.allocatedTokenAccount.equals(
        poolAccounts.allocatedTokenAccount
      )
    );
    assert.deepEqual(rewardPoolData.distributionType, distributionType);
    // constantRewardPerEpoch should be changed to 0.
    assert(rewardPoolData.constantRewardPerEpoch.isZero());
    assert(rewardPoolData.percentageRewardMbpsPerEpoch.eq(rewardPerEpoch));
    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
    );
  });

  it("Updating RewardPool succeeds in changing distribution strategy again.", async () => {
    const distributionType = { constant: {} };
    const rewardPerEpoch = new BN(1);

    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      1,
      startingEpoch
    );
    await program.methods
      .updateRewardPool(distributionType, rewardPerEpoch, true)
      .accounts({
        rewardPoolAuthority: rewardPoolAuthority.publicKey,
        stakePool,
        rewardPool: poolAccounts.rewardPool,
      })
      .signers([rewardPoolAuthority])
      .rpc();

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

    // Verify RewardPool account state.
    const rewardPoolData = accountData.rewardPool;
    assert(rewardPoolData.stakePool.equals(stakePool));
    assert(
      rewardPoolData.rewardPoolAuthority.equals(rewardPoolAuthority.publicKey)
    );
    assert.equal(rewardPoolData.poolId, 1);
    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));
    // percentageRewardMbpsPerEpoch should be changed to 0.
    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
    );
  });

  it("Updating RewardPool fails when max percentage is exceeded.", async () => {
    const distributionType = { percentage: {} };
    const rewardPerEpoch = new BN(10000 * 1000 + 1); // More than 10000 Mbps.

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

    try {
      await program.methods
        .updateRewardPool(distributionType, rewardPerEpoch, true)
        .accounts({
          rewardPoolAuthority: rewardPoolAuthority.publicKey,
          stakePool,
          rewardPool: poolAccounts.rewardPool,
        })
        .signers([rewardPoolAuthority])
        .rpc();
      assert(false);
    } catch (err) {
      assertParsedError(err, "Percentage chosen is invalid.");
    }
  });

  it("Updating RewardPool succeeds in setting RewardPool to inactive after pool starts.", async () => {
    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      1,
      startingEpoch
    );

    const distributionType = { constant: {} };
    const rewardPerEpoch = new BN(1);
    await program.methods
      .updateRewardPool(distributionType, rewardPerEpoch, false)
      .accounts({
        rewardPoolAuthority: rewardPoolAuthority.publicKey,
        stakePool,
        rewardPool: poolAccounts.rewardPool,
      })
      .signers([rewardPoolAuthority])
      .rpc();

    // Verify RewardPool account state.
    const accountData = await getAccountData(program, {
      stakePool,
      ...poolAccounts,
    });
    const rewardPoolData = accountData.rewardPool;
    const stakePoolData = accountData.stakePool;
    assert.isFalse(rewardPoolData.isActive);

    // Verify that current epoch of stakePool is after starting epoch.
    assert(stakePoolData.currentEpoch > startingEpoch);

    // Verify that inactive epoch vec is consistent.
    const expectedInactiveEpochVec = defaultVec();
    expectedInactiveEpochVec[rewardPoolData.poolId - 1] =
      stakePoolData.currentEpoch;
    assert.deepEqual(stakePoolData.inactiveEpochVec, expectedInactiveEpochVec);

    // Other fields should remain unchanged.
    assert(rewardPoolData.stakePool.equals(stakePool));
    assert(
      rewardPoolData.rewardPoolAuthority.equals(rewardPoolAuthority.publicKey)
    );
    assert.equal(rewardPoolData.poolId, 1);
    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.equal(rewardPoolData.epochRewardDecimals, epochRewardDecimals);
    assert.equal(rewardPoolData.rewardPoolBump, poolAccounts.rewardPoolBump);
    assert.equal(
      rewardPoolData.allocatedTokenAccountBump,
      poolAccounts.allocatedTokenAccountBump
    );
  });

  it("Updating RewardPool fails in setting an inactive RewardPool to active", async () => {
    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      1,
      startingEpoch
    );
    const distributionType = { constant: {} };
    const rewardPerEpoch = new BN(1);

    try {
      await program.methods
        .updateRewardPool(distributionType, rewardPerEpoch, true)
        .accounts({
          rewardPoolAuthority: rewardPoolAuthority.publicKey,
          stakePool,
          rewardPool: poolAccounts.rewardPool,
        })
        .signers([rewardPoolAuthority])
        .rpc();
    } catch (err) {
      assertParsedError(err, "RewardPool is no longer active.");
    }
  });
});