import { Program } from "@coral-xyz/anchor";
import { PumpAmm } from "../types/pump_amm";
import {
  AccountInfo,
  Connection,
  PublicKey,
  SystemProgram,
  TransactionInstruction,
} from "@solana/web3.js";
import {
  globalConfigPda,
  globalVolumeAccumulatorPda,
  lpMintPda,
  poolPda,
  PUMP_AMM_PROGRAM_ID,
  pumpAmmEventAuthorityPda,
  userVolumeAccumulatorPda,
} from "./pda";
import {
  AccountLayout,
  ASSOCIATED_TOKEN_PROGRAM_ID,
  createAssociatedTokenAccountIdempotentInstruction,
  createCloseAccountInstruction,
  createSyncNativeInstruction,
  getAccount,
  getAssociatedTokenAddressSync,
  NATIVE_MINT,
  TOKEN_2022_PROGRAM_ID,
  TOKEN_PROGRAM_ID,
} from "@solana/spl-token";
import { depositToken0Internal } from "./deposit";
import { withdrawInternal } from "./withdraw";
import { buyBaseInputInternal, buyQuoteInputInternal } from "./buy";
import { sellBaseInputInternal, sellQuoteInputInternal } from "./sell";
import {
  CollectCoinCreatorFeeSolanaState,
  CommonSolanaState,
  CreatePoolSolanaState,
  DepositBaseResult,
  DepositQuoteResult,
  GlobalConfig,
  GlobalVolumeAccumulator,
  LiquidityAccounts,
  LiquiditySolanaState,
  Pool,
  SwapAccounts,
  SwapSolanaState,
  UserVolumeAccumulator,
  WithdrawResult,
} from "../types/sdk";
import { getPumpAmmProgram } from "./util";
import BN from "bn.js";
import { currentDayTokens, totalUnclaimedTokens } from "./tokenIncentives";

export const POOL_ACCOUNT_NEW_SIZE = 300;

export class PumpAmmInternalSdk {
  public readonly connection: Connection;
  private readonly program: Program<PumpAmm>;
  private readonly offlineProgram: Program<PumpAmm>;
  private readonly globalConfig: PublicKey;

  constructor(connection: Connection, programId: string = PUMP_AMM_PROGRAM_ID) {
    this.connection = connection;

    this.program = getPumpAmmProgram(connection, programId);
    this.offlineProgram = getPumpAmmProgram(
      null as any as Connection,
      programId,
    );

    this.globalConfig = globalConfigPda(this.offlineProgram.programId)[0];
  }

  programId(): PublicKey {
    return this.offlineProgram.programId;
  }

  globalConfigKey(): PublicKey {
    return this.globalConfig;
  }

  poolKey(
    index: number,
    creator: PublicKey,
    baseMint: PublicKey,
    quoteMint: PublicKey,
  ): [PublicKey, number] {
    return poolPda(
      index,
      creator,
      baseMint,
      quoteMint,
      this.offlineProgram.programId,
    );
  }

  lpMintKey(pool: PublicKey): [PublicKey, number] {
    return lpMintPda(pool, this.offlineProgram.programId);
  }

  fetchGlobalConfigAccount(): Promise<GlobalConfig> {
    return this.program.account.globalConfig.fetch(this.globalConfig);
  }

  fetchPool(pool: PublicKey): Promise<Pool> {
    return this.program.account.pool.fetch(pool);
  }

  decodeGlobalConfig(
    globalConfigAccountInfo: AccountInfo<Buffer>,
  ): GlobalConfig {
    return this.offlineProgram.coder.accounts.decode<GlobalConfig>(
      "globalConfig",
      globalConfigAccountInfo.data,
    );
  }

  decodePool(poolAccountInfo: AccountInfo<Buffer>) {
    return this.offlineProgram.coder.accounts.decode<Pool>(
      "pool",
      poolAccountInfo.data,
    );
  }

  fetchGlobalVolumeAccumulator(): Promise<GlobalVolumeAccumulator> {
    return this.program.account.globalVolumeAccumulator.fetch(
      globalVolumeAccumulatorPda()[0],
    );
  }

  decodeGlobalVolumeAccumulator(
    globalVolumeAccumulatorAccountInfo: AccountInfo<Buffer>,
  ): GlobalVolumeAccumulator {
    return this.offlineProgram.coder.accounts.decode<GlobalVolumeAccumulator>(
      "globalVolumeAccumulator",
      globalVolumeAccumulatorAccountInfo.data,
    );
  }

  fetchUserVolumeAccumulator(
    user: PublicKey,
  ): Promise<UserVolumeAccumulator | null> {
    return this.program.account.userVolumeAccumulator.fetchNullable(
      userVolumeAccumulatorPda(user)[0],
    );
  }

  decodeUserVolumeAccumulator(
    userVolumeAccumulatorAccountInfo: AccountInfo<Buffer>,
  ): UserVolumeAccumulator {
    return this.offlineProgram.coder.accounts.decode<UserVolumeAccumulator>(
      "userVolumeAccumulator",
      userVolumeAccumulatorAccountInfo.data,
    );
  }

