import { Instruction as JupiterInstruction, QuoteResponse, SwapInstructionsResponse } from '@jup-ag/api';
import BN from 'bn.js';
import {
  compileRouteIxs,
  QuoteParams,
  RouteParams,
  RouterContext,
  RouterInstructions,
  RouterOutput,
} from '../swap_api_utils';
import { RouterType } from '../consts';
import { determineRole, getAllLookupTables } from '../utils';
import { address, Instruction, Rpc, SolanaRpcApi } from '@solana/kit';
import { extractOrAddWrapAndUnwrapSolIxs } from '../swap_api_utils/wsol';
import { extractOrAddAtaIxs } from '../swap_api_utils/ata';

export abstract class JupiterBase {
  protected constructor(
    readonly routerType: RouterType,
    readonly connection: Rpc<SolanaRpcApi>,
  ) {}

  async getRouterOutput(
    instructions: SwapInstructionsResponse,
    quoteResponse: QuoteResponse,
    params: RouteParams,
    ctx: RouterContext,
    timeAfterGetQuote: number,
    timeBeforeGetQuote: number,
  ): Promise<RouterOutput> {
    const { setupInstructions, swapInstruction, addressLookupTableAddresses, cleanupInstruction } = instructions;
    const lookupTablePubkeys = addressLookupTableAddresses.map((addr) => address(addr));
    const lookupTableAccounts = await getAllLookupTables(this.connection, lookupTablePubkeys);

    const timeAfterSwapIxsRequest = Date.now();

    const ixsSetupJupiter = setupInstructions.map((ix) => this.convertJupiterToTransactionInstruction(ix));
    const ixsCleanupJupiter = cleanupInstruction
      ? [this.convertJupiterToTransactionInstruction(cleanupInstruction)]
      : [];

    const {
      createInAtaIxs,
      createOutAtaIxs,
      otherIxs: nonAtaIxs,
    } = await extractOrAddAtaIxs(
      params,
      ixsSetupJupiter,
      ctx.inMintInfo.tokenProgramId,
      ctx.outMintInfo.tokenProgramId,
    );
    // ignore remaining ixs
    const { wrapSolIxs, unwrapSolIxs } = await extractOrAddWrapAndUnwrapSolIxs(params, [
      ...nonAtaIxs,
      ...ixsCleanupJupiter,
    ]);

    const priceImpactBps = Number(quoteResponse.priceImpactPct) * 100;

    const routeInstructions: RouterInstructions = {
      createInAtaIxs,
      createOutAtaIxs,
      wrapSolIxs: wrapSolIxs,
      swapIxs: [this.convertJupiterToTransactionInstruction(swapInstruction)],
      unwrapSolIxs: unwrapSolIxs,
    };

    return {
      instructions: routeInstructions,
      ixsRouter: compileRouteIxs(routeInstructions),
      amountsExactIn: {
        amountIn: params.swapType === 'exactIn' ? params.amount : new BN(0),
        amountOutGuaranteed: params.swapType === 'exactIn' ? new BN(quoteResponse.otherAmountThreshold) : new BN(0),
        amountOut: params.swapType === 'exactIn' ? new BN(quoteResponse.outAmount) : new BN(0),
      },
      amountsExactOut: {
        amountOut: params.swapType === 'exactOut' ? params.amount : new BN(0),
        amountInGuaranteed:
          params.swapType === 'exactOut' ? new BN(quoteResponse.otherAmountThreshold.toString()) : new BN(0),
        amountIn: params.swapType === 'exactOut' ? new BN(quoteResponse.inAmount.toString()) : new BN(0),
      },
      swapType: params.swapType,
      responseTimeGetQuoteMs: timeAfterGetQuote - timeBeforeGetQuote,
      responseTimeSwapIxsMs: timeAfterSwapIxsRequest - timeAfterGetQuote,
      lookupTableAccounts: lookupTableAccounts,
      routerType: this.routerType,
      expiryTime: 0,
      priceImpactBps,
    };
  }

  async getQuoteOutput(
    quoteResponse: QuoteResponse,
    params: QuoteParams,
    timeAfterGetQuoteData: number,
    timeBeforeGetQuoteData: number,
  ) {
    const priceImpactBps = Number(quoteResponse.priceImpactPct) * 100;

    return {
      amountsExactIn: {
        amountIn: params.swapType === 'exactIn' ? params.amount : new BN(0),
        amountOutGuaranteed: params.swapType === 'exactIn' ? new BN(quoteResponse.otherAmountThreshold) : new BN(0),
        amountOut: params.swapType === 'exactIn' ? new BN(quoteResponse.outAmount) : new BN(0),
      },
      amountsExactOut: {
        amountOut: params.swapType === 'exactOut' ? params.amount : new BN(0),
        amountInGuaranteed:
          params.swapType === 'exactOut' ? new BN(quoteResponse.otherAmountThreshold.toString()) : new BN(0),
        amountIn: params.swapType === 'exactOut' ? new BN(quoteResponse.outAmount.toString()) : new BN(0),
      },
      swapType: params.swapType,
      responseTimeGetQuoteMs: timeAfterGetQuoteData - timeBeforeGetQuoteData,
      responseTimeSwapIxsMs: 0,
      routerType: this.routerType,
      expiryTime: 0,
      priceImpactBps,
    };
  }

  private convertJupiterToTransactionInstruction(instruction: JupiterInstruction): Instruction {
    return {
      programAddress: address(instruction.programId),
      accounts: instruction.accounts.map((key) => ({
        address: address(key.pubkey),
        role: determineRole(key.isSigner, key.isWritable),
      })),
      data: Buffer.from(instruction.data, 'base64'),
    };
  }
}