import {Epoch} from "@lodestar/types";
import {routes} from "@lodestar/api";
import {MapDef} from "@lodestar/utils";
import {Metrics} from "../metrics/index.js";

const PROPOSER_PRESERVE_EPOCHS = 2;

export type ProposerPreparationData = routes.validator.ProposerPreparationData;

export class BeaconProposerCache {
  private readonly feeRecipientByValidatorIndex: MapDef<number, {epoch: Epoch; feeRecipient: string}>;
  constructor(
    opts: {suggestedFeeRecipient: string},
    private readonly metrics?: Metrics | null
  ) {
    this.feeRecipientByValidatorIndex = new MapDef(() => ({
      epoch: 0,
      feeRecipient: opts.suggestedFeeRecipient,
    }));
  }

  add(epoch: Epoch, {validatorIndex, feeRecipient}: ProposerPreparationData): void {
    this.feeRecipientByValidatorIndex.set(validatorIndex, {epoch, feeRecipient});
  }

  prune(epoch: Epoch): void {
    // This is not so optimized function, but could maintain a 2d array may be?
    for (const [validatorIndex, feeRecipientEntry] of this.feeRecipientByValidatorIndex.entries()) {
      // We only retain an entry for PROPOSER_PRESERVE_EPOCHS epochs
      if (feeRecipientEntry.epoch + PROPOSER_PRESERVE_EPOCHS < epoch) {
        this.feeRecipientByValidatorIndex.delete(validatorIndex);
      }
    }
  }

  getOrDefault(proposerIndex: number): string {
    return this.feeRecipientByValidatorIndex.getOrDefault(proposerIndex).feeRecipient;
  }

  get(proposerIndex: number): string | undefined {
    return this.feeRecipientByValidatorIndex.get(proposerIndex)?.feeRecipient;
  }
}