  async createPoolInstructionsInternal(
    createPoolSolanaState: CreatePoolSolanaState,
    baseIn: BN,
    quoteIn: BN,
  ): Promise<TransactionInstruction[]> {
    const {
      index,
      creator,
      baseMint,
      quoteMint,
      poolKey,
      baseTokenProgram,
      quoteTokenProgram,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
      userBaseAccountInfo,
      userQuoteAccountInfo,
      poolBaseAccountInfo,
      poolQuoteAccountInfo,
    } = createPoolSolanaState;

    return await this.withWsolAccounts(
      creator,
      baseMint,
      userBaseTokenAccount,
      this.accountExists(userBaseAccountInfo, baseTokenProgram),
      baseIn,
      quoteMint,
      userQuoteTokenAccount,
      this.accountExists(userQuoteAccountInfo, quoteTokenProgram),
      quoteIn,
      async () => {
        const instructions: TransactionInstruction[] = [];

        if (!this.accountExists(poolBaseAccountInfo, baseTokenProgram)) {
          instructions.push(
            createAssociatedTokenAccountIdempotentInstruction(
              creator,
              poolBaseTokenAccount,
              poolKey,
              baseMint,
              baseTokenProgram,
            ),
          );
        }

        if (!this.accountExists(poolQuoteAccountInfo, quoteTokenProgram)) {
          instructions.push(
            createAssociatedTokenAccountIdempotentInstruction(
              creator,
              poolQuoteTokenAccount,
              poolKey,
              quoteMint,
              quoteTokenProgram,
            ),
          );
        }

        instructions.push(
          await this.offlineProgram.methods
            .createPool(index, baseIn, quoteIn, SystemProgram.programId)
            .accountsPartial({
              globalConfig: this.globalConfig,
              baseMint,
              quoteMint,
              creator,
              userBaseTokenAccount,
              userQuoteTokenAccount,
              baseTokenProgram,
              quoteTokenProgram,
            })
            .instruction(),
        );

        return instructions;
      },
    );
  }

  async depositInstructionsInternal(
    liquiditySolanaState: LiquiditySolanaState,
    lpToken: BN,
    maxBase: BN,
    maxQuote: BN,
  ): Promise<TransactionInstruction[]> {
    const {
      pool,
      user,
      userPoolAccountInfo,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      userPoolTokenAccount,
      userBaseAccountInfo,
      userQuoteAccountInfo,
      baseTokenProgram,
      quoteTokenProgram,
    } = liquiditySolanaState;

    const { baseMint, quoteMint, lpMint } = pool;

    const liquidityAccounts = this.liquidityAccounts(liquiditySolanaState);

    return await this.withFixPoolInstructions(
      liquiditySolanaState,
      async () => {
        return await this.withWsolAccounts(
          user,
          baseMint,
          userBaseTokenAccount,
          this.accountExists(userBaseAccountInfo, baseTokenProgram),
          maxBase,
          quoteMint,
          userQuoteTokenAccount,
          this.accountExists(userQuoteAccountInfo, quoteTokenProgram),
          maxQuote,
          async () => {
            const instructions: TransactionInstruction[] = [];

            if (
              !this.accountExists(userPoolAccountInfo, TOKEN_2022_PROGRAM_ID)
            ) {
              instructions.push(
                createAssociatedTokenAccountIdempotentInstruction(
                  user,
                  userPoolTokenAccount,
                  user,
                  lpMint,
                  TOKEN_2022_PROGRAM_ID,
                ),
              );
            }

            instructions.push(
              await this.offlineProgram.methods
                .deposit(lpToken, maxBase, maxQuote)
                .accounts(liquidityAccounts)
                .instruction(),
            );

            return instructions;
          },
        );
      },
    );
  }

  private async withWsolAccounts(
    user: PublicKey,
    baseMint: PublicKey,
    userBaseAta: PublicKey,
    userBaseAtaExists: boolean,
    baseAmount: BN,
    quoteMint: PublicKey,
    userQuoteAta: PublicKey,
    userQuoteAtaExists: boolean,
    quoteAmount: BN,
    block: () => Promise<TransactionInstruction[]>,
  ) {
    return await this.withWsolAccount(
      user,
      user,
      baseMint,
      userBaseAta,
      userBaseAtaExists,
      baseAmount,
      async () =>
        this.withWsolAccount(
          user,
          user,
          quoteMint,
          userQuoteAta,
          userQuoteAtaExists,
          quoteAmount,
          block,
        ),
    );
  }

  private async withWsolAccount(
    payer: PublicKey,
    user: PublicKey,
    mint: PublicKey,
    ata: PublicKey,
    ataExists: boolean,
    amount: BN,
    block: () => Promise<TransactionInstruction[]>,
    closeWsolAccount: boolean = true,
  ): Promise<TransactionInstruction[]> {
    const instructions: TransactionInstruction[] = [];

    if (mint.equals(NATIVE_MINT)) {
      if (!ataExists) {
        instructions.push(
          createAssociatedTokenAccountIdempotentInstruction(
            payer,
            ata,
            user,
            NATIVE_MINT,
          ),
        );
      }

      if (amount.gtn(0)) {
        instructions.push(
          SystemProgram.transfer({
            fromPubkey: user,
            toPubkey: ata,
            lamports: BigInt(amount.toString()),
          }),
          createSyncNativeInstruction(ata),
        );
      }
    }

    const blockInstructions = await block();
    instructions.push(...blockInstructions);

    if (mint.equals(NATIVE_MINT) && closeWsolAccount) {
      instructions.push(
        createCloseAccountInstruction(
          ata,
          user,
          user,
          undefined,
          TOKEN_PROGRAM_ID,
        ),
      );
    }

    return instructions;
  }

  private accountExists(
    accountInfo: AccountInfo<Buffer> | null,
    owner: PublicKey,
  ): boolean {
    return accountInfo !== null && accountInfo.owner.equals(owner);
  }

