import { Connection, GetProgramAccountsFilter, PublicKey, Transaction } from "@solana/web3.js";
import { BN } from "@coral-xyz/anchor";
import {
  ConfigData,
  FLEXIBLE_STAKING_ACCOUNT_LEN,
  FlexibleStakingAccountData,
  FlexibleStakingAccountDataV2,
  FlexibleVaultData,
  FlexibleVaultDataV2,
  FlexibleVaultParams,
  getFlexibleStakingAccountMintOffsetV2,
  getFlexibleStakingAccountUserOffset,
  getFlexibleStakingAccountUserOffsetV2,
  getLockupStakingAccountMintOffsetV2,
  getLockupStakingAccountOwnerOffset,
  getLockupStakingAccountOwnerOffsetV2,
  getLockupStakingAccountUnlockedOffset,
  getLockupStakingAccountUnlockedOffsetV2,
  getLockupVaultMintOffsetV2,
  getWithdrawPositionClaimedOffset,
  getWithdrawPositionClaimedOffsetV2,
  getWithdrawPositionMintOffsetV2,
  getWithdrawPositionUserOffset,
  getWithdrawPositionUserOffsetV2,
  LOCKUP_STAKING_ACCOUNT_LEN,
  LOCKUP_VAULT_LEN,
  LockupVaultData,
  LockupVaultParams,
  PDA,
  ProtocolV2StakingInfo,
  SCALE_FACTOR,
  SOLCAT_FLEXIBLE_STAKING_ACCOUNT_LEN,
  SOLCAT_LOCKUP_STAKING_ACCOUNT_LEN,
  SOLCAT_LOCKUP_VAULT_LEN,
  SOLCAT_STAKING_PROGRAM_ID,
  SOLCAT_WITHDRAW_POSITION_LEN,
  SolCatStakingProtocolInfo,
  STAKING_PROGRAM_ID_V2,
  TokenStakingInfo,
  UserFlexibleStakingInfo,
  UserLockupStakingInfo,
  WITHDRAW_POSITION_LEN,
  WithdrawPositionData,
  WithdrawPositionDataV2,
  WithdrawPositionInfo,
  WithdrawPositionInfoV2,
} from "..";
import { AccountFetcher, ParsableLockupVault, ParsableLockupStakingAccount } from "../fetcher";
import {
  createAssociatedTokenAccountInstruction,
  createTransferCheckedInstruction,
  getAccount,
  getAssociatedTokenAddressSync,
  getMint,
} from "@solana/spl-token";
import { SolCatContext } from "../contexts/solcat-context";
import { V2Context } from "../contexts/v2-context";

function validateV2Mint(usingV2: boolean, mint?: PublicKey) {
  if (usingV2 && !mint) {
    throw new Error("Mint is required for V2 functions");
  }
}

export class StakingClient {
  ctx: SolCatContext;
  v2Ctx: V2Context;
  public pda: PDA;
  public v2Pda: PDA;

  constructor(
    ctx: SolCatContext,
    v2Ctx: V2Context,
    pda: PDA,
    v2Pda: PDA
  ) {
    this.ctx = ctx;
    this.v2Ctx = v2Ctx;
    this.pda = pda;
    this.v2Pda = v2Pda;
  }

  public async depositFund(amount: BN, usingV2 = false, mint?: PublicKey) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    // Get config data
    const config = pda.config();
    const configInfo = await context.connection.getAccountInfo(config.key);
    if (!configInfo) {
      throw new Error("Config account not found");
    }

    let mintInfo;
    if (usingV2) {
      validateV2Mint(usingV2, mint);
      mintInfo = await getMint(context.connection, mint!);
      if (!mintInfo) {
        throw new Error("Mint not found");
      }
    } else {
      const configData = await context.program.account.config.fetch(config.key) as ConfigData;
      mintInfo = await getMint(context.connection, configData.mint);
    }

