import {routes} from "@lodestar/api";
import {IForkChoice, ProtoBlock} from "@lodestar/fork-choice";
import {IBeaconStateView, computeEpochAtSlot} from "@lodestar/state-transition";
import {BeaconBlock, Epoch, RootHex, Slot, phase0} from "@lodestar/types";
import {Logger, toRootHex} from "@lodestar/utils";
import {Metrics} from "../../metrics/index.js";
import {JobItemQueue} from "../../util/queue/index.js";
import {BlockStateCache, CheckpointHex, CheckpointStateCache} from "../stateCache/types.js";
import {RegenError, RegenErrorCode} from "./errors.js";
import {
  IStateRegenerator,
  IStateRegeneratorInternal,
  RegenCaller,
  RegenFnName,
  StateRegenerationOpts,
} from "./interface.js";
import {RegenModules, StateRegenerator} from "./regen.js";

const REGEN_QUEUE_MAX_LEN = 256;
// TODO: Should this constant be lower than above? 256 feels high
const REGEN_CAN_ACCEPT_WORK_THRESHOLD = 16;

type QueuedStateRegeneratorModules = RegenModules & {
  signal: AbortSignal;
};

type RegenRequestKey = keyof IStateRegeneratorInternal;
type RegenRequestByKey = {[K in RegenRequestKey]: {key: K; args: Parameters<IStateRegeneratorInternal[K]>}};
export type RegenRequest = RegenRequestByKey[RegenRequestKey];

/**
 * Regenerates states that have already been processed by the fork choice
 *
 * All requests are queued so that only a single state at a time may be regenerated at a time
 */
export class QueuedStateRegenerator implements IStateRegenerator {
  readonly jobQueue: JobItemQueue<[RegenRequest], IBeaconStateView>;
  private readonly regen: StateRegenerator;

  private readonly forkChoice: IForkChoice;
  private readonly blockStateCache: BlockStateCache;
  private readonly checkpointStateCache: CheckpointStateCache;
  private readonly metrics: Metrics | null;
  private readonly logger: Logger;

  constructor(modules: QueuedStateRegeneratorModules) {
    this.regen = new StateRegenerator(modules);
    this.jobQueue = new JobItemQueue<[RegenRequest], IBeaconStateView>(
      this.jobQueueProcessor,
      {maxLength: REGEN_QUEUE_MAX_LEN, signal: modules.signal},
      modules.metrics ? modules.metrics.regenQueue : undefined
    );
    this.forkChoice = modules.forkChoice;
    this.blockStateCache = modules.blockStateCache;
    this.checkpointStateCache = modules.checkpointStateCache;
    this.metrics = modules.metrics;
    this.logger = modules.logger;
  }

  async init(): Promise<void> {
    if (this.checkpointStateCache.init) {
      return this.checkpointStateCache.init();
    }
  }

  canAcceptWork(): boolean {
    return this.jobQueue.jobLen < REGEN_CAN_ACCEPT_WORK_THRESHOLD;
  }

  dropCache(): void {
    this.blockStateCache.clear();
    this.checkpointStateCache.clear();
  }

  dumpCacheSummary(): routes.lodestar.StateCacheItem[] {
    return [...this.blockStateCache.dumpSummary(), ...this.checkpointStateCache.dumpSummary()];
  }

  /**
   * Get a state from block state cache.
   */
  getStateSync(stateRoot: RootHex): IBeaconStateView | null {
    return this.blockStateCache.get(stateRoot);
  }

  /**
   * Get state for block processing.
   */
  getPreStateSync(block: BeaconBlock): IBeaconStateView | null {
    const parentRoot = toRootHex(block.parentRoot);
    const parentBlock = this.forkChoice.getBlockHexDefaultStatus(parentRoot);
    if (!parentBlock) {
      throw new RegenError({
        code: RegenErrorCode.BLOCK_NOT_IN_FORKCHOICE,
        blockRoot: block.parentRoot,
      });
    }

    const parentEpoch = computeEpochAtSlot(parentBlock.slot);
    const blockEpoch = computeEpochAtSlot(block.slot);

    // Check the checkpoint cache (if the pre-state is a checkpoint state)
    if (parentEpoch < blockEpoch) {
      const checkpointState = this.checkpointStateCache.getLatest(parentRoot, blockEpoch);
      if (checkpointState && computeEpochAtSlot(checkpointState.slot) === blockEpoch) {
        return checkpointState;
      }
    }

    // Check the state cache, only if the state doesn't need to go through an epoch transition.
    // Otherwise the state transition may not be cached and wasted. Queue for regen since the
    // work required will still be significant.
    if (parentEpoch === blockEpoch) {
      const state = this.blockStateCache.get(parentBlock.stateRoot);
      if (state) {
        return state;
      }
    }

    return null;
  }

  async getCheckpointStateOrBytes(cp: CheckpointHex): Promise<IBeaconStateView | Uint8Array | null> {
    return this.checkpointStateCache.getStateOrBytes(cp);
  }

  /**
   * Get checkpoint state from cache
   */
  getCheckpointStateSync(cp: CheckpointHex): IBeaconStateView | null {
    return this.checkpointStateCache.get(cp);
  }

  /**
   * Get state closest to head
   */
  getClosestHeadState(head: ProtoBlock): IBeaconStateView | null {
    return this.checkpointStateCache.getLatest(head.blockRoot, Infinity) || this.blockStateCache.get(head.stateRoot);
  }

  pruneOnCheckpoint(finalizedEpoch: Epoch, justifiedEpoch: Epoch, headStateRoot: RootHex): void {
    this.checkpointStateCache.prune(finalizedEpoch, justifiedEpoch);
    this.blockStateCache.prune(headStateRoot);
  }

