import * as anchor from "@project-serum/anchor";
import { BN } from "@project-serum/anchor";
import { Keypair, Transaction, SystemProgram } from "@solana/web3.js";
import { assert } from "chai";
import { parseTranactionError } from "../deps/psyfi-euros/deps/tokenized-euros/packages/tokenized-euros/src";
import { updateStakePoolEpochRpc } from "../packages/psystake/src/instructions";
const serumCmn = require("@project-serum/common");
const TokenInstructions = require("@project-serum/serum").TokenInstructions;
const {
  TOKEN_PROGRAM_ID,
  ASSOCIATED_TOKEN_PROGRAM_ID,
  Token,
} = require("@solana/spl-token");
const { MAX_REWARD_POOLS } = require("./setup");

// Our own sleep function.
function sleep(ms) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

async function getTokenAccount(provider, addr) {
  return await serumCmn.getTokenAccount(provider, addr);
}

async function getMintInfo(provider, addr) {
  return await serumCmn.getMintInfo(provider, addr);
}

async function createMint(provider, authority) {
  if (authority === undefined) {
    authority = provider.wallet.publicKey;
  }
  const mint = Keypair.generate();
  const instructions = await createMintInstructions(
    provider,
    authority,
    mint.publicKey
  );

  const tx = new Transaction();
  tx.add(...instructions);

  await provider.send(tx, [mint]);

  return mint.publicKey;
}

async function createMintInstructions(provider, authority, mint) {
  let instructions = [
    SystemProgram.createAccount({
      fromPubkey: provider.wallet.publicKey,
      newAccountPubkey: mint,
      space: 82,
      lamports: await provider.connection.getMinimumBalanceForRentExemption(82),
      programId: TOKEN_PROGRAM_ID,
    }),
    TokenInstructions.initializeMint({
      mint,
      decimals: 6,
      mintAuthority: authority,
    }),
  ];
  return instructions;
}

async function createTokenAccount(provider, mint, owner) {
  const newAccount = Keypair.generate();
  const tx = new Transaction();
  tx.add(
    ...(await createTokenAccountInstrs(
      provider,
      newAccount.publicKey,
      mint,
      owner,
      await provider.connection.getMinimumBalanceForRentExemption(165)
    ))
  );
  await provider.send(tx, [newAccount]);
  return newAccount.publicKey;
}

async function createTokenAccountInstrs(
  provider,
  newAccountPubkey,
  mint,
  owner,
  lamports
) {
  return [
    SystemProgram.createAccount({
      fromPubkey: provider.wallet.publicKey,
      newAccountPubkey,
      space: 165,
      lamports,
      programId: TOKEN_PROGRAM_ID,
    }),
    TokenInstructions.initializeAccount({
      account: newAccountPubkey,
      mint,
      owner,
    }),
  ];
}

const getAccountData = async (program, accounts) => {
  let response = {};

  const provider = program.provider;

  if (accounts.stakePool) {
    const stakePoolData = await program.account.stakePool.fetch(
      accounts.stakePool
    );
    response["stakePool"] = stakePoolData;
  }
  if (accounts.rewardPool) {
    const rewardPoolData = await program.account.rewardPool.fetch(
      accounts.rewardPool
    );
    response["rewardPool"] = rewardPoolData;
  }
  if (accounts.stakingTokenAccount) {
    const tokenAccount = await getTokenAccount(
      provider,
      accounts.stakingTokenAccount
    );
    response["stakingTokenAccountAmount"] = tokenAccount.amount;
  }

  if (accounts.rewardTokenAccount) {
    const tokenAccount = await getTokenAccount(
      provider,
      accounts.rewardTokenAccount
    );
    response["rewardTokenAccountAmount"] = tokenAccount.amount;
  }
  if (accounts.allocatedTokenAccount) {
    const tokenAccount = await getTokenAccount(
      provider,
      accounts.allocatedTokenAccount
    );
    response["allocatedTokenAccountAmount"] = tokenAccount.amount;
  }
  if (accounts.rewardRecord) {
    response["rewardRecord"] = await program.account.rewardRecord.fetch(
      accounts.rewardRecord
    );
  }
  if (accounts.stakingRecord) {
    response["stakingRecord"] = await program.account.stakingRecord.fetch(
      accounts.stakingRecord
    );
  }
  if (accounts.ownerTokenAccount) {
    const tokenAccount = await getTokenAccount(
      provider,
      accounts.ownerTokenAccount
    );
    response["ownerTokenAccountAmount"] = tokenAccount.amount;
  }
  return response;
};