    const tx = new Transaction();
    const senderATA = getAssociatedTokenAddressSync(mintInfo.address, context.wallet.publicKey);
    const configATA = getAssociatedTokenAddressSync(mintInfo.address, config.key, true);
    const [senderATAInfo, configATAInfo] = await context.connection.getMultipleAccountsInfo([senderATA, configATA]);
    if (!senderATAInfo) {
      throw new Error("Sender associated token account not found");
    }
    const senderATAData = await getAccount(context.connection, senderATA);
    if ((new BN(senderATAData.amount.toString())).lt(amount)) {
      throw new Error("Insufficient token balance to deposit");
    }
    if (!configATAInfo) {
      tx.add(
        createAssociatedTokenAccountInstruction(
          context.wallet.publicKey,
          configATA,
          config.key,
          mintInfo.address,
        )
      );
    }
    tx.add(
      createTransferCheckedInstruction(
        senderATA,
        mintInfo.address,
        configATA,
        this.ctx.wallet.publicKey,
        BigInt(amount.toString()),
        mintInfo.decimals,
      )
    );
    return tx;
  }

  // Admin functions
  public async updateAdmin(newAdmin: PublicKey, usingV2 = false) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    const config = pda.config();
    const configData = await context.fetcher.getConfig(config.key, true, usingV2);
    if (!configData) {
      throw new Error("Config account not found");
    }

    if (context.wallet.publicKey.toBase58() !== configData.admin.toBase58()) {
      throw new Error("Unauthorized");
    }
    if (newAdmin.toBase58() === configData.admin.toBase58()) {
      throw new Error("New admin is the same as the current admin");
    }

    const tx = new Transaction();
    const updateAdminIx = await context.methods.updateAdmin({
      inputs: {
        newAdmin,
      },
    });
    if (!updateAdminIx.ix) {
      throw new Error("Failed to create updating admin instruction");
    }
    tx.add(updateAdminIx.ix);

    return tx;
  }

  public async withdrawFund(amount: BN, usingV2 = false, mint?: PublicKey) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    const config = pda.config();
    const configInfo = await context.connection.getAccountInfo(config.key);
    if (!configInfo) {
      throw new Error("Config account not found");
    }

    const configData = await context.program.account.config.fetch(config.key);
    if (context.wallet.publicKey.toBase58() !== configData.admin.toBase58()) {
      throw new Error("Unauthorized");
    }
    let mintInfo;
    if (usingV2) {
      validateV2Mint(usingV2, mint);
      mintInfo = await getMint(context.connection, mint!);
      if (!mintInfo) {
        throw new Error("Mint not found");
      }
    } else {
      mintInfo = await getMint(context.connection, (configData as ConfigData).mint);
    }

    const tx = new Transaction();
    const adminATA = getAssociatedTokenAddressSync(mintInfo.address, context.wallet.publicKey);
    const configATA = getAssociatedTokenAddressSync(mintInfo.address, config.key, true);
    const [adminATAInfo, configATAInfo] = await context.connection.getMultipleAccountsInfo([adminATA, configATA]);
    if (!adminATAInfo) {
      tx.add(
        createAssociatedTokenAccountInstruction(
          context.wallet.publicKey,
          adminATA,
          context.wallet.publicKey,
          mintInfo.address,
        )
      );
    }
    if (!configATAInfo) {
      throw new Error("Config associated token account not found");
    }
    const configATAData = await getAccount(context.connection, configATA);
    if ((new BN(configATAData.amount.toString())).lt(amount)) {
      throw new Error("Insufficient token balance to withdraw");
    }
    const withdrawFundIx = usingV2
      ? await context.methods.withdrawFund({
        accounts: {
          mint: mintInfo.address,
        },
        inputs: {
          amount,
        }
      })
      : await context.methods.withdrawFund({
        inputs: {
          amount,
        },
      });
    if (!withdrawFundIx.ix) {
      throw new Error("Failed to create withdrawing fund instruction");
    }
    tx.add(withdrawFundIx.ix);
    return tx;
  }

  public async withdrawToken(amount: BN, mint: PublicKey) {
    const context = this.v2Ctx;
    const pda = this.v2Pda;

    const config = pda.config();
    const treasury = pda.treasury();
    const configInfo = await context.connection.getAccountInfo(config.key);
    if (!configInfo) {
      throw new Error("Config account not found");
    }

    const configData = await context.program.account.config.fetch(config.key);
    if (context.wallet.publicKey.toBase58() !== configData.admin.toBase58()) {
      throw new Error("Unauthorized");
    }
    let mintInfo;
    mintInfo = await getMint(context.connection, mint!);
    if (!mintInfo) {
      throw new Error("Mint not found");
    }

    const tx = new Transaction();
    const adminATA = getAssociatedTokenAddressSync(mintInfo.address, context.wallet.publicKey);
    const treasuryATA = getAssociatedTokenAddressSync(mintInfo.address, treasury.key, true);
    const [adminATAInfo, treasuryATAInfo] = await context.connection.getMultipleAccountsInfo([adminATA, treasuryATA]);
    if (!adminATAInfo) {
      tx.add(
        createAssociatedTokenAccountInstruction(
          context.wallet.publicKey,
          adminATA,
          context.wallet.publicKey,
          mintInfo.address,
        )
      );
    }
    if (!treasuryATAInfo) {
      throw new Error("Treasury associated token account not found");
    }
    const treasuryATAData = await getAccount(context.connection, treasuryATA);
    if ((new BN(treasuryATAData.amount.toString())).lt(amount)) {
      throw new Error("Insufficient token balance to withdraw");
    }
    const withdrawTokenIx = await context.methods.withdrawRevenue({
      accounts: {
        mint: mintInfo.address,
      },
      inputs: {
        amount,
      }
    });
    if (!withdrawTokenIx.ix) {
      throw new Error("Failed to create withdrawing token instruction");
    }
    tx.add(withdrawTokenIx.ix);
    return tx;
  }

  public async updateLockupVault(
    index: number,
    apy: number,
    lockupDuration: BN,
    minStakeAmount: BN,
    usingV2 = false,
    mint?: PublicKey
  ) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    if (apy > 10000) {
      throw new Error("Invalid APY");
    }

    validateV2Mint(usingV2, mint);

    const config = pda.config();
    const lockupVault = usingV2 ? pda.lockupVaultV2(mint!, new BN(index)) : pda.lockupVault(new BN(index));
    const [configInfo, lockupVaultInfo] = await context.connection.getMultipleAccountsInfo([config.key, lockupVault.key]);
    if (!configInfo) {
      throw new Error("Config account not found");
    }
    const configData = await context.program.account.config.fetch(config.key);
    if (context.wallet.publicKey.toBase58() !== configData.admin.toBase58()) {
      throw new Error("Unauthorized");
    }
    if (!lockupVaultInfo) {
      throw new Error("Lockup vault is not initialized");
    }

    const tx = new Transaction();
    const updateLockupVaultIx = usingV2
      ? await context.methods.updateLockupVault({
        accounts: {
          mint: mint!,
          lockupVault: lockupVault.key,
        },
        inputs: {
          index,
          apy,
          lockupDuration,
          minStakeAmount,
        },
      })
      : await context.methods.updateLockupVault({
        inputs: {
          index,
          apy,
          lockupDuration,
          minStakeAmount,
        },
      });
    if (!updateLockupVaultIx.ix) {
      throw new Error("Failed to create updating lockup vault instruction");
    }
    tx.add(updateLockupVaultIx.ix);
    return tx;
  }

  public async updateFlexibleVault(rewardRate: BN, withdrawalDelayTime: BN, usingV2 = false, mint?: PublicKey) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    validateV2Mint(usingV2, mint);

    const config = pda.config();
    const flexibleVault = usingV2 ? pda.flexibleVaultV2(mint!) : pda.flexibleVault();
    const [configInfo, flexibleVaultInfo] = await context.connection.getMultipleAccountsInfo([config.key, flexibleVault.key]);
    if (!configInfo) {
      throw new Error("Config account not found");
    }
    const configData = await context.program.account.config.fetch(config.key);
    if (context.wallet.publicKey.toBase58() !== configData.admin.toBase58()) {
      throw new Error("Unauthorized");
    }
    if (!flexibleVaultInfo) {
      throw new Error("Flexible vault is not initialized");
    }

    const tx = new Transaction();
    const updateFlexibleVaultIx = usingV2
      ? await context.methods.updateFlexibleVault({
        accounts: {
          mint: mint!,
        },
        inputs: {
          rewardRate,
          withdrawalDelayTime,
        },
      })
      : await context.methods.updateFlexibleVault({
        inputs: {
          rewardRate,
          withdrawalDelayTime,
        },
      });
    if (!updateFlexibleVaultIx.ix) {
      throw new Error("Failed to create updating flexible vault instruction");
    }
    tx.add(updateFlexibleVaultIx.ix);
    return tx;
  }

  public async initTokenVaults(
    mint: PublicKey,
    flexibleVaultParam: FlexibleVaultParams | null,
    lockupVaultParams: LockupVaultParams[],
  ) {
    const context = this.v2Ctx;
    const pda = this.v2Pda;

    const config = pda.config();
    const configInfo = await context.connection.getAccountInfo(config.key);
    if (!configInfo) {
      throw new Error("Config account not found");
    }

    const configData = await context.program.account.config.fetch(config.key);
    if (context.wallet.publicKey.toBase58() !== configData.admin.toBase58()) {
      throw new Error("Unauthorized");
    }

    const mintInfo = await getMint(context.connection, mint);
    if (!mintInfo) {
      throw new Error("Mint not found");
    }

    const tx = new Transaction();
    const configATA = getAssociatedTokenAddressSync(mintInfo.address, config.key, true);
    const flexibleVault = pda.flexibleVaultV2(mint);
    const [configATAInfo, flexibleVaultInfo] = await context.connection.getMultipleAccountsInfo([
      configATA,
      flexibleVault.key
    ]);
    if (!configATAInfo) {
      tx.add(
        createAssociatedTokenAccountInstruction(
          context.wallet.publicKey,
          configATA,
          config.key,
          mintInfo.address,
        )
      );
    }
    if (flexibleVaultParam) {
      if (flexibleVaultInfo) {
        throw new Error("Flexible vault already initialized");
      }
      const initFlexibleVaultIx = await context.methods.initFlexibleVault({
        accounts: {
          admin: context.wallet.publicKey,
          mint,
          flexibleVault: flexibleVault.key,
        },
        inputs: {
          bump: flexibleVault.bump,
          rewardRate: flexibleVaultParam.rewardRate,
          withdrawDelayTime: flexibleVaultParam.withdrawalDelayTime,
        },
      });
      if (!initFlexibleVaultIx.ix) {
        throw new Error("Failed to create initializing flexible vault instruction");
      }
      tx.add(initFlexibleVaultIx.ix);
    }

    if (lockupVaultParams.length > 0) {
      const lockupVaults = lockupVaultParams.map(param => pda.lockupVaultV2(mint, new BN(param.index)));
      const lockVaultsInfo = await context.connection.getMultipleAccountsInfo(lockupVaults.map(vault => vault.key));
      for (let i = 0; i < lockupVaults.length; i++) {
        const vault = lockupVaults[i];
        const param = lockupVaultParams[i];
        if (lockVaultsInfo[i]) {
          throw new Error(`Lockup vault ${param.index} already initialized`);
        }
        const initLockupVaultIx = await context.methods.initLockupVault({
          accounts: {
            admin: context.wallet.publicKey,
            mint,
            lockupVault: vault.key,
          },
          inputs: {
            bump: vault.bump,
            index: param.index,
            apy: param.apy,
            lockupDuration: param.lockupDuration,
            minStakeAmount: param.minStakeAmount,
          },
        });
        if (!initLockupVaultIx.ix) {
          throw new Error("Failed to create initializing lockup vault instruction");
        }
        tx.add(initLockupVaultIx.ix);
      }
    }

    return tx;
  }

  // User functions
  public async stake(amount: BN, usingV2 = false, mint?: PublicKey) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    validateV2Mint(usingV2, mint);

    const tx = new Transaction();

    const config = pda.config();
    const configData = await (context as SolCatContext).program.account.config.fetch(config.key);
    const mintPubkey = usingV2 ? mint! : configData.mint;
    const associatedConfig = getAssociatedTokenAddressSync(mintPubkey, config.key, true);
    const flexibleAccount = usingV2
      ? pda.flexibleStakingAccountV2(mintPubkey, context.wallet.publicKey)
      : pda.flexibleStakingAccount(context.wallet.publicKey);
    const [
      associatedConfigInfo,
      flexibleAccountInfo
    ] = await context.connection.getMultipleAccountsInfo([associatedConfig, flexibleAccount.key]);
    if (usingV2 && !associatedConfigInfo) {
      tx.add(createAssociatedTokenAccountInstruction(
        context.wallet.publicKey,
        associatedConfig,
        config.key,
        mintPubkey,
      ))
    }
    if (!flexibleAccountInfo) {
      const initFlexibleAccountIx = usingV2
        ? await context.methods.initFlexibleAccount({
          accounts: {
            user: context.wallet.publicKey,
            mint: mint!,
          },
          inputs: {
            bump: flexibleAccount.bump,
          },
        })
        : await context.methods.initFlexibleAccount({
          accounts: {
            user: context.wallet.publicKey,
          },
          inputs: {
            bump: flexibleAccount.bump,
          },
        });
      if (!initFlexibleAccountIx.ix) {
        throw new Error("Failed to create initializing flexible account instruction");
      }
      tx.add(initFlexibleAccountIx.ix);
    }

    const stakeIx = usingV2
      ? await context.methods.stake({
        accounts: {
          user: context.wallet.publicKey,
          mint: mintPubkey,
        },
        inputs: {
          amount,
        },
      })
      : await context.methods.stake({
        accounts: {
          user: context.wallet.publicKey,
        },
        inputs: {
          amount,
        },
      });
    if (!stakeIx.ix) {
      throw new Error("Failed to create staking instruction");
    }
    tx.add(stakeIx.ix);

    return tx;
  }

  public async withdraw(amount: BN, usingV2 = false, mint?: PublicKey) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    const tx = new Transaction();
    let withdrawIndex = new BN(0);

    validateV2Mint(usingV2, mint);

    const userAccount = usingV2
      ? pda.userV2(mint!, context.wallet.publicKey)
      : pda.user(context.wallet.publicKey);
    const userAccountInfo = await context.connection.getAccountInfo(userAccount.key);

    if (!userAccountInfo) {
      const initUserIx = usingV2
        ? await (context as V2Context).methods.initUser({
          accounts: { user: context.wallet.publicKey, mint: mint! },
          inputs: { bump: userAccount.bump }
        })
        : await (context as SolCatContext).methods.initUser({
          accounts: { user: context.wallet.publicKey },
          inputs: { bump: userAccount.bump }
        });

      if (!initUserIx.ix) {
        throw new Error("Failed to create initializing user instruction");
      }
      tx.add(initUserIx.ix);
    } else {
      withdrawIndex = new BN((await context.program.account.user.fetch(userAccount.key)).latestWithdrawIndex);
    }

    const flexibleAccount = usingV2
      ? pda.flexibleStakingAccountV2(mint!, context.wallet.publicKey)
      : pda.flexibleStakingAccount(context.wallet.publicKey);
    const config = pda.config();
    const withdrawPosition = usingV2
      ? pda.withdrawPositionV2(mint!, context.wallet.publicKey, new BN(withdrawIndex))
      : pda.withdrawPosition(context.wallet.publicKey, new BN(withdrawIndex));
    const [configInfo, flexibleAccountInfo, withdrawPositionInfo] = await context.connection.getMultipleAccountsInfo([config.key, flexibleAccount.key, withdrawPosition.key]);
    if (!configInfo) {
      throw new Error("Config account not found");
    }
    if (!flexibleAccountInfo) {
      throw new Error("Ensure user has staked before withdrawing");
    }
    if (withdrawPositionInfo) {
      throw new Error("Withdraw position index is already taken");
    }

    const configData = await (context as SolCatContext).program.account.config.fetch(config.key);
    const mintPubkey = usingV2 ? mint! : configData.mint;
    const userATA = getAssociatedTokenAddressSync(mintPubkey, context.wallet.publicKey);

    const userATAInfo = await context.connection.getAccountInfo(userATA);
    if (!userATAInfo) {
      tx.add(
        createAssociatedTokenAccountInstruction(
          context.wallet.publicKey,
          userATA,
          context.wallet.publicKey,
          mintPubkey,
        )
      );
    }

    const withdrawIx = usingV2
      ? await (context as V2Context).methods.withdraw({
        accounts: {
          user: context.wallet.publicKey,
          withdrawPosition: withdrawPosition.key,
          mint: mint!,
        },
        inputs: {
          withdrawPositionBump: withdrawPosition.bump,
          index: withdrawIndex.toNumber(),
          amount,
        },
      })
      : await (context as SolCatContext).methods.withdraw({
        accounts: {
          user: context.wallet.publicKey,
        },
        inputs: {
          withdrawPositionBump: withdrawPosition.bump,
          index: withdrawIndex.toNumber(),
          amount,
        },
      });
    if (!withdrawIx.ix) {
      throw new Error("Failed to create withdrawing instruction");
    }
    tx.add(withdrawIx.ix);

    return tx;
  }

  public async withdrawAfterDelay(index: number, usingV2 = false, mint?: PublicKey) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    validateV2Mint(usingV2, mint);

    const tx = new Transaction();

    const config = pda.config();
    const withdrawPosition = usingV2
      ? pda.withdrawPositionV2(mint!, context.wallet.publicKey, new BN(index))
      : pda.withdrawPosition(context.wallet.publicKey, new BN(index));
    const [configInfo, withdrawPositionInfo] = await context.connection.getMultipleAccountsInfo([
      config.key,
      withdrawPosition.key
    ]);
    if (!configInfo) {
      throw new Error("Config account not found");
    }
    if (!withdrawPositionInfo) {
      throw new Error("Ensure create withdraw position before getting delay amount");
    }

    const configData = await (context as SolCatContext).program.account.config.fetch(config.key);
    const mintPubkey = usingV2 ? mint! : configData.mint;
    const userATA = getAssociatedTokenAddressSync(mintPubkey, context.wallet.publicKey);
    const userATAInfo = await context.connection.getAccountInfo(userATA);
    if (!userATAInfo) {
      tx.add(
        createAssociatedTokenAccountInstruction(
          context.wallet.publicKey,
          userATA,
          context.wallet.publicKey,
          mintPubkey,
        )
      );
    }

    const withdrawDelayIx = usingV2
      ? await context.methods.withdrawDelay({
        accounts: {
          user: context.wallet.publicKey,
          withdrawPosition: withdrawPosition.key,
          mint: mint!,
        },
        inputs: {
          index,
        },
      })
      : await context.methods.withdrawDelay({
        accounts: {
          user: context.wallet.publicKey,
        },
        inputs: {
          index,
        },
      });
    if (!withdrawDelayIx.ix) {
      throw new Error("Failed to create withdrawing after delay instruction");
    }
    tx.add(withdrawDelayIx.ix);

    return tx;
  }

  public async claim(usingV2 = false, mint?: PublicKey) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    validateV2Mint(usingV2, mint);

    const tx = new Transaction();

    const config = pda.config();
    const configInfo = await (context as SolCatContext).connection.getAccountInfo(config.key);
    if (!configInfo) {
      throw new Error("Config account not found");
    }
    const configData = await (context as SolCatContext).program.account.config.fetch(config.key);
    const mintPubkey = usingV2 ? mint! : configData.mint;
    const userATA = getAssociatedTokenAddressSync(mintPubkey, context.wallet.publicKey);
    const userATAInfo = await context.connection.getAccountInfo(userATA);
    if (!userATAInfo) {
      tx.add(
        createAssociatedTokenAccountInstruction(
          context.wallet.publicKey,
          userATA,
          context.wallet.publicKey,
          mintPubkey,
        )
      );
    }

    const treasury = pda.treasury();
    const treasuryATA = getAssociatedTokenAddressSync(mintPubkey, treasury.key, true);

    const flexibleAccount = usingV2
      ? pda.flexibleStakingAccountV2(mint!, context.wallet.publicKey)
      : pda.flexibleStakingAccount(context.wallet.publicKey);
    const [flexibleAccountInfo, treasuryATAInfo] = await context.connection.getMultipleAccountsInfo([
      flexibleAccount.key,
      treasuryATA
    ]);
    if (!flexibleAccountInfo) {
      throw new Error("Ensure user has staked before claiming");
    }
    if (usingV2 && !treasuryATAInfo) {
      tx.add(
        createAssociatedTokenAccountInstruction(
          context.wallet.publicKey,
          treasuryATA,
          treasury.key,
          mint!,
        )
      )
    }

    const claimIx = usingV2
      ? await context.methods.claim({
        accounts: {
          user: context.wallet.publicKey,
          mint: mint!,
        },
      })
      : await context.methods.claim({
        accounts: {
          user: context.wallet.publicKey,
        },
      });
    if (!claimIx.ix) {
      throw new Error("Failed to create claiming instruction");
    }
    tx.add(claimIx.ix);

    return tx;
  }

  public async lockup(amount: BN, vaultIndex: number, usingV2 = false, mint?: PublicKey) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    validateV2Mint(usingV2, mint);

    const tx = new Transaction();
    let lockupIndex = new BN(0);

    const config = pda.config();
    const userAccount = usingV2
      ? pda.userV2(mint!, context.wallet.publicKey)
      : pda.user(context.wallet.publicKey);
    const [configInfo, userAccountInfo] = await context.connection.getMultipleAccountsInfo([config.key, userAccount.key]);
    if (!configInfo) {
      throw new Error("Config account not found");
    }
    if (!userAccountInfo) {
      const initUserIx = usingV2
        ? await (context as V2Context).methods.initUser({
          accounts: { user: context.wallet.publicKey, mint: mint! },
          inputs: { bump: userAccount.bump },
        })
        : await (context as SolCatContext).methods.initUser({
          accounts: { user: context.wallet.publicKey },
          inputs: { bump: userAccount.bump },
        });
      if (!initUserIx.ix) {
        throw new Error("Failed to create initializing user instruction");
      }
      tx.add(initUserIx.ix);
    } else {
      lockupIndex = new BN((await context.program.account.user.fetch(userAccount.key)).latestLockupIndex);
    }
    const configData = await context.program.account.config.fetch(config.key);
    const configATA = usingV2
      ? getAssociatedTokenAddressSync(mint!, config.key, true)
      : getAssociatedTokenAddressSync((configData as ConfigData).mint, config.key, true);

    const userLockupStakingAccount = usingV2
      ? pda.lockupStakingAccount(
        context.wallet.publicKey,
        new BN(lockupIndex),
        new BN(vaultIndex),
        mint!
      )
      : pda.lockupStakingAccount(
        context.wallet.publicKey,
        new BN(lockupIndex),
        new BN(vaultIndex)
      );
    const [
      configATAInfo,
      userLockupStakingAccountInfo
    ] = await context.connection.getMultipleAccountsInfo([
      configATA,
      userLockupStakingAccount.key,
    ]);
    if (!configATAInfo) {
      tx.add(createAssociatedTokenAccountInstruction(
        context.wallet.publicKey,
        configATA,
        config.key,
        mint!,
      ));
    }
    if (userLockupStakingAccountInfo) {
      throw new Error("Lockup staking account is already initialized");
    }

    const lockUpIx = usingV2
      ? await context.methods.lockup({
        accounts: {
          user: context.wallet.publicKey,
          lockupVault: pda.lockupVaultV2(mint!, new BN(vaultIndex)).key,
          mint: mint!,
        },
        inputs: {
          userLockupStakingAccountBump: userLockupStakingAccount.bump,
          index: lockupIndex.toNumber(),
          vaultIndex,
          amount,
        },
      })
      : await context.methods.lockup({
        accounts: {
          user: context.wallet.publicKey,
        },
        inputs: {
          userLockupStakingAccountBump: userLockupStakingAccount.bump,
          index: lockupIndex.toNumber(),
          vaultIndex,
          amount,
        },
      });
    if (!lockUpIx.ix) {
      throw new Error("Failed to create locking up instruction");
    }
    tx.add(lockUpIx.ix);

    return tx;
  }

  public async unlock(index: number, vaultIndex: number, usingV2 = false, mint?: PublicKey) {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    validateV2Mint(usingV2, mint);

    const tx = new Transaction();

    const config = pda.config();
    const configInfo = await context.connection.getAccountInfo(config.key);
    if (!configInfo) {
      throw new Error("Config account not found");
    }
    const configData = await (context as SolCatContext).program.account.config.fetch(config.key);
    const mintPubkey = usingV2 ? mint! : configData.mint;
    const userATA = getAssociatedTokenAddressSync(mintPubkey, context.wallet.publicKey);
    const userATAInfo = await context.connection.getAccountInfo(userATA);
    if (!userATAInfo) {
      tx.add(
        createAssociatedTokenAccountInstruction(
          context.wallet.publicKey,
          userATA,
          context.wallet.publicKey,
          mintPubkey,
        )
      );
    }

    if (usingV2) {
      const treasury = this.v2Pda.treasury();
      const treasuryATA = getAssociatedTokenAddressSync(mint!, treasury.key, true);

      const treasuryATAInfo = await context.connection.getAccountInfo(treasuryATA);
      if (!treasuryATAInfo) {
        tx.add(
          createAssociatedTokenAccountInstruction(
            context.wallet.publicKey,
            treasuryATA,
            treasury.key,
            mint!,
          )
        )
      }
    }

    const unlockIx = usingV2
      ? await context.methods.unlock({
        accounts: {
          user: context.wallet.publicKey,
          mint: mint!,
        },
        inputs: {
          index,
          vaultIndex,
        },
      })
      : await context.methods.unlock({
        accounts: {
          user: context.wallet.publicKey,
        },
        inputs: {
          index,
          vaultIndex,
        },
      });
    if (!unlockIx.ix) {
      throw new Error("Failed to create unlocking instruction");
    }
    tx.add(unlockIx.ix);

    return tx;
  }

  // Getter functions
  public async getUserFlexibleStakingInfo(
    filter: {
      withdrew?: boolean,
      usingV2: boolean,
      mint?: PublicKey,
      user?: PublicKey
    }
  ): Promise<UserFlexibleStakingInfo> {
    const context = filter.usingV2 ? this.v2Ctx : this.ctx;
    const pda = filter.usingV2 ? this.v2Pda : this.pda;

    const configData = await context.program.account.config.fetch(pda.config().key);

    const flexibleStakingAccountsData = await this.getAllUserFlexibleStakingAccount(
      configData.admin,
      filter.user || context.wallet.publicKey,
      filter.usingV2,
      filter.mint
    );

    if (flexibleStakingAccountsData.length === 0) {
      return {};
    }

    const flexibleVaults = filter.usingV2
      ? (flexibleStakingAccountsData as FlexibleStakingAccountDataV2[]).map((account) => {
        return pda.flexibleVaultV2(account.mint);
      })
      : [pda.flexibleVault()];

    const flexibleVaultsData = await Promise.all(flexibleVaults.map(async (vault) => {
      return context.program.account.flexibleVault.fetch(vault.key);
    }));

    const withdrawPositionsData = await this.getAllUserWithdrawPosition(
      filter.user || context.wallet.publicKey,
      filter.usingV2,
      filter.withdrew,
      filter.mint
    );

    const now = Math.floor(Date.now() / 1000);
    let userFlexibleStakingInfo: UserFlexibleStakingInfo = {};
    flexibleVaultsData.map((vault, index) => {
      const pendingReward = this.getFlexiblePendingReward(now, vault, flexibleStakingAccountsData[index]);
      const totalFlexibleStaked = Number(BigInt(vault.totalSupply.toString()));
      const apr = calculateAPR(totalFlexibleStaked, Number(vault.rewardRate.toString()));
      if (filter.usingV2) {
        if (withdrawPositionsData.length === 0) {
          userFlexibleStakingInfo[(vault as FlexibleVaultDataV2).mint.toBase58()] = {
            totalStakedAmount: flexibleStakingAccountsData[index].balance.toString(),
            currentApr: apr,
            withdrawDelayTime: vault.withdrawalDelayTime.toString(),
            pendingReward: pendingReward.toString(),
            unstakePending: "0",
            redeemPending: "0",
            withdrawPositions: [],
          };
        } else {
          const {
            unstakePending,
            redeemPending,
            withdrawPositions
          } = this.getFlexibleStakingInfo(
            filter.usingV2,
            now,
            vault,
            withdrawPositionsData,
            (vault as FlexibleVaultDataV2).mint
          );
          userFlexibleStakingInfo[(vault as FlexibleVaultDataV2).mint.toBase58()] = {
            totalStakedAmount: flexibleStakingAccountsData[index].balance.toString(),
            currentApr: apr,
            withdrawDelayTime: vault.withdrawalDelayTime.toString(),
            pendingReward: pendingReward.toString(),
            unstakePending: unstakePending.toString(),
            redeemPending: redeemPending.toString(),
            withdrawPositions,
          };
        }
      } else {
        if (withdrawPositionsData.length === 0) {
          userFlexibleStakingInfo[(configData as ConfigData).mint.toBase58()] = {
            totalStakedAmount: flexibleStakingAccountsData[index].balance.toString(),
            currentApr: apr,
            withdrawDelayTime: vault.withdrawalDelayTime.toString(),
            pendingReward: pendingReward.toString(),
            unstakePending: "0",
            redeemPending: "0",
            withdrawPositions: [],
          };
        } else {
          const {
            unstakePending,
            redeemPending,
            withdrawPositions
          } = this.getFlexibleStakingInfo(filter.usingV2, now, vault, withdrawPositionsData);
          userFlexibleStakingInfo[(configData as ConfigData).mint.toBase58()] = {
            totalStakedAmount: flexibleStakingAccountsData[index].balance.toString(),
            currentApr: apr,
            withdrawDelayTime: vault.withdrawalDelayTime.toString(),
            pendingReward: pendingReward.toString(),
            unstakePending: unstakePending.toString(),
            redeemPending: redeemPending.toString(),
            withdrawPositions,
          }
        }
      }
    });

    return userFlexibleStakingInfo;
  }

  public async getUserLockupStakingInfo(
    filter: {
      unlocked?: boolean,
      usingV2: boolean,
      mint?: PublicKey,
      user?: PublicKey
    }
  ): Promise<UserLockupStakingInfo> {
    const context = filter.usingV2 ? this.v2Ctx : this.ctx;
    const pda = filter.usingV2 ? this.v2Pda : this.pda;

    const configData = await context.program.account.config.fetch(pda.config().key);

    const lockupVaultsData = await this.getAllLockupVaultProtocol(
      filter.usingV2,
      filter.usingV2
        ? filter.mint
        : (configData as ConfigData).mint
    );
    if (lockupVaultsData.length === 0) {
      return {};
    }

    let mintList: string[] = [];
    if (filter.usingV2 && !filter.mint) {
      mintList = [...new Set(lockupVaultsData.map(vault => vault.mint))];
    } else {
      mintList = filter.mint ? [filter.mint.toBase58()] : [(configData as ConfigData).mint.toBase58()]
    }
    if (mintList.length === 0) {
      return {};
    }

    const lockupStakingAccounts = await this.getAllUserLockupStakingAccount(
      filter.user || context.wallet.publicKey,
      filter.usingV2,
      filter.unlocked,
      filter.mint
    );
    if (lockupStakingAccounts.length === 0) {
      return {};
    }

    const userLockStakingInfo: UserLockupStakingInfo = {};
    mintList.map((mint) => {
      let totalStakedAmount = new BN(0);
      let maturedAmount = new BN(0);

      const lockupVaults = filter.usingV2
        ? lockupVaultsData.filter((vault) => vault.mint === mint)
        : lockupVaultsData;
      let lockupStakingAccountsByVault: any = [];

      lockupVaults.map((vault) => {
        const stakingAccounts = lockupStakingAccounts.filter((account) => {
          return account.lockupVault.toBase58() === vault.lockupVault
        })
        stakingAccounts.map((account) => {
          totalStakedAmount = totalStakedAmount.add(account.lockedAmount);
          const maturedTime = account.lockedAt.add(new BN(vault.lockupDuration.toString()));
          if (maturedTime.lte(new BN(Math.floor(Date.now() / 1000)))) {
            maturedAmount = maturedAmount.add(account.lockedAmount);
          }
          lockupStakingAccountsByVault.push({
            bump: account.bump,
            index: account.index,
            lockupVault: account.lockupVault.toBase58(),
            lockupVaultIndex: vault.index,
            lockupDuration: vault.lockupDuration.toString(),
            committedApy: account.committedApy,
            lockedAmount: account.lockedAmount.toString(),
            lockedAt: account.lockedAt.toString(),
            unlocked: account.unlocked,
          });
        });
      });

      lockupStakingAccountsByVault.sort((a: any, b: any) => a.index - b.index);

      userLockStakingInfo[mint] = {
        totalLockedAmount: totalStakedAmount.toString(),
        maturedAmount: maturedAmount.toString(),
        lockupStakingAccounts: lockupStakingAccountsByVault,
      };
    })

    return userLockStakingInfo;
  }

  private async getAllUserFlexibleStakingAccount(
    admin: PublicKey,
    user: PublicKey,
    usingV2: boolean,
    mint?: PublicKey
  ): Promise<FlexibleStakingAccountData[] | FlexibleStakingAccountDataV2[]> {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const pda = usingV2 ? this.v2Pda : this.pda;

    const filters: any[] = [
      { dataSize: usingV2 ? FLEXIBLE_STAKING_ACCOUNT_LEN : SOLCAT_FLEXIBLE_STAKING_ACCOUNT_LEN },
      {
        memcmp: {
          offset: usingV2 ? getFlexibleStakingAccountUserOffsetV2() : getFlexibleStakingAccountUserOffset(),
          bytes: user.toBase58()
        }
      }];
    if (usingV2 && mint) {
      filters.push({
        memcmp: {
          offset: getFlexibleStakingAccountMintOffsetV2(),
          bytes: mint.toBase58()
        }
      });
    }

    let flexibleStakingAccounts = await context.connection.getProgramAccounts(
      context.program.programId,
      { filters }
    );

    if (flexibleStakingAccounts.length === 0) {
      return [];
    }

    if (admin.toBase58() == user.toBase58() && !usingV2) {
      const config = pda.config().key;
      flexibleStakingAccounts = flexibleStakingAccounts.filter((account) => account.pubkey.toBase58() !== config.toBase58());
    }

    return await Promise.all(flexibleStakingAccounts.map(async (account) => {
      return await context.program.account.flexibleStakingAccount
        .fetch(account.pubkey);
    }));
  }

  private async getAllUserWithdrawPosition(
    user: PublicKey,
    usingV2: boolean,
    withdrew?: boolean,
    mint?: PublicKey
  ): Promise<any[]> {
    const context = usingV2 ? this.v2Ctx : this.ctx;

    let filters: any[] = [
      { dataSize: usingV2 ? WITHDRAW_POSITION_LEN : SOLCAT_WITHDRAW_POSITION_LEN },
      {
        memcmp: {
          offset: usingV2 ? getWithdrawPositionUserOffsetV2() : getWithdrawPositionUserOffset(),
          bytes: user.toBase58()
        }
      }];
    if (withdrew !== undefined) {
      filters.push({
        memcmp: {
          offset: usingV2 ? getWithdrawPositionClaimedOffsetV2() : getWithdrawPositionClaimedOffset(),
          bytes: withdrew ? "2" : "1",
        },
      });
    }
    if (usingV2 && mint) {
      filters.push({
        memcmp: {
          offset: getWithdrawPositionMintOffsetV2(),
          bytes: mint.toBase58()
        }
      });
    }

    const withdrawPositionAccounts = await context.connection.getProgramAccounts(
      context.program.programId,
      { filters }
    );
    if (withdrawPositionAccounts.length === 0) {
      return [];
    }

    let withdrawPositions: any[] = [];

    await Promise.all(withdrawPositionAccounts.map(async (account) => {
      const data = await context.program.account.withdrawPosition.fetch(account.pubkey)
      withdrawPositions.push(data);
    }));

    return withdrawPositions
  }

  private getFlexiblePendingReward(
    now: number,
    vault: FlexibleVaultData | FlexibleVaultDataV2,
    flexibleStakingAccount: FlexibleStakingAccountData | FlexibleStakingAccountDataV2
  ): BN {
    const timePassed = new BN(now).sub(vault.lastUpdateTime);

    let rewardPerToken = new BN(0);
    if (vault.totalSupply.eq(new BN(0))) {
      rewardPerToken = vault.rewardPerTokenStored;
    } else {
      rewardPerToken = vault.rewardPerTokenStored.add(
        vault.rewardRate
          .mul(timePassed)
          .mul(SCALE_FACTOR)
          .div(vault.totalSupply)
      );
    }

    return flexibleStakingAccount.balance
      .mul(
        rewardPerToken.sub(flexibleStakingAccount.userRewardPerTokenPaid)
      )
      .div(SCALE_FACTOR)
      .add(flexibleStakingAccount.reward);
  }

  private getFlexibleStakingInfo(
    usingV2: boolean,
    now: number,
    vault: FlexibleVaultData | FlexibleVaultDataV2,
    withdrawPositionsData: WithdrawPositionData[] | WithdrawPositionDataV2[],
    mint?: PublicKey
  ): {
    unstakePending: BN,
    redeemPending: BN,
    withdrawPositions: (WithdrawPositionInfo | WithdrawPositionInfoV2)[]
  } {
    let unstakePending = new BN(0);
    let redeemPending = new BN(0);

    if (withdrawPositionsData.length === 0) {
      return {
        unstakePending,
        redeemPending,
        withdrawPositions: []
      };
    }

    const withdrawPositions = withdrawPositionsData
      .filter((withdrawPosition) => {
        if (usingV2 && mint) {
          return (withdrawPosition as WithdrawPositionDataV2).mint.toBase58() === mint.toBase58();
        }
        return true;
      }).map((withdrawPosition) => {
        const releaseTime = withdrawPosition.startDelayTime.add(
          vault.withdrawalDelayTime
        );
        if (releaseTime.lt(new BN(now))) {
          redeemPending = redeemPending.add(withdrawPosition.amount);
        } else {
          unstakePending = unstakePending.add(withdrawPosition.amount);
        }
        return usingV2
          ? {
            mint: (withdrawPosition as WithdrawPositionDataV2).mint.toBase58(),
            positionIndex: (withdrawPosition as WithdrawPositionDataV2).positionIndex,
            amount: withdrawPosition.amount.toString(),
            startDelayTime: withdrawPosition.startDelayTime.toString(),
            unlockTime: releaseTime.toString(),
            claimed: withdrawPosition.claimed,
          }
          : {
            index: (withdrawPosition as WithdrawPositionData).index,
            amount: withdrawPosition.amount.toString(),
            startDelayTime: withdrawPosition.startDelayTime.toString(),
            unlockTime: releaseTime.toString(),
            claimed: withdrawPosition.claimed,
          };
      });

    // sort withdraw positions by index
    usingV2
      ? withdrawPositions.sort((a: any, b: any) => a.positionIndex - b.positionIndex)
      : withdrawPositions.sort((a: any, b: any) => a.index - b.index);
    return {
      unstakePending,
      redeemPending,
      withdrawPositions,
    }
  }

  private async getAllLockupVaultProtocol(usingV2: boolean, mint?: PublicKey): Promise<any[]> {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const fetcher = new AccountFetcher(context.connection);
    const lockupVaultFilter: GetProgramAccountsFilter[] = [{
      dataSize: usingV2 ? LOCKUP_VAULT_LEN : SOLCAT_LOCKUP_VAULT_LEN
    }];
    if (usingV2 && mint) {
      lockupVaultFilter.push({
        memcmp: {
          offset: getLockupVaultMintOffsetV2(),
          bytes: mint!.toBase58()
        }
      })
    }
    const lockupVaultAccounts = await context.connection.getProgramAccounts(
      context.program.programId,
      {
        filters: lockupVaultFilter,
      }
    );
    if (lockupVaultAccounts.length === 0) {
      return [];
    }

    return (await Promise.all(lockupVaultAccounts.map(async (info) => {
      const lockupVaultInfo: any = fetcher.parseAccountData(
        usingV2,
        info?.account?.data,
        ParsableLockupVault
      );
      if (lockupVaultInfo) {
        return {
          lockupVault: info.pubkey.toBase58(),
          bump: lockupVaultInfo.bump,
          mint: usingV2 ? lockupVaultInfo.mint.toBase58() : mint,
          index: lockupVaultInfo.index,
          apy: lockupVaultInfo.apy,
          lockupDuration: lockupVaultInfo.lockupDuration.toString(),
          minStakeAmount: lockupVaultInfo.minStakeAmount.toString(),
          totalRewardClaimed: lockupVaultInfo.totalRewardClaimed.toString(),
        }
      } else {
        return null;
      }
    })));
  }

  private async getAllUserLockupStakingAccount(
    user: PublicKey,
    usingV2: boolean,
    unlocked?: boolean,
    mint?: PublicKey
  ): Promise<any[]> {
    const context = usingV2 ? this.v2Ctx : this.ctx;
    const fetcher = new AccountFetcher(context.connection);

    const filters: any[] = [
      { dataSize: usingV2 ? LOCKUP_STAKING_ACCOUNT_LEN : SOLCAT_LOCKUP_STAKING_ACCOUNT_LEN },
      {
        memcmp: {
          offset: usingV2
            ? getLockupStakingAccountOwnerOffsetV2()
            : getLockupStakingAccountOwnerOffset(),
          bytes: user.toBase58()
        }
      },
    ];
    if (unlocked !== undefined) {
      filters.push({
        memcmp: {
          offset: usingV2 ? getLockupStakingAccountUnlockedOffsetV2() : getLockupStakingAccountUnlockedOffset(),
          bytes: unlocked ? "2" : "1",
        },
      });
    }
    if (usingV2 && mint) {
      filters.push({
        memcmp: {
          offset: getLockupStakingAccountMintOffsetV2(),
          bytes: mint.toBase58()
        }
      });
    }
    const lockupStakingAccounts = await context.connection.getProgramAccounts(
      context.program.programId,
      { filters }
    );
    if (lockupStakingAccounts.length === 0) {
      return [];
    }

    return (await Promise.all(lockupStakingAccounts.map(async (info) => {
      const lockupStakingAccountInfo: any = fetcher.parseAccountData(
        usingV2,
        info?.account?.data,
        ParsableLockupStakingAccount
      );
      if (lockupStakingAccountInfo) {
        return {
          bump: lockupStakingAccountInfo.bump,
          index: lockupStakingAccountInfo.index,
          lockupVault: lockupStakingAccountInfo.lockupVault,
          committedApy: lockupStakingAccountInfo.committedApy,
          lockedAmount: lockupStakingAccountInfo.lockedAmount,
          lockedAt: lockupStakingAccountInfo.lockedAt,
          unlocked: lockupStakingAccountInfo.unlocked,
        }
      } else {
        return null;
      }
    })));
  }
}

