import {pubkeyCache} from "@chainsafe/lodestar-z/pubkeys";
import {ChainForkConfig, createBeaconConfig} from "@lodestar/config";
import {config, config as minimalConfig} from "@lodestar/config/default";
import {
  EPOCHS_PER_HISTORICAL_VECTOR,
  EPOCHS_PER_SLASHINGS_VECTOR,
  GENESIS_EPOCH,
  GENESIS_SLOT,
  SLOTS_PER_HISTORICAL_ROOT,
} from "@lodestar/params";
import {phase0, ssz} from "@lodestar/types";
import {EpochCacheOpts} from "../cache/epochCache.js";
import {BeaconStateCache} from "../cache/stateCache.js";
import {ZERO_HASH} from "../constants/index.js";
import {BeaconStateAllForks, BeaconStatePhase0, CachedBeaconStateAllForks, createCachedBeaconState} from "../index.js";
import {newZeroedArray} from "../util/index.js";

/**
 * Copy of BeaconState, but all fields are marked optional to allow for swapping out variables as needed.
 */
type TestBeaconState = Partial<phase0.BeaconState>;

/**
 * Generate beaconState, by default it will use the initial state defined when the `ChainStart` log is emitted.
 * NOTE: All fields can be overridden through `opts`.
 * @param {TestBeaconState} opts
 * @returns {BeaconState}
 */
export function generateState(opts?: TestBeaconState): BeaconStatePhase0 {
  return ssz.phase0.BeaconState.toViewDU({
    genesisTime: Math.floor(Date.now() / 1000),
    genesisValidatorsRoot: ZERO_HASH,
    slot: GENESIS_SLOT,
    fork: {
      previousVersion: config.GENESIS_FORK_VERSION,
      currentVersion: config.GENESIS_FORK_VERSION,
      epoch: GENESIS_EPOCH,
    },
    latestBlockHeader: {
      slot: 0,
      proposerIndex: 0,
      parentRoot: Buffer.alloc(32),
      stateRoot: Buffer.alloc(32),
      bodyRoot: ssz.phase0.BeaconBlockBody.hashTreeRoot(ssz.phase0.BeaconBlockBody.defaultValue()),
    },
    blockRoots: Array.from({length: SLOTS_PER_HISTORICAL_ROOT}, () => ZERO_HASH),
    stateRoots: Array.from({length: SLOTS_PER_HISTORICAL_ROOT}, () => ZERO_HASH),
    historicalRoots: [],
    eth1Data: {
      depositRoot: Buffer.alloc(32),
      blockHash: Buffer.alloc(32),
      depositCount: 0n,
    },
    eth1DataVotes: [],
    eth1DepositIndex: 0,
    validators: [],
    balances: [],
    randaoMixes: Array.from({length: EPOCHS_PER_HISTORICAL_VECTOR}, () => ZERO_HASH),
    slashings: newZeroedArray(EPOCHS_PER_SLASHINGS_VECTOR),
    previousEpochAttestations: [],
    currentEpochAttestations: [],
    justificationBits: ssz.phase0.JustificationBits.defaultValue(),
    previousJustifiedCheckpoint: {
      epoch: GENESIS_EPOCH,
      root: ZERO_HASH,
    },
    currentJustifiedCheckpoint: {
      epoch: GENESIS_EPOCH,
      root: ZERO_HASH,
    },
    finalizedCheckpoint: {
      epoch: GENESIS_EPOCH,
      root: ZERO_HASH,
    },
    ...opts,
  });
}

export function generateCachedState(
  config: ChainForkConfig = minimalConfig,
  opts: TestBeaconState = {}
): CachedBeaconStateAllForks {
  const state = generateState(opts);
  pubkeyCache.reset();
  return createCachedBeaconState(state, {
    config: createBeaconConfig(config, state.genesisValidatorsRoot),
    pubkeyCache,
  });
}

export function createCachedBeaconStateTest<T extends BeaconStateAllForks>(
  state: T,
  configCustom: ChainForkConfig = config,
  opts?: EpochCacheOpts
): T & BeaconStateCache {
  return createCachedBeaconState<T>(
    state,
    {
      config: createBeaconConfig(configCustom, state.genesisValidatorsRoot),
      pubkeyCache,
    },
    opts
  );
}