import {ResponseIncoming} from "@lodestar/reqresp";
import {SignedBeaconBlock, phase0} from "@lodestar/types";
import {LodestarError} from "@lodestar/utils";
import {SerializedCache} from "../../../util/serializedCache.js";
import {ReqRespMethod, responseSszTypeByMethod} from "../types.js";
import {sszDeserializeResponse} from "./collect.js";

/**
 * Asserts a response from BeaconBlocksByRange respects the request and is sequential
 * Note: MUST allow missing block for skipped slots.
 */
export async function collectSequentialBlocksInRange(
  blockStream: AsyncIterable<ResponseIncoming>,
  {count, startSlot}: Pick<phase0.BeaconBlocksByRangeRequest, "count" | "startSlot">,
  serializedCache?: SerializedCache
): Promise<SignedBeaconBlock[]> {
  const blocks: SignedBeaconBlock[] = [];

  for await (const chunk of blockStream) {
    const blockType = responseSszTypeByMethod[ReqRespMethod.BeaconBlocksByRange](chunk.fork, chunk.protocolVersion);
    const block = sszDeserializeResponse(blockType, chunk.data);

    const blockSlot = block.message.slot;

    // Note: step is deprecated and assumed to be 1
    if (blockSlot >= startSlot + count) {
      throw new BlocksByRangeError({code: BlocksByRangeErrorCode.OVER_MAX_SLOT});
    }

    if (blockSlot < startSlot) {
      throw new BlocksByRangeError({code: BlocksByRangeErrorCode.UNDER_START_SLOT});
    }

    const prevBlock = blocks.at(-1);
    if (prevBlock && prevBlock.message.slot >= blockSlot) {
      throw new BlocksByRangeError({code: BlocksByRangeErrorCode.BAD_SEQUENCE});
    }

    blocks.push(block);
    // optionally cache the serialized response if the cache is available
    serializedCache?.set(block, chunk.data);
    if (blocks.length >= count) {
      break; // Done, collected all blocks
    }
  }

  return blocks;
}

export enum BlocksByRangeErrorCode {
  UNDER_START_SLOT = "BLOCKS_BY_RANGE_ERROR_UNDER_START_SLOT",
  OVER_MAX_SLOT = "BLOCKS_BY_RANGE_ERROR_OVER_MAX_SLOT",
  BAD_SEQUENCE = "BLOCKS_BY_RANGE_ERROR_BAD_SEQUENCE",
}

type BlocksByRangeErrorType =
  | {code: BlocksByRangeErrorCode.UNDER_START_SLOT}
  | {code: BlocksByRangeErrorCode.OVER_MAX_SLOT}
  | {code: BlocksByRangeErrorCode.BAD_SEQUENCE};

export class BlocksByRangeError extends LodestarError<BlocksByRangeErrorType> {}
