import { Connection, PublicKey } from "@solana/web3.js";
import { Address, translateAddress } from "@coral-xyz/anchor";
import invariant from "tiny-invariant";
import { ParsableConfig, ParsableEntity, ParsableFlexibleVault } from "./parsing";

/**
 * Supported accounts
 */
type CachedValue = {};

/**
 * Include both the entity (i.e. type) of the stored value, and the value itself
 */
interface CachedContent<T extends CachedValue> {
  entity: ParsableEntity<T>;
  value: CachedValue | null;
}

/**
 * Type for rpc batch request response
 */
type GetMultipleAccountsResponse = {
  error?: string;
  result?: {
    value?: ({ data: [string, string] } | null)[];
  };
};

export class AccountFetcher {
  private readonly connection: Connection;
  private readonly _cache: Record<string, CachedContent<CachedValue>> = {};

  constructor(connection: Connection, cache?: Record<string, CachedContent<CachedValue>>) {
    this.connection = connection;
    this._cache = cache ?? {};
  }

  /*** Public Methods ***/

  /**
   * Update the cached value of all entities currently in the cache.
   * Uses batched rpc request for network efficient fetch.
   */
  public async refreshAll(usingV2: boolean): Promise<void> {
    const addresses: string[] = Object.keys(this._cache);
    const data = await this.bulkRequest(addresses);

    for (const [idx, [key, cachedContent]] of Object.entries(this._cache).entries()) {
      const entity = cachedContent.entity;
      const value = entity.parse(usingV2, data[idx]);

      this._cache[key] = { entity, value };
    }
  }

  /*** Private Methods ***/

  /**
   * Retrieve from cache or fetch from rpc, an account
   */
  private async get<T extends CachedValue>(
    usingV2: boolean,
    address: PublicKey,
    entity: ParsableEntity<T>,
    refresh: boolean,
  ): Promise<T | null> {
    const key = address.toBase58();
    const cachedValue: CachedValue | null | undefined = this._cache[key]?.value;

    if (cachedValue !== undefined && !refresh) {
      return cachedValue as T | null;
    }

    const accountInfo = await this.connection.getAccountInfo(address);
    if (!accountInfo) {
      return null;
    }
    const accountData = accountInfo?.data;
    const parsedValue = entity.parse(usingV2, accountData);

    // Convert snake_case to camelCase
    const value = parsedValue ? Object.fromEntries(
      Object.entries(parsedValue).map(([k, v]) => [
        k.replace(/_([a-z])/g, (_, letter) => letter.toUpperCase()),
        v
      ])
    ) : null;

    this._cache[key] = { entity, value };

    return value as T | null;
  }


  /**
   * Make batch rpc request
   */
  private async bulkRequest(addresses: string[]): Promise<(Buffer | null)[]> {
    const responses: Promise<GetMultipleAccountsResponse>[] = [];
    const chunk = 100; // getMultipleAccounts has limitation of 100 accounts per request

    for (let i = 0; i < addresses.length; i += chunk) {
      const addressesSubset = addresses.slice(i, i + chunk);
      const res = (this.connection as any)._rpcRequest("getMultipleAccounts", [
        addressesSubset,
        { commitment: this.connection.commitment },
      ]);
      responses.push(res);
    }

    const combinedResult: (Buffer | null)[] = [];

    (await Promise.all(responses)).forEach((res) => {
      invariant(!res.error, `bulkRequest result error: ${res.error}`);
      invariant(!!res.result?.value, "bulkRequest no value");

      res.result.value.forEach((account) => {
        if (!account || account.data[1] !== "base64") {
          combinedResult.push(null);
        } else {
          combinedResult.push(Buffer.from(account.data[0], account.data[1]));
        }
      });
    });

    invariant(combinedResult.length === addresses.length, "bulkRequest not enough results");
    return combinedResult;
  }

  // YOUR CODE:
  /**
   * Parse data from account info
   */
  public parseAccountData(usingV2: boolean, data: Buffer | null | undefined, entity: ParsableEntity<CachedValue>) {
    if (!data) {
      return null;
    }

    const parsedValue = entity.parse(usingV2, data);

    // Convert snake_case to camelCase
    const value = parsedValue ? Object.fromEntries(
      Object.entries(parsedValue).map(([k, v]) => [
        k.replace(/_([a-z])/g, (_, letter) => letter.toUpperCase()),
        v
      ])
    ) : null;

    return value;
  }

  public async getFlexibleVault(address: Address, refresh: boolean, usingV2 = false) {
    return this.get(usingV2, translateAddress(address), ParsableFlexibleVault, refresh);
  }

  public async getConfig(address: Address, refresh: boolean, usingV2 = false) {
    return this.get(usingV2, translateAddress(address), ParsableConfig, refresh);
  }
}