  pruneOnFinalized(finalizedEpoch: number): void {
    this.checkpointStateCache.pruneFinalized(finalizedEpoch);
    this.blockStateCache.deleteAllBeforeEpoch(finalizedEpoch);
  }

  processState(blockRootHex: RootHex, postState: IBeaconStateView): void {
    this.blockStateCache.add(postState);
    this.checkpointStateCache.processState(blockRootHex, postState).catch((e) => {
      this.logger.debug("Error processing block state", {blockRootHex, slot: postState.slot}, e);
    });
  }

  addCheckpointState(cp: phase0.Checkpoint, item: IBeaconStateView): void {
    this.checkpointStateCache.add(cp, item);
  }

  updateHeadState(newHead: ProtoBlock, maybeHeadState: IBeaconStateView): void {
    const {stateRoot: newHeadStateRoot, blockRoot: newHeadBlockRoot, slot: newHeadSlot} = newHead;
    const maybeHeadStateRoot = toRootHex(maybeHeadState.hashTreeRoot());
    const logCtx = {
      newHeadSlot,
      newHeadBlockRoot,
      newHeadStateRoot,
      maybeHeadSlot: maybeHeadState.slot,
      maybeHeadStateRoot,
    };
    const headState =
      newHeadStateRoot === maybeHeadStateRoot ? maybeHeadState : this.blockStateCache.get(newHeadStateRoot);

    if (headState) {
      this.blockStateCache.setHeadState(headState);
    } else {
      // Trigger regen on head change if necessary
      this.logger.warn("Head state not available, triggering regen", logCtx);
      // for the old BlockStateCacheImpl only
      //     - head has changed, so the existing cached head state is no longer useful. Set strong reference to null to free
      //       up memory for regen step below. During regen, node won't be functional but eventually head will be available
      // for the new FIFOBlockStateCache, this has no affect
      this.blockStateCache.setHeadState(null);

      // for the new FIFOBlockStateCache, it's important to reload state to regen head state here if needed
      const allowDiskReload = true;
      this.regen.getState(newHeadStateRoot, RegenCaller.processBlock, allowDiskReload).then(
        (headStateRegen) => this.blockStateCache.setHeadState(headStateRegen),
        (e) => this.logger.error("Error on head state regen", logCtx, e)
      );
    }
  }

  updatePreComputedCheckpoint(rootHex: RootHex, epoch: Epoch): number | null {
    return this.checkpointStateCache.updatePreComputedCheckpoint(rootHex, epoch);
  }

  /**
   * Get the state to run with `block`.
   * - State after `block.parentRoot` dialed forward to block.slot
   */
  async getPreState(block: BeaconBlock, opts: StateRegenerationOpts, rCaller: RegenCaller): Promise<IBeaconStateView> {
    this.metrics?.regenFnCallTotal.inc({caller: rCaller, entrypoint: RegenFnName.getPreState});

    // First attempt to fetch the state from caches before queueing
    const cachedState = this.getPreStateSync(block);

    if (cachedState !== null) {
      return cachedState;
    }

    // The state is not immediately available in the caches, enqueue the job
    this.metrics?.regenFnQueuedTotal.inc({caller: rCaller, entrypoint: RegenFnName.getPreState});
    return this.jobQueue.push({key: "getPreState", args: [block, opts, rCaller]});
  }

  /**
   * Get state of provided `blockRoot` and dial forward to `slot`
   * Use this api with care because we don't want the queue to be busy
   * For the context, gossip block validation uses this api so we want it to be as fast as possible
   * @returns
   */
  async getBlockSlotState(
    block: ProtoBlock,
    slot: Slot,
    opts: StateRegenerationOpts,
    rCaller: RegenCaller
  ): Promise<IBeaconStateView> {
    this.metrics?.regenFnCallTotal.inc({caller: rCaller, entrypoint: RegenFnName.getBlockSlotState});

    // The state is not immediately available in the caches, enqueue the job
    return this.jobQueue.push({key: "getBlockSlotState", args: [block, slot, opts, rCaller]});
  }

  async getState(stateRoot: RootHex, rCaller: RegenCaller): Promise<IBeaconStateView> {
    this.metrics?.regenFnCallTotal.inc({caller: rCaller, entrypoint: RegenFnName.getState});

    // First attempt to fetch the state from cache before queueing
    const state = this.blockStateCache.get(stateRoot);
    if (state) {
      return state;
    }

    // The state is not immediately available in the cache, enqueue the job
    this.metrics?.regenFnQueuedTotal.inc({caller: rCaller, entrypoint: RegenFnName.getState});
    return this.jobQueue.push({key: "getState", args: [stateRoot, rCaller]});
  }

  private jobQueueProcessor = async (regenRequest: RegenRequest): Promise<IBeaconStateView> => {
    const metricsLabels = {
      caller: regenRequest.args.at(-1) as RegenCaller,
      entrypoint: regenRequest.key as RegenFnName,
    };
    let timer: (() => number) | undefined;
    try {
      timer = this.metrics?.regenFnCallDuration.startTimer(metricsLabels);
      switch (regenRequest.key) {
        case "getPreState":
          return await this.regen.getPreState(...regenRequest.args);
        case "getBlockSlotState":
          return await this.regen.getBlockSlotState(...regenRequest.args);
        case "getState":
          return await this.regen.getState(...regenRequest.args);
      }
    } catch (e) {
      this.metrics?.regenFnTotalErrors.inc(metricsLabels);
      throw e;
    } finally {
      if (timer) timer();
    }
  };
}