import {
  type MAX_NULLIFIERS_PER_TX,
  MAX_NULLIFIER_READ_REQUESTS_PER_TX,
  type NULLIFIER_TREE_HEIGHT,
} from '@aztec/constants';
import type { Fr } from '@aztec/foundation/curves/bn254';
import { MembershipWitness } from '@aztec/foundation/trees';

import type { NullifierLeafPreimage } from '../../trees/nullifier_leaf.js';
import type { ClaimedLengthArray } from '../claimed_length_array.js';
import type { ScopedNullifier } from '../nullifier.js';
import { NullifierReadRequestHintsBuilder } from './nullifier_read_request_hints.js';
import { ScopedReadRequest } from './read_request.js';
import { PendingReadHint, ReadRequestActionEnum, ReadRequestResetActions } from './read_request_hints.js';

export function isValidNullifierReadRequest(readRequest: ScopedReadRequest, nullifier: ScopedNullifier) {
  return (
    readRequest.value.equals(nullifier.value) &&
    nullifier.contractAddress.equals(readRequest.contractAddress) &&
    readRequest.counter > nullifier.counter
  );
}

interface NullifierMembershipWitnessWithPreimage {
  membershipWitness: MembershipWitness<typeof NULLIFIER_TREE_HEIGHT>;
  leafPreimage: NullifierLeafPreimage;
}

export function getNullifierReadRequestResetActions(
  nullifierReadRequests: ClaimedLengthArray<ScopedReadRequest, typeof MAX_NULLIFIER_READ_REQUESTS_PER_TX>,
  nullifiers: ClaimedLengthArray<ScopedNullifier, typeof MAX_NULLIFIERS_PER_TX>,
): ReadRequestResetActions<typeof MAX_NULLIFIER_READ_REQUESTS_PER_TX> {
  const resetActions = ReadRequestResetActions.empty(MAX_NULLIFIER_READ_REQUESTS_PER_TX);

  const nullifierMap: Map<bigint, { nullifier: ScopedNullifier; index: number }[]> = new Map();
  nullifiers.getActiveItems().forEach((nullifier, index) => {
    const value = nullifier.value.toBigInt();
    const arr = nullifierMap.get(value) ?? [];
    arr.push({ nullifier, index });
    nullifierMap.set(value, arr);
  });

  for (let i = 0; i < nullifierReadRequests.claimedLength; ++i) {
    const readRequest = nullifierReadRequests.array[i];

    if (readRequest.contractAddress.isZero()) {
      // Settled read: empty contract address means resolve against the nullifier tree.
      resetActions.actions[i] = ReadRequestActionEnum.READ_AS_SETTLED;
    } else {
      // Pending read: non-empty contract address means match against a pending nullifier.
      const pendingNullifier = nullifierMap
        .get(readRequest.value.toBigInt())
        ?.find(({ nullifier }) => isValidNullifierReadRequest(readRequest, nullifier));

      if (pendingNullifier) {
        resetActions.actions[i] = ReadRequestActionEnum.READ_AS_PENDING;
        resetActions.pendingReadHints.push(new PendingReadHint(i, pendingNullifier.index));
      }
      // Otherwise, the read request may be resolved by a future nullifier. Leave as NOOP.
    }
  }

  return resetActions;
}

export async function buildNullifierReadRequestHintsFromResetActions<PENDING extends number, SETTLED extends number>(
  oracle: {
    getNullifierMembershipWitness(nullifier: Fr): Promise<NullifierMembershipWitnessWithPreimage>;
  },
  nullifierReadRequests: ClaimedLengthArray<ScopedReadRequest, typeof MAX_NULLIFIER_READ_REQUESTS_PER_TX>,
  resetActions: ReadRequestResetActions<typeof MAX_NULLIFIER_READ_REQUESTS_PER_TX>,
  maxPending: PENDING = MAX_NULLIFIER_READ_REQUESTS_PER_TX as PENDING,
  maxSettled: SETTLED = MAX_NULLIFIER_READ_REQUESTS_PER_TX as SETTLED,
) {
  const builder = new NullifierReadRequestHintsBuilder(maxPending, maxSettled);

  resetActions.pendingReadHints.forEach(hint => {
    builder.addPendingReadRequest(hint.readRequestIndex, hint.pendingValueIndex);
  });

  // Collect all settled read requests
  const settledRequests: { index: number; readRequest: ScopedReadRequest }[] = [];
  for (let i = 0; i < resetActions.actions.length; i++) {
    if (resetActions.actions[i] === ReadRequestActionEnum.READ_AS_SETTLED) {
      settledRequests.push({ index: i, readRequest: nullifierReadRequests.array[i] });
    }
  }

  const siloedValues = settledRequests.map(({ readRequest }) => readRequest.value);

  // Fetch all membership witnesses in parallel
  const membershipWitnesses = await Promise.all(siloedValues.map(value => oracle.getNullifierMembershipWitness(value)));

  // Add settled read requests to builder
  for (let i = 0; i < settledRequests.length; i++) {
    builder.addSettledReadRequest(
      settledRequests[i].index,
      membershipWitnesses[i].membershipWitness,
      membershipWitnesses[i].leafPreimage,
    );
  }

  return builder.toHints();
}

export async function buildNullifierReadRequestHints<PENDING extends number, SETTLED extends number>(
  oracle: {
    getNullifierMembershipWitness(nullifier: Fr): Promise<NullifierMembershipWitnessWithPreimage>;
  },
  nullifierReadRequests: ClaimedLengthArray<ScopedReadRequest, typeof MAX_NULLIFIER_READ_REQUESTS_PER_TX>,
  nullifiers: ClaimedLengthArray<ScopedNullifier, typeof MAX_NULLIFIERS_PER_TX>,
  maxPending: PENDING = MAX_NULLIFIER_READ_REQUESTS_PER_TX as PENDING,
  maxSettled: SETTLED = MAX_NULLIFIER_READ_REQUESTS_PER_TX as SETTLED,
) {
  const resetActions = getNullifierReadRequestResetActions(nullifierReadRequests, nullifiers);
  return await buildNullifierReadRequestHintsFromResetActions(
    oracle,
    nullifierReadRequests,
    resetActions,
    maxPending,
    maxSettled,
  );
}