import { Connection, PublicKey, TransactionInstruction, LAMPORTS_PER_SOL } from '@solana/web3.js';
import BN from 'bn.js';
import { Raydium, TxVersion } from '@raydium-io/raydium-sdk-v2';
import { mints } from '../../helpers/constants';
import { assertAmmPoolHasMintAndWsol, findAmmPoolInfo, findAmmPoolInfoById } from './pool-utils';

export class RaydiumAmmClient {
  private readonly connection: Connection;
  private raydiumPromise: Promise<Raydium> | null = null;

  constructor(connection: Connection) {
    this.connection = connection;
  }

  private async getRaydium(owner: PublicKey): Promise<Raydium> {
    if (!this.raydiumPromise) {
      this.raydiumPromise = Raydium.load({ connection: this.connection, owner });
    }
    return this.raydiumPromise;
  }

  async getBuyInstructions(params: { mintAddress: PublicKey; wallet: PublicKey; solAmount: number; slippage: number; poolAddress?: PublicKey; }): Promise<TransactionInstruction[]> {
    const { mintAddress, wallet, solAmount, slippage, poolAddress } = params;
    const raydium = await this.getRaydium(wallet);

    const poolInfo: any = poolAddress
      ? await findAmmPoolInfoById(raydium, poolAddress)
      : await findAmmPoolInfo(raydium, mintAddress);
    assertAmmPoolHasMintAndWsol(poolInfo, mintAddress);
    const poolKeys: any = await raydium.liquidity.getAmmPoolKeys(poolInfo.id);
    const rpcData: any = await raydium.liquidity.getRpcPoolInfo(poolInfo.id);

    const inputMint = new PublicKey(mints.WSOL).toBase58();
    const baseIn = poolInfo.mintA.address === inputMint;
    const mintIn = baseIn ? poolInfo.mintA : poolInfo.mintB;
    const mintOut = baseIn ? poolInfo.mintB : poolInfo.mintA;

    const amountIn = new BN(Math.round(solAmount * LAMPORTS_PER_SOL));

    const out = raydium.liquidity.computeAmountOut({
      poolInfo: {
        ...poolInfo,
        baseReserve: rpcData.baseReserve,
        quoteReserve: rpcData.quoteReserve,
        status: rpcData.status.toNumber(),
        version: 4,
      },
      amountIn,
      mintIn: mintIn.address,
      mintOut: mintOut.address,
      slippage,
    });

    const make = await raydium.liquidity.swap({
      poolInfo,
      poolKeys,
      amountIn,
      amountOut: out.minAmountOut,
      fixedSide: 'in',
      inputMint: mintIn.address,
      txVersion: TxVersion.LEGACY,
    });

    return make.transaction.instructions;
  }

  async getSellInstructions(params: { mintAddress: PublicKey; wallet: PublicKey; tokenAmount: number; slippage: number; poolAddress?: PublicKey; }): Promise<TransactionInstruction[]> {
    const { mintAddress, wallet, tokenAmount, slippage, poolAddress } = params;
    const raydium = await this.getRaydium(wallet);

    const poolInfo: any = poolAddress
      ? await findAmmPoolInfoById(raydium, poolAddress)
      : await findAmmPoolInfo(raydium, mintAddress);
    assertAmmPoolHasMintAndWsol(poolInfo, mintAddress);
    const poolKeys: any = await raydium.liquidity.getAmmPoolKeys(poolInfo.id);
    const rpcData: any = await raydium.liquidity.getRpcPoolInfo(poolInfo.id);

    const baseIn = poolInfo.mintA.address === mintAddress.toBase58();
    const mintIn = baseIn ? poolInfo.mintA : poolInfo.mintB;
    const mintOut = baseIn ? poolInfo.mintB : poolInfo.mintA;

    const decimals = mintIn.decimals;
    const amountIn = new BN(Math.round(tokenAmount * Math.pow(10, decimals)));

    const out = raydium.liquidity.computeAmountOut({
      poolInfo: {
        ...poolInfo,
        baseReserve: rpcData.baseReserve,
        quoteReserve: rpcData.quoteReserve,
        status: rpcData.status.toNumber(),
        version: 4,
      },
      amountIn,
      mintIn: mintIn.address,
      mintOut: mintOut.address,
      slippage,
    });

    const make = await raydium.liquidity.swap({
      poolInfo,
      poolKeys,
      amountIn,
      amountOut: out.minAmountOut,
      fixedSide: 'in',
      inputMint: mintIn.address,
      txVersion: TxVersion.LEGACY,
    });

    return make.transaction.instructions;
  }

}