import {
  GetJupiterPriceParams,
  GetJupiterPriceResponse,
  IRouterContext,
  QuoteParams,
  RouteOutput,
  RouteParams,
  safeBNToNumber,
  SwapType,
} from '../swap_api_utils';
import axios from 'axios';
import BN from 'bn.js';
import { RouterType, JUP_ULTRA_BASE_URL, ROUTER_PROGRAM_MAP, JUPITER_PRICE_API } from '../consts';
import { removeComputeBudgetProgram } from '../utils';
import { base64ToTransaction, extractInstructionsAndLookupTablesFromTxn } from '../utils/decodeTransaction';
import { Router } from '../swap_api_utils';
import { Logger } from '../utils/Logger';
import { Account, Instruction, Rpc, SolanaRpcApi } from '@solana/kit';
import { AddressLookupTable } from '@solana-program/address-lookup-table';

export class JupiterUltraRouter implements Router {
  readonly connection: Rpc<SolanaRpcApi>;
  private readonly baseUrl = JUP_ULTRA_BASE_URL;
  readonly routerType: RouterType;
  private logger: Logger;

  constructor(connection: Rpc<SolanaRpcApi>, logger: Logger) {
    this.connection = connection;
    this.routerType = 'jupiterU';
    this.logger = logger;
  }

  async quote(params: QuoteParams): Promise<RouteOutput[] | undefined> {
    try {
      if (params.swapType === 'exactOut') {
        return undefined;
      }
      const timeBeforeGetQuoteData = Date.now();
      const swapQuote = await this.getSwapQuote({
        inputMint: params.tokenIn.toString(),
        outputMint: params.tokenOut.toString(),
        amount: safeBNToNumber(params.amount),
      });
      const timeAfterGetQuoteData = Date.now();

      if (swapQuote.inAmount !== params.amount.toString()) {
        return undefined;
      }

      return [this.buildRouteOutput(swapQuote, params.swapType, timeAfterGetQuoteData - timeBeforeGetQuoteData, 0)];
    } catch (error) {
      this.logger.error(`JupUltra: Failed to get quote`, error);
      return undefined;
    }
  }

  async route(params: RouteParams, _ctx: IRouterContext): Promise<RouteOutput[] | undefined> {
    try {
      if (params.swapType === 'exactOut') {
        return undefined;
      }
      if (params.destinationTokenAccount) {
        return undefined;
      }
      if (params.wrapAndUnwrapSol === false || params.includeSetupIxs === false) {
        return undefined;
      }
      const timeBeforeGetQuote = Date.now();
      const swapQuoteResponse = await this.getSwapQuote({
        inputMint: params.tokenIn.toString(),
        outputMint: params.tokenOut.toString(),
        amount: safeBNToNumber(params.amount),
        taker: params.executor.toString(),
      });
      const timeAfterGetQuote = Date.now();
      if (swapQuoteResponse.inAmount !== params.amount.toString()) {
        return undefined;
      }

      if (!swapQuoteResponse.transaction) {
        this.logger.error(`JupUltra: Swap transaction is null`);
        return undefined;
      }
      const priceImpactBps = Number(swapQuoteResponse.priceImpactPct) * 100;
      const transaction = base64ToTransaction(swapQuoteResponse.transaction);

      let instructionsWithoutCu: Instruction[] | undefined = undefined;
      let lookupTablesAccounts: Account<AddressLookupTable>[] | undefined = undefined;

      // In case of aggregator swap, we need to keep the transaction and execute it
      const { instructions, lookupTables } = await extractInstructionsAndLookupTablesFromTxn(
        this.connection,
        transaction,
      );
      instructionsWithoutCu = removeComputeBudgetProgram(instructions);
      let foundValidSwap = false;
      for (const ix of instructionsWithoutCu) {
        if (
          ix.programAddress === ROUTER_PROGRAM_MAP['jupiterU'] ||
          ix.programAddress === ROUTER_PROGRAM_MAP['jupiterZ']
        ) {
          foundValidSwap = true;
          break;
        }
      }
      if (!foundValidSwap) {
        return undefined;
      }
      lookupTablesAccounts = lookupTables;

      const timeAfterSwapIxs = Date.now();

      return [
        {
          ...this.buildRouteOutput(
            swapQuoteResponse,
            params.swapType,
            timeAfterGetQuote - timeBeforeGetQuote,
            timeAfterSwapIxs - timeAfterGetQuote,
          ),
          transaction: transaction,
          ixsRouter: instructionsWithoutCu,
          lookupTableAccounts: lookupTablesAccounts,
          priceImpactBps,
        },
      ];
    } catch (error) {
      this.logger.error(`${this.routerType}: Failed to route swap`, error);
      return undefined;
    }
  }

