import { Address, Account, Rpc, SolanaRpcApi, AddressesByLookupTableAddress } from '@solana/kit';
import { AddressLookupTable, fetchAllAddressLookupTable } from '@solana-program/address-lookup-table';

// In-memory cache of decoded LUT accounts. Fetching+decoding a LUT base58-encodes up to 256
// addresses, and popular LUTs (Jupiter's, the kswap LUT) are re-fetched on every route/request
// — that was the single biggest main-thread cost. LUTs are append-only, so existing addresses
// never move; a short TTL only risks missing freshly-appended addresses until it refreshes,
// which is harmless for routing/simulation. Concurrent misses for the same address may double-
// fetch on a cold cache (acceptable); the warm cache dedupes.
const LUT_CACHE_TTL_MS = 60_000;
const LUT_CACHE_MAX_ENTRIES = 5_000;
const lutCache = new Map<string, { account: Account<AddressLookupTable>; expiresAt: number }>();

export async function getAllLookupTables(
  connection: Rpc<SolanaRpcApi>,
  addresses: Address[],
): Promise<Account<AddressLookupTable>[]> {
  if (addresses.length === 0) {
    return [];
  }

  const now = Date.now();
  const missing: Address[] = [];
  for (const address of addresses) {
    const cached = lutCache.get(address);
    if (!cached || cached.expiresAt <= now) {
      missing.push(address);
    }
  }

  if (missing.length > 0) {
    // De-dupe the misses so a route referencing the same LUT twice doesn't fetch it twice.
    const uniqueMissing = [...new Set(missing)];
    const fetched = await fetchAllAddressLookupTable(connection, uniqueMissing);
    const expiresAt = Date.now() + LUT_CACHE_TTL_MS;
    for (const account of fetched) {
      lutCache.set(account.address, { account, expiresAt });
    }
    if (lutCache.size > LUT_CACHE_MAX_ENTRIES) {
      for (const [key, entry] of lutCache) {
        if (entry.expiresAt <= now) lutCache.delete(key);
      }
    }
  }

  // Return one account per input address, in input order, skipping any that don't exist
  // (mirrors fetchAllAddressLookupTable, which only returns existing tables).
  const result: Account<AddressLookupTable>[] = [];
  for (const address of addresses) {
    const cached = lutCache.get(address);
    if (cached) {
      result.push(cached.account);
    }
  }
  return result;
}

export function transformFromLookupTableAccounts(
  lookupTableAccounts: Account<AddressLookupTable>[],
): AddressesByLookupTableAddress {
  const addressesByLookupTable: AddressesByLookupTableAddress = {};
  for (const account of lookupTableAccounts) {
    addressesByLookupTable[account.address] = account.data.addresses;
  }
  return addressesByLookupTable;
}