import * as anchor from "@project-serum/anchor";
import { Program, BN } from "@project-serum/anchor";
import { sleep } from "@project-serum/common";
import { assert } from "chai";
import { Keypair, PublicKey } from "@solana/web3.js";
import { PsyStake } from "../target/types/psy_stake";
import {
  deriveStakePool,
  deriveStakingTokenAccount,
  deriveRewardPoolPdas,
  deriveStakingRecord,
} from "./pdas";
const { setupTestAccounts, MIN_EPOCH_DURATION } = require("./setup");
const {
  getAccountData,
  sleep,
  waitForStakePoolEpoch,
  assertParsedError,
} = require("./utils");
import {
  stakeTokenRpc,
  claimRewardRpc,
  unstakeTokenRpc,
  createRewardPoolRpc,
  updateRewardPoolRpc,
  createStakingRecordRpc,
  createRewardRecordRpc,
} from "../packages/psystake/src/instructions";
import {
  DistributionType,
  calculateRewardUnits,
} from "../packages/psystake/src";

describe("Test unstake token 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,
    stakePoolAuthority: Keypair,
    rewardPoolAuthority: Keypair;

  const uniqueSeed = 1;

  let stakePool: PublicKey,
    stakingTokenAccount: PublicKey,
    stakingTokenMint: PublicKey,
    rewardPool: PublicKey,
    stakingRecord: PublicKey,
    stakingRecordBump: number;

  const startingEpoch = 1;
  let poolId = 1;
  const epochDuration = new BN(MIN_EPOCH_DURATION);

  const distributionType = DistributionType.CONSTANT;
  const rewardPerEpoch = new BN(1000000);
  const epochRewardDecimals = 8;

  const createStakingRecordHelper = async (
    recordOwnerKeypair = userWalletKeypair
  ) => {
    await createStakingRecordRpc(
      program,
      recordOwnerKeypair.publicKey,
      stakePool,
      [recordOwnerKeypair]
    );
  };

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

  const unstakeTokenHelper = async (unstakeAmount, ownerTokenAccount) => {
    await unstakeTokenRpc(
      program,
      userWalletKeypair.publicKey,
      stakePool,
      ownerTokenAccount,
      unstakeAmount,
      [userWalletKeypair]
    );
  };

  const createRewardPoolHelper = async (startingEpoch, poolId) => {
    await createRewardPoolRpc(
      program,
      stakePoolAuthority.publicKey,
      rewardPoolAuthority.publicKey,
      stakePool,
      network.btcToken.publicKey,
      poolId,
      startingEpoch,
      distributionType,
      rewardPerEpoch,
      epochRewardDecimals,
      [stakePoolAuthority]
    );
  };

  const updateRewardPoolHelper = async (isActive, poolId) => {
    await updateRewardPoolRpc(
      program,
      rewardPoolAuthority.publicKey,
      stakePool,
      poolId,
      distributionType,
      rewardPerEpoch,
      isActive,
      [rewardPoolAuthority]
    );
  };

  const claimRewardHelper = async (
    startRewardRecordNumber,
    endRewardRecordNumber
  ) => {
    await claimRewardRpc(
      program,
      userWalletKeypair.publicKey,
      stakePool,
      userBtcTokenAccount,
      startRewardRecordNumber,
      endRewardRecordNumber,
      poolId,
      [userWalletKeypair]
    );
  };

  const createRewardRecordHelper = async (epochToCreate) => {
    await createRewardRecordRpc(
      program,
      userWalletKeypair.publicKey,
      stakePool,
      poolId,
      epochToCreate,
      [userWalletKeypair]
    );
  };

  before(async () => {
    ({
      userWalletKeypair,
      userUsdcTokenAccount,
      userBtcTokenAccount,
      userWalletKeypair2,
      userUsdcTokenAccount2,
      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 + 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 } = poolAccounts);
    await createRewardPoolHelper(startingEpoch, poolId);

    // Create a StakingRecord owned by first user.
    [stakingRecord, stakingRecordBump] = await deriveStakingRecord(
      program,
      userWalletKeypair.publicKey,
      stakePool
    );
    await createStakingRecordHelper();
  });

  it("Unstake token fails when owner token account is from different mint.", async () => {
    const stakeAmount = new BN(500);
    const lockUpPeriod = 0;
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    try {
      await unstakeTokenHelper(stakeAmount, userBtcTokenAccount);
    } catch (err) {
      assertParsedError(err, "Owner token account is invalid.");
    }
  });

  it("Unstake token fails when staked amount is insufficient.", async () => {
    const unstakeAmount = new BN(10 ** 12);

    try {
      await unstakeTokenHelper(unstakeAmount, userUsdcTokenAccount);
    } catch (err) {
      assertParsedError(err, "Insufficient amount for staking or unstaking.");
    }
  });

  it("Unstake all tokens successfully.", async () => {
    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    await unstakeTokenHelper(
      stakingRecordBefore.stakedAmount,
      userUsdcTokenAccount
    );

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.isZero());

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.isZero());
    assert(stakingRecordAfter.stakedAmount.isZero());

    // Verify transfer of tokens.
    assert(
      accountDataBefore.stakingTokenAccountAmount
        .sub(accountDataAfter.stakingTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);
    assert(
      stakingRecordAfter.lockUpExpiry.eq(stakingRecordBefore.lockUpExpiry)
    );
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      stakingRecordBefore.lastEpochClaimedVec
    );
  });

  it("Unstake token should decrement total reward units in StakePool.", async () => {
    // Stake some tokens for first User.
    const stakeAmount1 = new BN(100);
    await stakeTokenHelper(new BN(100), 0, userUsdcTokenAccount);

    // Create a StakingRecord owned by second user.
    const [stakingRecord2, stakingRecordBump2] = await deriveStakingRecord(
      program,
      userWalletKeypair2.publicKey,
      stakePool
    );
    await createStakingRecordHelper(userWalletKeypair2);

    // Stake tokens for second user.
    const stakeAmount2 = new BN(10 ** 9);
    await program.methods
      .stakeToken(stakeAmount2, 0)
      .accounts({
        recordOwner: userWalletKeypair2.publicKey,
        stakePool,
        stakingRecord: stakingRecord2,
        ownerTokenAccount: userUsdcTokenAccount2,
        stakingTokenAccount,
      })
      .signers([userWalletKeypair2])
      .rpc();

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord: stakingRecord2,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount2,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    // Expect reward units in StakePool to equal amount staked by user 1 and 2.
    assert(
      accountDataBefore.stakePool.totalRewardUnits.eq(
        calculateRewardUnits(stakeAmount1, 0).add(
          calculateRewardUnits(stakeAmount2, 0)
        )
      )
    );

    // Unstake tokens for second user.
    await program.methods
      .unstakeToken(stakingRecordBefore.stakedAmount)
      .accounts({
        recordOwner: userWalletKeypair2.publicKey,
        stakePool,
        stakingRecord: stakingRecord2,
        ownerTokenAccount: userUsdcTokenAccount2,
        stakingTokenAccount,
      })
      .signers([userWalletKeypair2])
      .rpc();

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord: stakingRecord2,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount2,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Expect reward units in StakePool to equal amount staked by user 1.
    assert(
      accountDataAfter.stakePool.totalRewardUnits.eq(
        calculateRewardUnits(stakeAmount1, 0)
      )
    );

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.isZero());
    assert(stakingRecordAfter.stakedAmount.isZero());

    // Verify transfer of tokens.
    assert(
      accountDataBefore.stakingTokenAccountAmount
        .sub(accountDataAfter.stakingTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair2.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump2);
    assert(
      stakingRecordAfter.lockUpExpiry.eq(stakingRecordBefore.lockUpExpiry)
    );

    // Last epoch claim vec should not change after unstake.
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      stakingRecordBefore.lastEpochClaimedVec
    );
  });

  it("Unstake token fails when rewards are not claimed.", async () => {
    const stakeAmount = new BN(500);
    const lockUpPeriod = 0;
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

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

    // Create RewardRecords for epoch 2.
    await createRewardRecordHelper(2);

    // Expect unstake token to fail since rewards in epoch 1 has not been claimed.
    try {
      await unstakeTokenHelper(stakeAmount, userUsdcTokenAccount);
    } catch (err) {
      assertParsedError(err, "Some RewardPool is unclaimed.");
    }
  });

  it("Unstake token succeeds after rewards are claimed.", async () => {
    // Claim rewards for epoch 1.
    await claimRewardHelper(1, 1);

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    await unstakeTokenHelper(
      stakingRecordBefore.stakedAmount,
      userUsdcTokenAccount
    );

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.isZero());

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.isZero());
    assert(stakingRecordAfter.stakedAmount.isZero());

    // Verify transfer of tokens.
    assert(
      accountDataBefore.stakingTokenAccountAmount
        .sub(accountDataAfter.stakingTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);
    assert(
      stakingRecordAfter.lockUpExpiry.eq(stakingRecordBefore.lockUpExpiry)
    );
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      stakingRecordBefore.lastEpochClaimedVec
    );
  });

  it("Unstake token fails when lockup has yet to expire.", async () => {
    const stakeAmount = new BN(5 * 10 ** 9);
    const lockUpPeriod = 1;

    // Stake token with non-zero lockup duration.
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    try {
      await unstakeTokenHelper(stakeAmount, userUsdcTokenAccount);
    } catch (err) {
      assertParsedError(err, "Lockup period has not expire.");
    }
  });

  it("Unstake partial position after lockup expiry.", async () => {
    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    const lockupHasExpired = async () => {
      const sleepDurationSeconds = stakingRecordBefore.lockUpExpiry
        .sub(new BN(new Date().getTime() / 1000))
        .toNumber();

      if (sleepDurationSeconds < -1) return true;

      const sleepDurationMs = sleepDurationSeconds * 1000 + 2000;
      console.log(`Sleeping for ${sleepDurationMs}ms`);
      const slept = await sleep(sleepDurationMs).then(() => true);
      if (slept) return true;
    };

    assert(await lockupHasExpired());

    const unstakeAmount = new BN(5);
    await unstakeTokenHelper(unstakeAmount, userUsdcTokenAccount);

    // Since 5*10**9 tokens were staked previous with 30 days lock up,
    // 5*10**9 * 1.15 = 5.75*10**9 reward units were issued.
    // Remaining 4999999995 tokens should be equivalent to 5749999994 reward units,
    // due to the flooring mechanism for reward units calculation.
    const remainingAmount = stakingRecordBefore.stakedAmount.sub(unstakeAmount);
    const expectedRewardUnits = calculateRewardUnits(
      stakingRecordBefore.stakedAmount.sub(unstakeAmount),
      1
    );

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(stakingRecordAfter.stakedAmount.eq(remainingAmount));

    // Verify transfer of tokens.
    assert(
      accountDataBefore.stakingTokenAccountAmount
        .sub(accountDataAfter.stakingTokenAccountAmount)
        .eq(unstakeAmount)
    );
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(unstakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);
    assert(
      stakingRecordAfter.lockUpExpiry.eq(stakingRecordBefore.lockUpExpiry)
    );
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      stakingRecordBefore.lastEpochClaimedVec
    );
  });

  it("Unstake token after RewardPool is made inactive.", async () => {
    // Change first RewardPool to inactive.
    await updateRewardPoolHelper(false, 1);

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

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    await unstakeTokenHelper(
      stakingRecordBefore.stakedAmount,
      userUsdcTokenAccount
    );

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.isZero());

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.isZero());
    assert(stakingRecordAfter.stakedAmount.isZero());

    // Verify transfer of tokens.
    assert(
      accountDataBefore.stakingTokenAccountAmount
        .sub(accountDataAfter.stakingTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);
    assert(
      stakingRecordAfter.lockUpExpiry.eq(stakingRecordBefore.lockUpExpiry)
    );
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      stakingRecordBefore.lastEpochClaimedVec
    );
  });

  it("Unstake token after a new RewardPool is added.", async () => {
    // Stake some tokens.
    await stakeTokenHelper(new BN(100), 0, userUsdcTokenAccount);

    // Add a new RewardPool to start in Epoch 5.
    await createRewardPoolHelper(5, 2);

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    await unstakeTokenHelper(
      stakingRecordBefore.stakedAmount,
      userUsdcTokenAccount
    );

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
      ownerTokenAccount: userUsdcTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.isZero());

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.isZero());
    assert(stakingRecordAfter.stakedAmount.isZero());

    // Verify transfer of tokens.
    assert(
      accountDataBefore.stakingTokenAccountAmount
        .sub(accountDataAfter.stakingTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );
    assert(
      accountDataAfter.ownerTokenAccountAmount
        .sub(accountDataBefore.ownerTokenAccountAmount)
        .eq(stakingRecordBefore.stakedAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);
    assert(
      stakingRecordAfter.lockUpExpiry.eq(stakingRecordBefore.lockUpExpiry)
    );

    // Last epoch claim vec should not change after unstake.
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      stakingRecordBefore.lastEpochClaimedVec
    );
  });
});