import { MintInfo, SOL_MINT_INFO } from './types';
import { WRAPPED_SOL_MINT } from '@kamino-finance/limo-sdk';
import { Address, Rpc, SolanaRpcApi } from '@solana/kit';
import { fetchAllMaybeMint, fetchMaybeMint } from '@solana-program/token-2022';

export interface IRouterContext {
  get inMintInfo(): MintInfo;
  get outMintInfo(): MintInfo;
}

export class RouterContext implements IRouterContext {
  constructor(
    public inMintInfo: MintInfo,
    public outMintInfo: MintInfo,
  ) {}
}

export class NoRouterContext implements IRouterContext {
  get inMintInfo(): MintInfo {
    throw new Error(`NoRouterContext: inMintInfo is not available`);
  }
  get outMintInfo(): MintInfo {
    throw new Error(`NoRouterContext: outMintInfo is not available`);
  }
}

export async function loadRouterContext(
  connection: Rpc<SolanaRpcApi>,
  fromToken: Address,
  toToken: Address,
): Promise<IRouterContext> {
  let fromTokenInfo: MintInfo | undefined;
  let toTokenInfo: MintInfo | undefined;
  if (fromToken === WRAPPED_SOL_MINT) {
    fromTokenInfo = SOL_MINT_INFO;
  }
  if (toToken === WRAPPED_SOL_MINT) {
    toTokenInfo = SOL_MINT_INFO;
  }
  if (fromTokenInfo && toTokenInfo) {
    return new RouterContext(fromTokenInfo, toTokenInfo);
  } else if (toTokenInfo) {
    const fromTokenInfo = await fetchSingleMintInfo(connection, fromToken);
    return new RouterContext(fromTokenInfo, toTokenInfo);
  } else if (fromTokenInfo) {
    const toTokenInfo = await fetchSingleMintInfo(connection, toToken);
    return new RouterContext(fromTokenInfo, toTokenInfo);
  }
  const mintAccounts = await fetchAllMaybeMint(connection, [fromToken, toToken]);
  const fromMintAcc = mintAccounts[0];
  const toMintAcc = mintAccounts[1];
  if (!fromMintAcc.exists) {
    throw new Error(`Failed to fetch mint account info ${fromToken}`);
  }
  if (!toMintAcc.exists) {
    throw new Error(`Failed to fetch mint account info ${toToken}`);
  }
  return new RouterContext(
    { tokenProgramId: fromMintAcc.programAddress, decimals: fromMintAcc.data.decimals },
    { tokenProgramId: toMintAcc.programAddress, decimals: toMintAcc.data.decimals },
  );
}

async function fetchSingleMintInfo(connection: Rpc<SolanaRpcApi>, mint: Address): Promise<MintInfo> {
  const mintInfo = await fetchMaybeMint(connection, mint);
  if (!mintInfo.exists) {
    throw new Error(`Failed to fetch mint account info ${mint}`);
  }
  return {
    tokenProgramId: mintInfo.programAddress,
    decimals: mintInfo.data.decimals,
  };
}