  depositBaseInputInternal(
    liquiditySolanaState: LiquiditySolanaState,
    base: BN,
    slippage: number,
  ): DepositBaseResult {
    const { pool, poolBaseTokenAccount, poolQuoteTokenAccount } =
      liquiditySolanaState;

    const { token1, lpToken, maxToken0, maxToken1 } = depositToken0Internal(
      base,
      slippage,
      new BN(poolBaseTokenAccount.amount.toString()),
      new BN(poolQuoteTokenAccount.amount.toString()),
      pool.lpSupply,
    );

    return {
      quote: token1,
      lpToken,
      maxBase: maxToken0,
      maxQuote: maxToken1,
    };
  }

  depositQuoteInputInternal(
    liquiditySolanaState: LiquiditySolanaState,
    quote: BN,
    slippage: number,
  ): DepositQuoteResult {
    const { pool, poolBaseTokenAccount, poolQuoteTokenAccount } =
      liquiditySolanaState;

    const { token1, lpToken, maxToken0, maxToken1 } = depositToken0Internal(
      quote,
      slippage,
      new BN(poolQuoteTokenAccount.amount.toString()),
      new BN(poolBaseTokenAccount.amount.toString()),
      pool.lpSupply,
    );

    return {
      base: token1,
      lpToken,
      maxBase: maxToken1,
      maxQuote: maxToken0,
    };
  }

  async withdrawInstructionsInternal(
    liquiditySolanaState: LiquiditySolanaState,
    lpTokenAmountIn: BN,
    minBaseAmountOut: BN,
    minQuoteAmountOut: BN,
  ): Promise<TransactionInstruction[]> {
    const {
      pool,
      baseTokenProgram,
      quoteTokenProgram,
      user,
      userBaseAccountInfo,
      userQuoteAccountInfo,
      userBaseTokenAccount,
      userQuoteTokenAccount,
    } = liquiditySolanaState;

    const { baseMint, quoteMint } = pool;

    const liquidityAccounts = this.liquidityAccounts(liquiditySolanaState);

    return await this.withFixPoolInstructions(
      liquiditySolanaState,
      async () => {
        const instructions: TransactionInstruction[] = [];

        let baseWsolAtaCreated = false;

        if (!this.accountExists(userBaseAccountInfo, baseTokenProgram)) {
          instructions.push(
            createAssociatedTokenAccountIdempotentInstruction(
              user,
              userBaseTokenAccount,
              user,
              baseMint,
              baseTokenProgram,
            ),
          );

          if (baseMint.equals(NATIVE_MINT)) {
            baseWsolAtaCreated = true;
          }
        }

        let quoteWsolAtaCreated = false;

        if (!this.accountExists(userQuoteAccountInfo, quoteTokenProgram)) {
          instructions.push(
            createAssociatedTokenAccountIdempotentInstruction(
              user,
              userQuoteTokenAccount,
              user,
              quoteMint,
              quoteTokenProgram,
            ),
          );

          if (quoteMint.equals(NATIVE_MINT)) {
            quoteWsolAtaCreated = true;
          }
        }

        instructions.push(
          await this.offlineProgram.methods
            .withdraw(lpTokenAmountIn, minBaseAmountOut, minQuoteAmountOut)
            .accounts(liquidityAccounts)
            .instruction(),
        );

        if (baseWsolAtaCreated) {
          instructions.push(
            createCloseAccountInstruction(
              userBaseTokenAccount,
              user,
              user,
              undefined,
              TOKEN_PROGRAM_ID,
            ),
          );
        }

        if (quoteWsolAtaCreated) {
          instructions.push(
            createCloseAccountInstruction(
              userQuoteTokenAccount,
              user,
              user,
              undefined,
              TOKEN_PROGRAM_ID,
            ),
          );
        }

        return instructions;
      },
    );
  }

  withdrawInputsInternal(
    liquiditySolanaState: LiquiditySolanaState,
    lpAmount: BN,
    slippage: number,
  ): WithdrawResult {
    const { pool, poolBaseTokenAccount, poolQuoteTokenAccount } =
      liquiditySolanaState;

    return withdrawInternal(
      lpAmount,
      slippage,
      new BN(poolBaseTokenAccount.amount.toString()),
      new BN(poolQuoteTokenAccount.amount.toString()),
      pool.lpSupply,
    );
  }

  private liquidityAccounts(
    liquiditySolanaState: LiquiditySolanaState,
  ): LiquidityAccounts {
    const {
      poolKey,
      pool,
      user,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      userPoolTokenAccount,
    } = liquiditySolanaState;

    const {
      baseMint,
      quoteMint,
      lpMint,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
    } = pool;

    let program = this.programId();
    let [eventAuthority] = pumpAmmEventAuthorityPda(program);

    return {
      pool: poolKey,
      globalConfig: this.globalConfig,
      user,
      baseMint,
      quoteMint,
      lpMint,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      userPoolTokenAccount,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
      tokenProgram: TOKEN_PROGRAM_ID,
      token2022Program: TOKEN_2022_PROGRAM_ID,
      eventAuthority,
      program,
    };
  }

  async buyInstructionsInternal(
    swapSolanaState: SwapSolanaState,
    baseOut: BN,
    maxQuoteIn: BN,
  ): Promise<TransactionInstruction[]> {
    return await this.withFixPoolInstructions(swapSolanaState, async () => {
      return await this.buyInstructionsInternalNoPool(
        swapSolanaState,
        baseOut,
        maxQuoteIn,
      );
    });
  }

  async createPoolSolanaState(
    index: number,
    creator: PublicKey,
    baseMint: PublicKey,
    quoteMint: PublicKey,
    userBaseTokenAccount: PublicKey | undefined = undefined,
    userQuoteTokenAccount: PublicKey | undefined = undefined,
  ): Promise<CreatePoolSolanaState> {
    const [globalConfigAccountInfo, baseMintAccountInfo, quoteMintAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        this.globalConfig,
        baseMint,
        quoteMint,
      ]);

