/**
 * Methods from jup-lock:
 * - createLockerProgram
 * - deriveEscrow
 * - createVestingPlanV2
 * Reference: https://github.com/jup-ag/jup-lock/blob/main/tests/locker_utils/index.ts
 */

import assert from 'assert';
import 'dotenv/config';

import {
  ASSOCIATED_TOKEN_PROGRAM_ID,
  createAssociatedTokenAccountInstruction,
  getAssociatedTokenAddressSync,
  TOKEN_2022_PROGRAM_ID,
  TOKEN_PROGRAM_ID,
} from '@solana/spl-token';
import {
  AnchorProvider,
  BN,
  Program,
  Wallet,
  web3,
  workspace,
} from '@coral-xyz/anchor';
import { AccountMeta, Connection, TransactionExpiredTimeoutError } from '@solana/web3.js';

// TODO: Generate type file from IDL json
import { Locker } from '../../target/types/locker';
import { TokenExtensionUtil } from './token-2022/token-extensions';
import {
  OptionRemainingAccountsInfoData,
  RemainingAccountsBuilder,
  RemainingAccountsType,
} from './token-2022/remaining-accounts';

assert(process.env.RPC_ENDPOINT);

const connection = new Connection(process.env.RPC_ENDPOINT);

const ESCROW_USE_SPL_TOKEN = 0;

const MEMO_PROGRAM = new web3.PublicKey(
  "MemoSq4gqABAXKb96qnH8TysNcWxMyWCqXgDLGmfcHr"
);

export function createLockerProgram(wallet: Wallet): Program<Locker> {
  const provider = new AnchorProvider(connection, wallet, {
    maxRetries: 3,
  });
  provider.opts.commitment = 'confirmed';

  return workspace.Locker as Program<Locker>;
}

export function deriveEscrow(base: web3.PublicKey, programId: web3.PublicKey) {
  return web3.PublicKey.findProgramAddressSync(
    [Buffer.from('escrow'), base.toBuffer()],
    programId
  );
}

export interface CreateVestingPlanParams {
  ownerKeypair: web3.Keypair;
  tokenMint: web3.PublicKey;
  isAssertion: boolean;
  vestingStartTime: BN;
  cliffTime: BN;
  frequency: BN;
  cliffUnlockAmount: BN;
  amountPerPeriod: BN;
  numberOfPeriod: BN;
  recipient: web3.PublicKey;
  updateRecipientMode: number;
  cancelMode: number;
  tokenProgram?: web3.PublicKey;
}