  private async getSwapQuote(params: SwapQuoteParams): Promise<SwapQuoteResponse> {
    try {
      const response = await axios.get<SwapQuoteResponse>(`${this.baseUrl}/order`, {
        params,
        headers: {},
      });
      return response.data;
    } catch (error) {
      this.logger.error(`${this.routerType}: Failed to fetch swap quote`, error);
      throw new Error(`${this.routerType}: Failed to fetch swap quote`, { cause: error });
    }
  }

  private buildRouteOutput(
    swapQuote: SwapQuoteResponse,
    swapType: SwapType,
    responseTimeGetQuoteMs: number,
    responseTimeSwapIxsMs: number,
  ): RouteOutput {
    return {
      amountsExactIn: {
        amountIn: new BN(swapQuote.inAmount),
        amountOutGuaranteed: new BN(swapQuote.otherAmountThreshold),
        amountOut: new BN(swapQuote.outAmount),
      },
      amountsExactOut: {
        amountOut: new BN(0),
        amountInGuaranteed: new BN(0),
        amountIn: new BN(0),
      },
      swapType: swapType,
      responseTimeGetQuoteMs,
      responseTimeSwapIxsMs,
      routerType: swapQuote.swapType === 'aggregator' ? 'jupiterU' : 'jupiterZ',
      expiryTime: 0,
      jupRequestId: swapQuote.requestId,
    };
  }
}

export async function executeJupiterZTransaction(params: ExecuteParams): Promise<ExecuteResponse> {
  try {
    const response = await axios.post<ExecuteResponse>(
      `${JUP_ULTRA_BASE_URL}/execute`,
      {
        signedTransaction: params.signedTransaction,
        requestId: params.requestId,
      },
      {
        headers: {
          'Content-Type': 'application/json',
          Accept: 'application/json',
        },
      },
    );

    return response.data;
  } catch (error) {
    params.logger.error('Execute transaction error:', error);
    throw new Error(`Failed to execute transaction`, { cause: error });
  }
}

export async function fetchJupiterPrice(query: GetJupiterPriceParams): Promise<GetJupiterPriceResponse> {
  const response = await axios.get<GetJupiterPriceResponse>(JUPITER_PRICE_API, {
    params: query,
    headers: { 'Content-Type': 'application/json' },
  });
  return response.data;
}

interface SwapQuoteParams {
  inputMint: string;
  outputMint: string;
  amount: number;
  taker?: string;
}

interface SwapQuoteResponse {
  inputMint: string;
  outputMint: string;
  inAmount: string;
  outAmount: string;
  otherAmountThreshold: string;
  swapMode: 'ExactIn' | 'ExactOut';
  slippageBps: number;
  priceImpactPct: string;
  routePlan: RoutePlan[];
  feeBps: number;
  transaction: string | null;
  gasless: boolean;
  prioritizationType: 'None' | 'ComputeBudget' | 'Jito';
  prioritizationFeeLamports: number;
  requestId: string;
  swapType: 'rfq' | 'aggregator';
  quoteId: string;
  maker?: string | null;
  taker?: string | null;
  expireAt?: number | null;
  platformFee?: {
    amount: string;
    feeBps: number;
  };
  totalTime?: number | null;
}

interface RoutePlan {
  swapInfo: SwapInfo;
  percent: number;
}

interface SwapInfo {
  ammKey: string;
  label: string;
  inputMint: string;
  outputMint: string;
  inAmount: string;
  outAmount: string;
  feeAmount: string;
  feeMint: string;
}

interface ExecuteParams {
  signedTransaction: string;
  requestId: string;
  logger: Logger;
}

interface SwapEvent {
  inputMint: string;
  inputAmount: string;
  outputMint: string;
  outputAmount: string;
}

interface ExecuteResponse {
  status: 'Success' | 'Failed';
  signature?: string;
  slot?: string;
  error?: string;
  code: number;
  totalInputAmount?: string;
  totalOutputAmount?: string;
  inputAmountResult?: string;
  outputAmountResult?: string;
  swapEvents?: SwapEvent[];
}