export async function getProtocolInfo(
  connection: Connection,
  usingV2 = false,
  mints: PublicKey[]
): Promise<SolCatStakingProtocolInfo | ProtocolV2StakingInfo> {
  if (usingV2 && mints?.length === 0) {
    throw new Error("Mints array is empty, please provide for V2 protocol");
  }
  const programId = usingV2 ? STAKING_PROGRAM_ID_V2 : SOLCAT_STAKING_PROGRAM_ID;
  const pda = new PDA(programId);
  const fetcher = new AccountFetcher(connection);
  const config = pda.config();
  const configData = await fetcher.getConfig(config.key, true, usingV2);
  if (!configData) {
    throw new Error("Config account not found");
  }

  if (usingV2) {
    const flexibleVaults = mints.map(mint => pda.flexibleVaultV2(mint));
    let stakingInfo: TokenStakingInfo = {};

    await Promise.all(flexibleVaults.map(async (flexibleVault) => {
      const flexibleVaultInfo = (await fetcher.getFlexibleVault(flexibleVault.key, true, usingV2)) as FlexibleVaultDataV2;
      if (!flexibleVaultInfo) {
        throw new Error("Flexible vault account not found");
      }

      const lockupVaultAccounts = await connection.getProgramAccounts(
        programId,
        {
          filters: [
            {
              dataSize: LOCKUP_VAULT_LEN
            },
            {
              memcmp: {
                offset: getLockupVaultMintOffsetV2(),
                bytes: flexibleVaultInfo.mint.toBase58()
              }
            }
          ],
        }
      );
      if (lockupVaultAccounts.length === 0) {
        throw new Error("Lockup vault account not found");
      }
      const lockupVaultsInfo: LockupVaultData[] = await Promise.all(lockupVaultAccounts.map(async (info) => {
        const lockupVaultInfo = fetcher.parseAccountData(usingV2, info?.account?.data, ParsableLockupVault);
        return lockupVaultInfo as LockupVaultData;
      }));
      // sort lockup vaults info by index
      lockupVaultsInfo.sort((a: any, b: any) => a.index - b.index);

      const totalFlexibleStaked = Number(BigInt(flexibleVaultInfo.totalSupply.toString()));
      const apr = calculateAPR(totalFlexibleStaked, Number(flexibleVaultInfo.rewardRate.toString()));

      stakingInfo[flexibleVaultInfo.mint.toBase58()] = {
        flexibleRewardRate: flexibleVaultInfo.rewardRate.toString(),
        flexibleApr: apr,
        flexibleWithdrawalDelayTime: flexibleVaultInfo.withdrawalDelayTime.toString(),
        ...lockupVaultsInfo.reduce<{
          lockupApy: number[];
          lockupDuration: string[];
          minLockupAmount: string[];
        }>((acc, vault) => {
          if (vault) {
            acc.lockupApy.push(Number(vault.apy));
            acc.lockupDuration.push(vault.lockupDuration.toString());
            acc.minLockupAmount.push(vault.minStakeAmount.toString());
          }
          return acc;
        }, {
          lockupApy: [],
          lockupDuration: [],
          minLockupAmount: [],
        }),
      }
    }));

    return {
      admin: configData.admin.toBase58(),
      stakingInfo,
    }
  } else {
    const flexibleVault = pda.flexibleVault();
    const flexibleVaultInfo = await fetcher.getFlexibleVault(flexibleVault.key, true, false);
    if (!flexibleVaultInfo) {
      throw new Error("Flexible vault account not found");
    }

    const lockupVaultAccounts = await connection.getProgramAccounts(
      programId,
      {
        filters: [
          {
            dataSize: SOLCAT_LOCKUP_VAULT_LEN
          },
        ],
      }
    );
    if (lockupVaultAccounts.length === 0) {
      throw new Error("Lockup vault account not found");
    }
    const lockupVaultsInfo: LockupVaultData[] = await Promise.all(lockupVaultAccounts.map(async (info) => {
      const lockupVaultInfo = fetcher.parseAccountData(false, info?.account?.data, ParsableLockupVault);
      return lockupVaultInfo as LockupVaultData;
    }));
    // sort lockup vaults info by index
    lockupVaultsInfo.sort((a: any, b: any) => a.index - b.index);

    const totalFlexibleStaked = Number(BigInt(flexibleVaultInfo.totalSupply.toString()));
    const apr = calculateAPR(totalFlexibleStaked, Number(flexibleVaultInfo.rewardRate.toString()));

    return {
      admin: configData.admin.toBase58(),
      mint: (configData as ConfigData).mint.toBase58(),
      flexibleRewardRate: flexibleVaultInfo.rewardRate.toString(),
      flexibleApr: apr,
      flexibleWithdrawalDelayTime: flexibleVaultInfo.withdrawalDelayTime.toString(),
      totalFlexibleStaked: totalFlexibleStaked.toString(),
      totalFlexibleReward: flexibleVaultInfo.totalRewardClaimed.toString(),
      ...lockupVaultsInfo.reduce<{
        lockupApy: number[];
        lockupDuration: string[];
        minLockupAmount: string[];
        totalLockupStaked: string[];
        totalLockupReward: string[];
      }>((acc, vault) => {
        if (vault) {
          acc.lockupApy.push(Number(vault.apy));
          acc.lockupDuration.push(vault.lockupDuration.toString());
          acc.minLockupAmount.push(vault.minStakeAmount.toString());
          acc.totalLockupStaked.push(vault.totalLockedAmount.toString());
          acc.totalLockupReward.push(vault.totalRewardClaimed.toString());
        }
        return acc;
      }, {
        lockupApy: [],
        lockupDuration: [],
        minLockupAmount: [],
        totalLockupStaked: [],
        totalLockupReward: [],
      }),
    };
  }
}

function calculateAPR(totalStaked: number, rewardPerSecond: number) {
  if (totalStaked === 0) return 0;
  const secondsInYear = 365 * 24 * 60 * 60;
  const annualReward = rewardPerSecond * secondsInYear;
  const apr = (annualReward / totalStaked) * 100;
  return apr;
}