// V2 instructions
export async function createVestingPlanV2(params: CreateVestingPlanParams) {
  let {
    tokenMint,
    ownerKeypair,
    vestingStartTime,
    cliffTime,
    frequency,
    cliffUnlockAmount,
    amountPerPeriod,
    numberOfPeriod,
    recipient,
    updateRecipientMode,
    cancelMode,
    tokenProgram,
  } = params;

  const program = createLockerProgram(new Wallet(ownerKeypair));

  const baseKP = web3.Keypair.generate();

  let [escrow] = deriveEscrow(baseKP.publicKey, program.programId);

  const senderToken = getAssociatedTokenAddressSync(
    tokenMint,
    ownerKeypair.publicKey,
    false,
    tokenProgram,
    ASSOCIATED_TOKEN_PROGRAM_ID
  );

  const escrowToken = getAssociatedTokenAddressSync(
    tokenMint,
    escrow,
    true,
    tokenProgram,
    ASSOCIATED_TOKEN_PROGRAM_ID
  );

  let remainingAccountsInfo = null;
  let remainingAccounts: AccountMeta[] = [];
  if (tokenProgram == TOKEN_2022_PROGRAM_ID) {
    let inputTransferHookAccounts =
      await TokenExtensionUtil.getExtraAccountMetasForTransferHook(
        program.provider.connection,
        tokenMint,
        senderToken,
        escrowToken,
        ownerKeypair.publicKey,
        TOKEN_2022_PROGRAM_ID
      );

    [remainingAccountsInfo, remainingAccounts] = new RemainingAccountsBuilder()
      .addSlice(
        RemainingAccountsType.TransferHookEscrow,
        inputTransferHookAccounts
      )
      .build() as [OptionRemainingAccountsInfoData, AccountMeta[]];
  }

  assert(tokenProgram);

  try {
    await program.methods
      .createVestingEscrowV2(
        {
          vestingStartTime,
          cliffTime,
          frequency,
          cliffUnlockAmount,
          amountPerPeriod,
          numberOfPeriod,
          updateRecipientMode,
          cancelMode,
        },
        remainingAccountsInfo
      )
      .accounts({
        base: baseKP.publicKey,
        senderToken,
        escrowToken,
        recipient,
        tokenMint,
        sender: ownerKeypair.publicKey,
        tokenProgram,
        systemProgram: web3.SystemProgram.programId,
        escrow,
        // TODO: Fix type error for escrowToken
      } as any)
      .remainingAccounts(remainingAccounts ? remainingAccounts : [])
      .preInstructions([
        createAssociatedTokenAccountInstruction(
          ownerKeypair.publicKey,
          escrowToken,
          escrow,
          tokenMint,
          tokenProgram,
          ASSOCIATED_TOKEN_PROGRAM_ID
        ),
      ])
      .signers([baseKP, ownerKeypair])
      .rpc();

    return escrow;
  } catch (error) {
    if (error instanceof TransactionExpiredTimeoutError) {
      console.error('Transaction confirmation delayed for',error.signature);
      console.log('Confirming the transaction again...');
      const confirmedTransaction = await connection.getTransaction(error.signature, {
        commitment: 'confirmed',
        maxSupportedTransactionVersion: 0
      });

      if(confirmedTransaction === null) {
        console.error('Transaction failed for',error.signature);
      }

      return escrow;
    }
  }
}

export interface ClaimTokenParamsV2 {
  isAssertion: boolean;
  escrow: web3.PublicKey;
  recipient: web3.Keypair;
  maxAmount: BN;
  recipientToken: web3.PublicKey;
  tokenProgram: web3.PublicKey;
}

export async function claimTokenV2(params: ClaimTokenParamsV2) {
  let { isAssertion, escrow, recipient, maxAmount, recipientToken } = params;
  const program = createLockerProgram(new Wallet(recipient));
  const escrowState = await program.account.vestingEscrow.fetch(escrow);
  const tokenProgram =
    escrowState.tokenProgramFlag == ESCROW_USE_SPL_TOKEN
      ? TOKEN_PROGRAM_ID
      : TOKEN_2022_PROGRAM_ID;

  const escrowToken = getAssociatedTokenAddressSync(
    escrowState.tokenMint,
    escrow,
    true,
    tokenProgram,
    ASSOCIATED_TOKEN_PROGRAM_ID
  );

  let remainingAccountsInfo = null;
  let remainingAccounts: AccountMeta[] | undefined = [];
  if (tokenProgram == TOKEN_2022_PROGRAM_ID) {
    let claimTransferHookAccounts =
      await TokenExtensionUtil.getExtraAccountMetasForTransferHook(
        program.provider.connection,
        escrowState.tokenMint,
        escrowToken,
        recipientToken,
        escrow,
        TOKEN_2022_PROGRAM_ID
      );

    [remainingAccountsInfo, remainingAccounts] = new RemainingAccountsBuilder()
      .addSlice(
        RemainingAccountsType.TransferHookEscrow,
        claimTransferHookAccounts
      )
      .build();
  }

  const tx = await program.methods
    .claimV2(maxAmount, remainingAccountsInfo)
    .accounts({
      tokenProgram,
      tokenMint: escrowState.tokenMint,
      memoProgram: MEMO_PROGRAM,
      escrow,
      escrowToken,
      recipient: recipient.publicKey,
      recipientToken,
    } as any)
    .remainingAccounts(remainingAccounts ? remainingAccounts : [])
    .signers([recipient])
    .rpc();
}