import {
  type MAX_NOTE_HASHES_PER_TX,
  MAX_NOTE_HASH_READ_REQUESTS_PER_TX,
  type NOTE_HASH_TREE_HEIGHT,
} from '@aztec/constants';
import type { Fr } from '@aztec/foundation/curves/bn254';
import type { MembershipWitness } from '@aztec/foundation/trees';

import type { ClaimedLengthArray } from '../claimed_length_array.js';
import type { ScopedNoteHash } from '../note_hash.js';
import { NoteHashReadRequestHintsBuilder } from './note_hash_read_request_hints.js';
import type { ScopedReadRequest } from './read_request.js';
import { PendingReadHint, ReadRequestActionEnum, ReadRequestResetActions } from './read_request_hints.js';

export function isValidNoteHashReadRequest(readRequest: ScopedReadRequest, noteHash: ScopedNoteHash) {
  return (
    noteHash.value.equals(readRequest.value) &&
    noteHash.contractAddress.equals(readRequest.contractAddress) &&
    readRequest.counter > noteHash.counter
  );
}

export function getNoteHashReadRequestResetActions(
  noteHashReadRequests: ClaimedLengthArray<ScopedReadRequest, typeof MAX_NOTE_HASH_READ_REQUESTS_PER_TX>,
  noteHashes: ClaimedLengthArray<ScopedNoteHash, typeof MAX_NOTE_HASHES_PER_TX>,
): ReadRequestResetActions<typeof MAX_NOTE_HASH_READ_REQUESTS_PER_TX> {
  const resetActions = ReadRequestResetActions.empty(MAX_NOTE_HASH_READ_REQUESTS_PER_TX);

  const noteHashMap: Map<bigint, { noteHash: ScopedNoteHash; index: number }[]> = new Map();
  noteHashes.getActiveItems().forEach((noteHash, index) => {
    const value = noteHash.value.toBigInt();
    const arr = noteHashMap.get(value) ?? [];
    arr.push({ noteHash, index });
    noteHashMap.set(value, arr);
  });

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

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

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

  return resetActions;
}

export async function buildNoteHashReadRequestHintsFromResetActions<PENDING extends number, SETTLED extends number>(
  oracle: {
    getNoteHashMembershipWitness(noteHash: Fr): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined>;
  },
  noteHashReadRequests: ClaimedLengthArray<ScopedReadRequest, typeof MAX_NOTE_HASH_READ_REQUESTS_PER_TX>,
  noteHashes: ClaimedLengthArray<ScopedNoteHash, typeof MAX_NOTE_HASHES_PER_TX>,
  resetActions: ReadRequestResetActions<typeof MAX_NOTE_HASH_READ_REQUESTS_PER_TX>,
  maxPending: PENDING = MAX_NOTE_HASH_READ_REQUESTS_PER_TX as PENDING,
  maxSettled: SETTLED = MAX_NOTE_HASH_READ_REQUESTS_PER_TX as SETTLED,
) {
  const builder = new NoteHashReadRequestHintsBuilder(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: noteHashReadRequests.array[i] });
    }
  }

  // Fetch all membership witnesses in parallel
  const membershipWitnesses = await Promise.all(
    settledRequests.map(({ readRequest }) => oracle.getNoteHashMembershipWitness(readRequest.value)),
  );

  // Add settled read requests to builder
  for (let i = 0; i < settledRequests.length; i++) {
    const membershipWitness = membershipWitnesses[i];
    if (!membershipWitness) {
      throw new Error('Read request is reading an unknown note hash.');
    }
    builder.addSettledReadRequest(settledRequests[i].index, membershipWitness, settledRequests[i].readRequest.value);
  }

  const noteHashMap: Map<bigint, { noteHash: ScopedNoteHash; index: number }[]> = new Map();
  noteHashes.getActiveItems().forEach((noteHash, index) => {
    const value = noteHash.value.toBigInt();
    const arr = noteHashMap.get(value) ?? [];
    arr.push({ noteHash, index });
    noteHashMap.set(value, arr);
  });

  return builder.toHints();
}

export async function buildNoteHashReadRequestHints<PENDING extends number, SETTLED extends number>(
  oracle: {
    getNoteHashMembershipWitness(noteHash: Fr): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>>;
  },
  noteHashReadRequests: ClaimedLengthArray<ScopedReadRequest, typeof MAX_NOTE_HASH_READ_REQUESTS_PER_TX>,
  noteHashes: ClaimedLengthArray<ScopedNoteHash, typeof MAX_NOTE_HASHES_PER_TX>,
  maxPending: PENDING = MAX_NOTE_HASH_READ_REQUESTS_PER_TX as PENDING,
  maxSettled: SETTLED = MAX_NOTE_HASH_READ_REQUESTS_PER_TX as SETTLED,
) {
  const resetActions = getNoteHashReadRequestResetActions(noteHashReadRequests, noteHashes);
  return await buildNoteHashReadRequestHintsFromResetActions(
    oracle,
    noteHashReadRequests,
    noteHashes,
    resetActions,
    maxPending,
    maxSettled,
  );
}