const getTokenAndMintAccounts = async (
  provider,
  vaultCollateralTokenAccountAddr,
  userCollateralTokenAccountAddr,
  userQuoteTokenAccountAddr,
  userVaultOwnershipTokenAccountAddr,
  vaultOwnershipTokenMintAddr,
  opts
) => {
  const vaultCollateralTokenAccount = await getTokenAccount(
    provider,
    vaultCollateralTokenAccountAddr
  );
  const userUsdcTokenAccount = await getTokenAccount(
    provider,
    userCollateralTokenAccountAddr
  );
  const userQuoteTokenAccount = await getTokenAccount(
    provider,
    userQuoteTokenAccountAddr
  );
  const userVaultOwnershipTokenAccount = await getTokenAccount(
    provider,
    userVaultOwnershipTokenAccountAddr
  );
  const vaultOwnershipTokenMint = await getMintInfo(
    provider,
    vaultOwnershipTokenMintAddr
  );
  let vaultTokenSupply = vaultOwnershipTokenMint.supply;

  let vaultAccount,
    curEpochHistory,
    prevEpochHistory,
    withdrawalCollateralAssetAccount;
  let writerTokenAmount = 0;
  let writerTokenSupply = 0;
  let optionTokenAmount = 0;
  let optionTokenSupply = 0;
  let collateralPoolAmount = 0;
  let vaultOwnershipTokenAmount = 0;
  let feesCollected = 0;

  const psyfiEurosProgram = anchor.workspace.PsyfiEuros;

  if (opts.withdrawalCollateralAssetAccount) {
    withdrawalCollateralAssetAccount = await getTokenAccount(
      provider,
      opts.withdrawalCollateralAssetAccount
    );
  }

  if (opts.writerTokenAccount) {
    const writerTokenAccount = await getTokenAccount(
      provider,
      opts.writerTokenAccount
    );
    writerTokenAmount = writerTokenAccount.amount;
  }

  if (opts.writerTokenMint) {
    const writerTokenMint = await getMintInfo(provider, opts.writerTokenMint);
    writerTokenSupply = writerTokenMint.supply;
  }

  if (opts.optionTokenAccount) {
    const optionTokenAccount = await getTokenAccount(
      provider,
      opts.optionTokenAccount
    );
    optionTokenAmount = optionTokenAccount.amount;
  }

  if (opts.optionTokenMint) {
    const optionTokenMint = await getMintInfo(provider, opts.optionTokenMint);
    optionTokenSupply = optionTokenMint.supply;
  }

  if (opts.vaultOwnershipTokenAccount) {
    const vaultOwnershipTokenAccount = await getTokenAccount(
      provider,
      opts.vaultOwnershipTokenAccount
    );
    vaultOwnershipTokenAmount = vaultOwnershipTokenAccount.amount;
  }

  if (opts.vaultAccount) {
    vaultAccount = await psyfiEurosProgram.account.vaultAccount.fetch(
      opts.vaultAccount
    );
  }
  if (opts.curEpochHistory) {
    curEpochHistory = await psyfiEurosProgram.account.epochHistory.fetch(
      opts.curEpochHistory
    );
  }
  if (opts.prevEpochHistory) {
    prevEpochHistory = await psyfiEurosProgram.account.epochHistory.fetch(
      opts.prevEpochHistory
    );
  }
  if (opts.collateralAssetPool) {
    const collateralAssetPool = await getTokenAccount(
      provider,
      opts.collateralAssetPool
    );
    collateralPoolAmount = collateralAssetPool.amount;
  }

  if (opts.feeTokenAccount) {
    const feeTokenAccount = await getTokenAccount(
      provider,
      opts.feeTokenAccount
    );
    feesCollected = feeTokenAccount.amount;
  }

  return {
    vaultAccount,
    curEpochHistory,
    prevEpochHistory,
    vaultCollateral: vaultCollateralTokenAccount.amount,
    withdrawalCollateral: withdrawalCollateralAssetAccount
      ? withdrawalCollateralAssetAccount.amount
      : null,
    userCollateral: userUsdcTokenAccount.amount,
    userQuote: userQuoteTokenAccount.amount,
    userVaultToken: userVaultOwnershipTokenAccount.amount,
    vaultTokenSupply,
    writerToken: writerTokenAmount,
    writerTokenSupply,
    optionToken: optionTokenAmount,
    optionTokenSupply,
    collateralPoolAmount,
    vaultOwnershipTokenAmount,
    feesCollected,
  };
};

const defaultVec = () => new Array(MAX_REWARD_POOLS).fill(0);

// Util function to wait and call instruction to update a StakePool to reach a target epoch.
// Returns true when target epoch is reached.
const waitForStakePoolEpoch = async (program, stakePool, targetEpoch) => {
  let stakePoolData = await program.account.stakePool.fetch(stakePool);
  console.log("stakePoolData.currentEpoch", stakePoolData.currentEpoch);

  if (stakePoolData.currentEpoch >= targetEpoch) return true;

  const sleepDurationSeconds = stakePoolData.nextEpochStartTime
    .sub(new BN(new Date().getTime() / 1000))
    .toNumber();

  if (sleepDurationSeconds >= -1) {
    const sleepDurationMs = sleepDurationSeconds * 1000 + 2000;
    console.log(`Sleeping for ${sleepDurationMs}ms`);
    const slept = await sleep(sleepDurationMs).then(() => true);
    if (slept) {
      await updateStakePoolEpochRpc(program, stakePool, []);
      stakePoolData = await program.account.stakePool.fetch(stakePool);
      return waitForStakePoolEpoch(program, stakePool, targetEpoch);
    }
  } else {
    await updateStakePoolEpochRpc(program, stakePool, []);
    stakePoolData = await program.account.stakePool.fetch(stakePool);
    return waitForStakePoolEpoch(program, stakePool, targetEpoch);
  }
};

// Day duration is set to 1s for testing.
const TESTING_DAY_DURATION = 1;

const calculateLockupDuration = (lockUpPeriod) => {
  let duration;
  switch (lockUpPeriod) {
    case 0:
      duration = 0;
      break;
    case 1:
      duration = 30 * TESTING_DAY_DURATION;
      break;
    case 2:
      duration = 60 * TESTING_DAY_DURATION;
      break;
    case 3:
      duration = 180 * TESTING_DAY_DURATION;
      break;
    case 4:
      duration = 360 * TESTING_DAY_DURATION;
      break;
    default:
      duration = 0;
  }
  return duration;
};

const assertParsedError = (err: any, msg: String) => {
  const programError = parseTranactionError(err);
  if (programError != null) {
    assert.equal(programError.msg, msg);
  }
};

module.exports = {
  TOKEN_PROGRAM_ID,
  sleep,
  createMint,
  createTokenAccount,
  getTokenAccount,
  getMintInfo,
  getTokenAndMintAccounts,
  getAccountData,
  defaultVec,
  waitForStakePoolEpoch,
  calculateLockupDuration,
  assertParsedError,
};