    if (globalConfigAccountInfo === null) {
      throw new Error("Global config account not found");
    }

    if (baseMintAccountInfo === null) {
      throw new Error(`baseMint=${baseMint.toString()} not found`);
    }

    if (quoteMintAccountInfo === null) {
      throw new Error(`quoteMint=${quoteMint.toString()} not found`);
    }

    const globalConfig = this.decodeGlobalConfig(globalConfigAccountInfo);

    const [baseTokenProgram, quoteTokenProgram] = [
      baseMintAccountInfo.owner,
      quoteMintAccountInfo.owner,
    ];

    const [poolKey] = poolPda(
      index,
      creator,
      baseMint,
      quoteMint,
      this.offlineProgram.programId,
    );

    const poolBaseTokenAccount = getAssociatedTokenAddressSync(
      baseMint,
      poolKey,
      true,
      baseTokenProgram,
    );

    const poolQuoteTokenAccount = getAssociatedTokenAddressSync(
      quoteMint,
      poolKey,
      true,
      quoteTokenProgram,
    );

    const [poolBaseAccountInfo, poolQuoteAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        poolBaseTokenAccount,
        poolQuoteTokenAccount,
      ]);

    if (userBaseTokenAccount === undefined) {
      userBaseTokenAccount = getAssociatedTokenAddressSync(
        baseMint,
        creator,
        true,
        baseTokenProgram,
      );
    }

    if (userQuoteTokenAccount === undefined) {
      userQuoteTokenAccount = getAssociatedTokenAddressSync(
        quoteMint,
        creator,
        true,
        quoteTokenProgram,
      );
    }

    const [userBaseAccountInfo, userQuoteAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        userBaseTokenAccount,
        userQuoteTokenAccount,
      ]);

    return {
      index,
      creator,
      baseMint,
      quoteMint,
      globalConfig,
      poolKey,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
      baseTokenProgram,
      quoteTokenProgram,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      userBaseAccountInfo,
      userQuoteAccountInfo,
      poolBaseAccountInfo,
      poolQuoteAccountInfo,
    };
  }

  async swapSolanaState(
    poolKey: PublicKey,
    user: PublicKey,
    userBaseTokenAccount: PublicKey | undefined = undefined,
    userQuoteTokenAccount: PublicKey | undefined = undefined,
  ): Promise<SwapSolanaState> {
    const [globalConfigAccountInfo, poolAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        this.globalConfig,
        poolKey,
      ]);

    if (globalConfigAccountInfo === null) {
      throw new Error("Global config account not found");
    }

    if (poolAccountInfo === null) {
      throw new Error("Pool account not found");
    }

    const globalConfig = this.decodeGlobalConfig(globalConfigAccountInfo);
    const pool = this.decodePool(poolAccountInfo);

    const { baseMint, quoteMint, poolBaseTokenAccount, poolQuoteTokenAccount } =
      pool;

    const [
      baseMintAccountInfo,
      quoteMintAccountInfo,
      poolBaseAccountInfo,
      poolQuoteAccountInfo,
    ] = await this.connection.getMultipleAccountsInfo([
      baseMint,
      quoteMint,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
    ]);

    if (baseMintAccountInfo === null) {
      throw new Error(`baseMint=${baseMint.toString()} not found`);
    }

    if (quoteMintAccountInfo === null) {
      throw new Error(`quoteMint=${quoteMint.toString()} not found`);
    }

    if (poolBaseAccountInfo === null) {
      throw new Error(
        `Pool base token account ${poolBaseTokenAccount.toString()} not found`,
      );
    }

    if (poolQuoteAccountInfo === null) {
      throw new Error(
        `Pool quote token account ${poolQuoteTokenAccount.toString()} not found`,
      );
    }

    const [baseTokenProgram, quoteTokenProgram] = [
      baseMintAccountInfo.owner,
      quoteMintAccountInfo.owner,
    ];

    const decodedPoolBaseTokenAccount = AccountLayout.decode(
      poolBaseAccountInfo.data,
    );
    const decodedPoolQuoteTokenAccount = AccountLayout.decode(
      poolQuoteAccountInfo.data,
    );

    if (userBaseTokenAccount === undefined) {
      userBaseTokenAccount = getAssociatedTokenAddressSync(
        baseMint,
        user,
        true,
        baseTokenProgram,
      );
    }

    if (userQuoteTokenAccount === undefined) {
      userQuoteTokenAccount = getAssociatedTokenAddressSync(
        quoteMint,
        user,
        true,
        quoteTokenProgram,
      );
    }

    const [userBaseAccountInfo, userQuoteAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        userBaseTokenAccount,
        userQuoteTokenAccount,
      ]);

    return {
      globalConfig,
      poolKey,
      poolAccountInfo,
      pool,
      poolBaseAmount: new BN(decodedPoolBaseTokenAccount.amount.toString()),
      poolQuoteAmount: new BN(decodedPoolQuoteTokenAccount.amount.toString()),
      baseTokenProgram,
      quoteTokenProgram,
      user,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      userBaseAccountInfo,
      userQuoteAccountInfo,
    };
  }

  async swapSolanaStateNoPool(
    poolKey: PublicKey,
    user: PublicKey,
    userBaseTokenAccount: PublicKey | undefined = undefined,
    userQuoteTokenAccount: PublicKey | undefined = undefined,
  ): Promise<SwapSolanaState> {
    const [globalConfigAccountInfo, poolAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        this.globalConfig,
        poolKey,
      ]);

    if (globalConfigAccountInfo === null) {
      throw new Error("Global config account not found");
    }

    if (poolAccountInfo === null) {
      throw new Error("Pool account not found");
    }

    const globalConfig = this.decodeGlobalConfig(globalConfigAccountInfo);
    const pool = this.decodePool(poolAccountInfo);

    const { baseMint, quoteMint, poolBaseTokenAccount, poolQuoteTokenAccount } =
      pool;

    const [
      baseMintAccountInfo,
      quoteMintAccountInfo,
      poolBaseAccountInfo,
      poolQuoteAccountInfo,
    ] = await this.connection.getMultipleAccountsInfo([
      baseMint,
      quoteMint,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
    ]);

    if (baseMintAccountInfo === null) {
      throw new Error(`baseMint=${baseMint.toString()} not found`);
    }

    if (quoteMintAccountInfo === null) {
      throw new Error(`quoteMint=${quoteMint.toString()} not found`);
    }

    if (poolBaseAccountInfo === null) {
      throw new Error(
        `Pool base token account ${poolBaseTokenAccount.toString()} not found`,
      );
    }

    if (poolQuoteAccountInfo === null) {
      throw new Error(
        `Pool quote token account ${poolQuoteTokenAccount.toString()} not found`,
      );
    }

    const [baseTokenProgram, quoteTokenProgram] = [
      baseMintAccountInfo.owner,
      quoteMintAccountInfo.owner,
    ];

    const decodedPoolBaseTokenAccount = AccountLayout.decode(
      poolBaseAccountInfo.data,
    );
    const decodedPoolQuoteTokenAccount = AccountLayout.decode(
      poolQuoteAccountInfo.data,
    );

    if (userBaseTokenAccount === undefined) {
      userBaseTokenAccount = getAssociatedTokenAddressSync(
        baseMint,
        user,
        true,
        baseTokenProgram,
      );
    }

    if (userQuoteTokenAccount === undefined) {
      userQuoteTokenAccount = getAssociatedTokenAddressSync(
        quoteMint,
        user,
        true,
        quoteTokenProgram,
      );
    }

    const [userBaseAccountInfo, userQuoteAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        userBaseTokenAccount,
        userQuoteTokenAccount,
      ]);

    return {
      globalConfig,
      poolKey,
      poolAccountInfo,
      pool,
      poolBaseAmount: new BN(decodedPoolBaseTokenAccount.amount.toString()),
      poolQuoteAmount: new BN(decodedPoolQuoteTokenAccount.amount.toString()),
      baseTokenProgram,
      quoteTokenProgram,
      user,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      userBaseAccountInfo,
      userQuoteAccountInfo,
    };
  }

  async liquiditySolanaState(
    poolKey: PublicKey,
    user: PublicKey,
    userBaseTokenAccount: PublicKey | undefined = undefined,
    userQuoteTokenAccount: PublicKey | undefined = undefined,
    userPoolTokenAccount: PublicKey | undefined = undefined,
  ): Promise<LiquiditySolanaState> {
    const [globalConfigAccountInfo, poolAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        this.globalConfig,
        poolKey,
      ]);

    if (globalConfigAccountInfo === null) {
      throw new Error("Global config account not found");
    }

    if (poolAccountInfo === null) {
      throw new Error("Pool account not found");
    }

    const globalConfig = this.decodeGlobalConfig(globalConfigAccountInfo);
    const pool = this.decodePool(poolAccountInfo);

    const {
      baseMint,
      quoteMint,
      lpMint,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
    } = pool;

    const [
      baseMintAccountInfo,
      quoteMintAccountInfo,
      poolBaseAccountInfo,
      poolQuoteAccountInfo,
    ] = await this.connection.getMultipleAccountsInfo([
      baseMint,
      quoteMint,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
    ]);

    if (baseMintAccountInfo === null) {
      throw new Error(`baseMint=${baseMint.toString()} not found`);
    }

    if (quoteMintAccountInfo === null) {
      throw new Error(`quoteMint=${quoteMint.toString()} not found`);
    }

    if (poolBaseAccountInfo === null) {
      throw new Error(
        `Pool base token account ${poolBaseTokenAccount.toString()} not found`,
      );
    }

    if (poolQuoteAccountInfo === null) {
      throw new Error(
        `Pool quote token account ${poolQuoteTokenAccount.toString()} not found`,
      );
    }

    const [baseTokenProgram, quoteTokenProgram] = [
      baseMintAccountInfo.owner,
      quoteMintAccountInfo.owner,
    ];

    const decodedPoolBaseTokenAccount = AccountLayout.decode(
      poolBaseAccountInfo.data,
    );
    const decodedPoolQuoteTokenAccount = AccountLayout.decode(
      poolQuoteAccountInfo.data,
    );

    if (userBaseTokenAccount === undefined) {
      userBaseTokenAccount = getAssociatedTokenAddressSync(
        baseMint,
        user,
        true,
        baseTokenProgram,
      );
    }

    if (userQuoteTokenAccount === undefined) {
      userQuoteTokenAccount = getAssociatedTokenAddressSync(
        quoteMint,
        user,
        true,
        quoteTokenProgram,
      );
    }

    if (userPoolTokenAccount === undefined) {
      userPoolTokenAccount = getAssociatedTokenAddressSync(
        lpMint,
        user,
        true,
        TOKEN_2022_PROGRAM_ID,
      );
    }

    const [userBaseAccountInfo, userQuoteAccountInfo, userPoolAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        userBaseTokenAccount,
        userQuoteTokenAccount,
        userPoolTokenAccount,
      ]);

    return {
      globalConfig,
      poolKey,
      poolAccountInfo,
      pool,
      poolBaseTokenAccount: decodedPoolBaseTokenAccount,
      poolQuoteTokenAccount: decodedPoolQuoteTokenAccount,
      baseTokenProgram,
      quoteTokenProgram,
      user,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      userPoolTokenAccount,
      userBaseAccountInfo,
      userQuoteAccountInfo,
      userPoolAccountInfo,
    };
  }

  async buyInstructionsInternalNoPool(
    swapSolanaState: SwapSolanaState,
    baseOut: BN,
    maxQuoteIn: BN,
  ): Promise<TransactionInstruction[]> {
    const { userBaseAccountInfo, userQuoteAccountInfo } = swapSolanaState;

    const swapAccounts = this.swapAccounts(swapSolanaState);

    const {
      user,
      baseMint,
      quoteMint,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      baseTokenProgram,
      quoteTokenProgram,
    } = swapAccounts;

    return this.withWsolAccount(
      user,
      user,
      quoteMint,
      userQuoteTokenAccount,
      this.accountExists(userQuoteAccountInfo, quoteTokenProgram),
      maxQuoteIn,
      async () => {
        const instructions = [];

        if (!this.accountExists(userBaseAccountInfo, baseTokenProgram)) {
          instructions.push(
            createAssociatedTokenAccountIdempotentInstruction(
              user,
              userBaseTokenAccount,
              user,
              baseMint,
              baseTokenProgram,
            ),
          );
        }

        instructions.push(
          await this.offlineProgram.methods
            .buy(baseOut, maxQuoteIn, { 0: true })
            .accounts(swapAccounts)
            .instruction(),
        );

        if (baseMint.equals(NATIVE_MINT)) {
          instructions.push(
            createCloseAccountInstruction(
              userBaseTokenAccount,
              user,
              user,
              undefined,
              TOKEN_PROGRAM_ID,
            ),
          );
        }

        return instructions;
      },
    );
  }

  async buyBaseInput(
    swapSolanaState: SwapSolanaState,
    base: BN,
    slippage: number,
  ): Promise<TransactionInstruction[]> {
    const { pool, globalConfig, poolBaseAmount, poolQuoteAmount } =
      swapSolanaState;

    const { maxQuote } = buyBaseInputInternal(
      base,
      slippage,
      poolBaseAmount,
      poolQuoteAmount,
      globalConfig,
      pool.coinCreator,
    );

    return this.buyInstructionsInternal(swapSolanaState, base, maxQuote);
  }

  async buyQuoteInput(
    swapSolanaState: SwapSolanaState,
    quote: BN,
    slippage: number,
  ): Promise<TransactionInstruction[]> {
    const { globalConfig, pool, poolBaseAmount, poolQuoteAmount } =
      swapSolanaState;

    const { base, maxQuote } = buyQuoteInputInternal(
      quote,
      slippage,
      poolBaseAmount,
      poolQuoteAmount,
      globalConfig,
      pool.coinCreator,
    );

    return this.buyInstructionsInternal(swapSolanaState, base, maxQuote);
  }

  async sellInstructionsInternal(
    swapSolanaState: SwapSolanaState,
    baseAmountIn: BN,
    minQuoteAmountOut: BN,
  ): Promise<TransactionInstruction[]> {
    return await this.withFixPoolInstructions(swapSolanaState, async () => {
      return await this.sellInstructionsInternalNoPool(
        swapSolanaState,
        baseAmountIn,
        minQuoteAmountOut,
      );
    });
  }

  private async withFixPoolInstructions(
    commonSolanaState: CommonSolanaState,
    block: () => Promise<TransactionInstruction[]>,
  ): Promise<TransactionInstruction[]> {
    const { poolAccountInfo, poolKey, user } = commonSolanaState;

    const instructions: TransactionInstruction[] = [];

    if (
      poolAccountInfo === null ||
      poolAccountInfo.data.length < POOL_ACCOUNT_NEW_SIZE
    ) {
      instructions.push(
        await this.offlineProgram.methods
          .extendAccount()
          .accountsPartial({
            account: poolKey,
            user,
          })
          .instruction(),
      );
    }

    return [...instructions, ...(await block())];
  }

  async sellInstructionsInternalNoPool(
    swapSolanaState: SwapSolanaState,
    baseAmountIn: BN,
    minQuoteAmountOut: BN,
  ): Promise<TransactionInstruction[]> {
    const { userBaseAccountInfo, userQuoteAccountInfo } = swapSolanaState;

    const swapAccounts = this.swapAccounts(swapSolanaState);

    const {
      user,
      baseMint,
      quoteMint,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      baseTokenProgram,
      quoteTokenProgram,
    } = swapAccounts;

    return this.withWsolAccount(
      user,
      user,
      baseMint,
      userBaseTokenAccount,
      this.accountExists(userBaseAccountInfo, baseTokenProgram),
      baseAmountIn,
      async () => {
        const instructions = [];

        if (!this.accountExists(userQuoteAccountInfo, quoteTokenProgram)) {
          instructions.push(
            createAssociatedTokenAccountIdempotentInstruction(
              user,
              userQuoteTokenAccount,
              user,
              quoteMint,
              quoteTokenProgram,
            ),
          );
        }

        instructions.push(
          await this.offlineProgram.methods
            .sell(baseAmountIn, minQuoteAmountOut)
            .accounts(swapAccounts)
            .instruction(),
        );

        if (quoteMint.equals(NATIVE_MINT)) {
          instructions.push(
            createCloseAccountInstruction(
              userQuoteTokenAccount,
              user,
              user,
              undefined,
              TOKEN_PROGRAM_ID,
            ),
          );
        }

        return instructions;
      },
    );
  }

  async sellBaseInput(
    swapSolanaState: SwapSolanaState,
    base: BN,
    slippage: number,
  ): Promise<TransactionInstruction[]> {
    const { globalConfig, pool, poolBaseAmount, poolQuoteAmount } =
      swapSolanaState;

    const { minQuote } = sellBaseInputInternal(
      base,
      slippage,
      poolBaseAmount,
      poolQuoteAmount,
      globalConfig,
      pool.coinCreator,
    );

    return this.sellInstructionsInternal(swapSolanaState, base, minQuote);
  }

  async sellQuoteInput(
    swapSolanaState: SwapSolanaState,
    quote: BN,
    slippage: number,
  ): Promise<TransactionInstruction[]> {
    const { globalConfig, pool, poolBaseAmount, poolQuoteAmount } =
      swapSolanaState;

    const { base, minQuote } = sellQuoteInputInternal(
      quote,
      slippage,
      poolBaseAmount,
      poolQuoteAmount,
      globalConfig,
      pool.coinCreator,
    );

    return this.sellInstructionsInternal(swapSolanaState, base, minQuote);
  }

  async extendAccount(
    account: PublicKey,
    user: PublicKey,
  ): Promise<TransactionInstruction> {
    return this.offlineProgram.methods
      .extendAccount()
      .accountsPartial({
        account,
        user,
      })
      .instruction();
  }

  async collectCoinCreatorFeeSolanaState(
    coinCreator: PublicKey,
    coinCreatorTokenAccount: PublicKey | undefined = undefined,
  ): Promise<CollectCoinCreatorFeeSolanaState> {
    const quoteMint = NATIVE_MINT;
    const quoteTokenProgram = TOKEN_PROGRAM_ID;

    let coinCreatorVaultAuthority =
      this.coinCreatorVaultAuthorityPda(coinCreator);

    let coinCreatorVaultAta = this.coinCreatorVaultAta(
      coinCreatorVaultAuthority,
      quoteMint,
      quoteTokenProgram,
    );

    if (coinCreatorTokenAccount === undefined) {
      coinCreatorTokenAccount = getAssociatedTokenAddressSync(
        quoteMint,
        coinCreator,
        true,
        quoteTokenProgram,
      );
    }

    const [coinCreatorVaultAtaAccountInfo, coinCreatorTokenAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        coinCreatorVaultAta,
        coinCreatorTokenAccount,
      ]);

    return {
      coinCreator,
      quoteMint,
      quoteTokenProgram,
      coinCreatorVaultAuthority,
      coinCreatorVaultAta,
      coinCreatorTokenAccount,
      coinCreatorVaultAtaAccountInfo,
      coinCreatorTokenAccountInfo,
    };
  }

  async collectCoinCreatorFee(
    collectCoinCreatorFeeSolanaState: CollectCoinCreatorFeeSolanaState,
  ): Promise<TransactionInstruction[]> {
    const {
      coinCreator,
      quoteMint,
      quoteTokenProgram,
      coinCreatorVaultAuthority,
      coinCreatorVaultAta,
      coinCreatorTokenAccount,
      coinCreatorVaultAtaAccountInfo,
      coinCreatorTokenAccountInfo,
    } = collectCoinCreatorFeeSolanaState;

    return await this.withWsolAccount(
      coinCreator,
      coinCreatorVaultAuthority,
      quoteMint,
      coinCreatorVaultAta,
      this.accountExists(coinCreatorVaultAtaAccountInfo, quoteTokenProgram),
      new BN(0),
      async () => {
        return await this.withWsolAccount(
          coinCreator,
          coinCreator,
          quoteMint,
          coinCreatorTokenAccount,
          this.accountExists(coinCreatorTokenAccountInfo, quoteTokenProgram),
          new BN(0),
          async () => {
            return [
              await this.offlineProgram.methods
                .collectCoinCreatorFee()
                .accountsPartial({
                  coinCreator,
                  coinCreatorTokenAccount,
                  quoteMint,
                  quoteTokenProgram,
                })
                .instruction(),
            ];
          },
        );
      },
      false,
    );
  }

  async getCoinCreatorVaultBalance(coinCreator: PublicKey): Promise<BN> {
    const quoteMint = NATIVE_MINT;
    const quoteTokenProgram = TOKEN_PROGRAM_ID;

    const coinCreatorVaultAuthority =
      this.coinCreatorVaultAuthorityPda(coinCreator);

    const coinCreatorVaultAta = this.coinCreatorVaultAta(
      coinCreatorVaultAuthority,
      quoteMint,
      quoteTokenProgram,
    );

    try {
      const tokenAccount = await getAccount(
        this.connection,
        coinCreatorVaultAta,
        undefined,
        quoteTokenProgram,
      );
      return new BN(tokenAccount.amount.toString());
    } catch (e) {
      console.error(`Error fetching token account ${coinCreatorVaultAta}:`, e);
      return new BN(0);
    }
  }

  async setCoinCreator(pool: PublicKey): Promise<TransactionInstruction> {
    return this.offlineProgram.methods
      .setCoinCreator()
      .accountsPartial({
        pool,
      })
      .instruction();
  }

  private swapAccounts(swapSolanaState: SwapSolanaState): SwapAccounts {
    const {
      globalConfig,
      poolKey,
      pool,
      baseTokenProgram,
      quoteTokenProgram,
      user,
      userBaseTokenAccount,
      userQuoteTokenAccount,
    } = swapSolanaState;

    const { protocolFeeRecipients } = globalConfig;

    const protocolFeeRecipient =
      protocolFeeRecipients[
        Math.floor(Math.random() * protocolFeeRecipients.length)
      ];

    const {
      baseMint,
      quoteMint,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
      coinCreator,
    } = pool;

    const coinCreatorVaultAuthority =
      this.coinCreatorVaultAuthorityPda(coinCreator);

    let program = this.programId();
    let [eventAuthority] = pumpAmmEventAuthorityPda(program);

    return {
      pool: poolKey,
      globalConfig: this.globalConfig,
      user,
      baseMint,
      quoteMint,
      userBaseTokenAccount,
      userQuoteTokenAccount,
      poolBaseTokenAccount,
      poolQuoteTokenAccount,
      protocolFeeRecipient,
      protocolFeeRecipientTokenAccount: getAssociatedTokenAddressSync(
        quoteMint,
        protocolFeeRecipient,
        true,
        quoteTokenProgram,
      ),
      baseTokenProgram,
      quoteTokenProgram,
      systemProgram: SystemProgram.programId,
      associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
      eventAuthority,
      program,
      coinCreatorVaultAta: this.coinCreatorVaultAta(
        coinCreatorVaultAuthority,
        quoteMint,
        quoteTokenProgram,
      ),
      coinCreatorVaultAuthority,
    };
  }

  coinCreatorVaultAuthorityPda(coinCreator: PublicKey) {
    const [coinCreatorVaultAuthority] = PublicKey.findProgramAddressSync(
      [Buffer.from("creator_vault"), coinCreator.toBuffer()],
      this.programId(),
    );
    return coinCreatorVaultAuthority;
  }

  coinCreatorVaultAta(
    coinCreatorVaultAuthority: PublicKey,
    quoteMint: PublicKey,
    quoteTokenProgram: PublicKey,
  ) {
    return getAssociatedTokenAddressSync(
      quoteMint,
      coinCreatorVaultAuthority,
      true,
      quoteTokenProgram,
    );
  }

  async claimTokenIncentivesInternal(
    user: PublicKey,
    payer: PublicKey,
  ): Promise<TransactionInstruction[]> {
    const { mint } = await this.fetchGlobalVolumeAccumulator();

    if (mint.equals(PublicKey.default)) {
      return [];
    }

    const [mintAccountInfo, userAccumulatorAccountInfo] =
      await this.connection.getMultipleAccountsInfo([
        mint,
        userVolumeAccumulatorPda(user)[0],
      ]);

    if (!mintAccountInfo) {
      return [];
    }

    if (!userAccumulatorAccountInfo) {
      return [];
    }

    return [
      await this.offlineProgram.methods
        .claimTokenIncentives()
        .accountsPartial({
          user,
          payer,
          mint,
          tokenProgram: mintAccountInfo.owner,
        })
        .instruction(),
    ];
  }

  async getTotalUnclaimedTokens(user: PublicKey): Promise<BN> {
    const [
      globalVolumeAccumulatorAccountInfo,
      userVolumeAccumulatorAccountInfo,
    ] = await this.connection.getMultipleAccountsInfo([
      globalVolumeAccumulatorPda()[0],
      userVolumeAccumulatorPda(user)[0],
    ]);

    if (
      !globalVolumeAccumulatorAccountInfo ||
      !userVolumeAccumulatorAccountInfo
    ) {
      return new BN(0);
    }

    const globalVolumeAccumulator = this.decodeGlobalVolumeAccumulator(
      globalVolumeAccumulatorAccountInfo,
    );
    const userVolumeAccumulator = this.decodeUserVolumeAccumulator(
      userVolumeAccumulatorAccountInfo,
    );

    return totalUnclaimedTokens(globalVolumeAccumulator, userVolumeAccumulator);
  }

  async getCurrentDayTokens(user: PublicKey): Promise<BN> {
    const [
      globalVolumeAccumulatorAccountInfo,
      userVolumeAccumulatorAccountInfo,
    ] = await this.connection.getMultipleAccountsInfo([
      globalVolumeAccumulatorPda()[0],
      userVolumeAccumulatorPda(user)[0],
    ]);

    if (
      !globalVolumeAccumulatorAccountInfo ||
      !userVolumeAccumulatorAccountInfo
    ) {
      return new BN(0);
    }

    const globalVolumeAccumulator = this.decodeGlobalVolumeAccumulator(
      globalVolumeAccumulatorAccountInfo,
    );
    const userVolumeAccumulator = this.decodeUserVolumeAccumulator(
      userVolumeAccumulatorAccountInfo,
    );

    return currentDayTokens(globalVolumeAccumulator, userVolumeAccumulator);
  }

  async syncUserVolumeAccumulator(
    user: PublicKey,
  ): Promise<TransactionInstruction> {
    return await this.offlineProgram.methods
      .syncUserVolumeAccumulator()
      .accountsPartial({ user })
      .instruction();
  }

  async initUserVolumeAccumulator({
    payer,
    user,
  }: {
    payer: PublicKey;
    user: PublicKey;
  }): Promise<TransactionInstruction> {
    return await this.offlineProgram.methods
      .initUserVolumeAccumulator()
      .accountsPartial({ payer, user })
      .instruction();
  }

  async closeUserVolumeAccumulator(
    user: PublicKey,
  ): Promise<TransactionInstruction> {
    return await this.offlineProgram.methods
      .closeUserVolumeAccumulator()
      .accountsPartial({ user })
      .instruction();
  }
}