import { padArrayEnd } from '@aztec/foundation/collection';
import type { Tuple } from '@aztec/foundation/serialize';

import type { ClaimedLengthArray } from '../claimed_length_array.js';
import type { ScopedNoteHash } from '../note_hash.js';
import type { ScopedNullifier } from '../nullifier.js';
import type { PrivateLogData, ScopedPrivateLogData } from '../private_log_data.js';
import { isValidNoteHashReadRequest } from './build_note_hash_read_request_hints.js';
import { isValidNullifierReadRequest } from './build_nullifier_read_request_hints.js';
import type { ScopedReadRequest } from './read_request.js';
import { ScopedValueCache } from './scoped_value_cache.js';
import { TransientDataSquashingHint } from './transient_data_squashing_hint.js';

export function buildTransientDataHints<NOTE_HASHES_LEN extends number, NULLIFIERS_LEN extends number>(
  noteHashes: ClaimedLengthArray<ScopedNoteHash, NOTE_HASHES_LEN>,
  nullifiers: ClaimedLengthArray<ScopedNullifier, NULLIFIERS_LEN>,
  futureNoteHashReads: ScopedReadRequest[],
  futureNullifierReads: ScopedReadRequest[],
  futureLogs: PrivateLogData[],
  noteHashNullifierCounterMap: Map<number, number>,
  splitCounter: number,
): { numTransientData: number; hints: Tuple<TransientDataSquashingHint, NULLIFIERS_LEN> } {
  const futureNoteHashReadsMap = new ScopedValueCache(futureNoteHashReads);
  const futureNullifierReadsMap = new ScopedValueCache(futureNullifierReads);
  const futureLogNoteHashCounters = new Set(futureLogs.filter(l => l.noteHashCounter > 0).map(l => l.noteHashCounter));

  const nullifierIndexMap: Map<number, number> = new Map();
  nullifiers.getActiveItems().forEach((n, i) => nullifierIndexMap.set(n.counter, i));

  const hints = [];
  for (let noteHashIndex = 0; noteHashIndex < noteHashes.claimedLength; noteHashIndex++) {
    const noteHash = noteHashes.array[noteHashIndex];
    const noteHashNullifierCounter = noteHashNullifierCounterMap.get(noteHash.counter);
    // The note hash might not be linked to a nullifier, or it might be read in the future, or a future log might be
    // linked to it.
    if (
      !noteHashNullifierCounter ||
      futureNoteHashReadsMap.get(noteHash).find(read => isValidNoteHashReadRequest(read, noteHash)) ||
      futureLogNoteHashCounters.has(noteHash.counter)
    ) {
      continue;
    }

    const nullifierIndex = nullifierIndexMap.get(noteHashNullifierCounter);
    // We might not have the corresponding nullifier yet
    if (nullifierIndex === undefined) {
      continue;
    }

    const nullifier = nullifiers.array[nullifierIndex];
    // Cannot nullify a non-revertible note hash with a revertible nullifier.
    if (noteHash.counter < splitCounter && nullifier.counter >= splitCounter) {
      continue;
    }

    // If the following errors show up, something's wrong with how we build the noteHashNullifierCounterMap in client_execution_context.ts.
    if (nullifier.counter < noteHash.counter) {
      throw new Error('Hinted nullifier has smaller counter than note hash.');
    }
    if (!nullifier.contractAddress.equals(noteHash.contractAddress)) {
      throw new Error('Contract address of hinted note hash does not match.');
    }

    if (!nullifier.nullifiedNoteHash.equals(noteHash.value)) {
      // If the hinted note hash has a different value, it means the nullifier is nullifying a siloed note hash.
      // We don't squash them and both values will be emitted.
      // The private kernel tail will check that the nullified note hash matches a siloed note hash in the same tx.
      continue;
    }

    // The nullifier might be read in the future
    if (futureNullifierReadsMap.get(nullifier).find(read => isValidNullifierReadRequest(read, nullifier))) {
      continue;
    }

    hints.push(new TransientDataSquashingHint(nullifierIndex, noteHashIndex));
  }

  const noActionHint = new TransientDataSquashingHint(nullifiers.array.length, noteHashes.array.length);
  return {
    numTransientData: hints.length,
    hints: padArrayEnd(hints, noActionHint, nullifiers.array.length as NULLIFIERS_LEN),
  };
}

/** Counts private logs that are linked to squashed note hashes and would be removed along with them. */
export function countSquashedLogs<NOTE_HASHES_LEN extends number, LOGS_LEN extends number>(
  noteHashes: ClaimedLengthArray<ScopedNoteHash, NOTE_HASHES_LEN>,
  privateLogs: ClaimedLengthArray<ScopedPrivateLogData, LOGS_LEN>,
  squashingHints: TransientDataSquashingHint[],
): number {
  const squashedNoteHashCounters = new Set(squashingHints.map(h => noteHashes.array[h.noteHashIndex].counter));
  return privateLogs.getActiveItems().filter(l => squashedNoteHashCounters.has(l.inner.noteHashCounter)).length;
}