import initACVM, { type ExecutionError, type ForeignCallHandler, executeCircuit } from '@aztec-foundation/noir-acvm_js';
import initAbi, { abiDecode, abiEncode } from '@aztec-foundation/noir-noirc_abi';
import type { InputMap } from '@aztec-foundation/noir-types';

import { type Logger, type LoggerBindings, resolveLogger } from '@aztec-labs/foundation/log';
import { Timer } from '@aztec-labs/foundation/timer';
import type { FunctionArtifactWithContractName } from '@aztec-labs/stdlib/abi';
import type { NoirCompiledCircuitWithName } from '@aztec-labs/stdlib/noir';
import { Buffer } from 'buffer';

import { type ACIRCallback, type ACIRExecutionResult, acvm } from './acvm/acvm.js';
import type { ACVMWitness } from './acvm/acvm_types.js';
import type { ACVMSuccess } from './acvm_native.js';
import { type CircuitSimulator, type ProtocolCircuitResult, enrichNoirError } from './circuit_simulator.js';

let wasmInitPromise: Promise<unknown> | undefined;

export class WASMSimulator implements CircuitSimulator {
  protected log: Logger;

  constructor(loggerOrBindings?: Logger | LoggerBindings) {
    this.log = resolveLogger('wasm-simulator', loggerOrBindings);
  }

  async init(): Promise<void> {
    // If these are available, then we are in the
    // web environment. For the node environment, this
    // is a no-op.
    if (typeof initAbi === 'function') {
      /** @ts-expect-error The node bundle doesn't include these default imports, so TS complains */
      wasmInitPromise ??= Promise.all([initAbi(), initACVM()]);
      await wasmInitPromise;
    }
  }

  async executeProtocolCircuit<ReturnType>(
    inputs: InputMap,
    artifact: NoirCompiledCircuitWithName,
    callback: ForeignCallHandler,
  ): Promise<ProtocolCircuitResult<ReturnType>> {
    await this.init();
    // `executeCircuit` only speaks witness maps, so the ABI encode/decode happens here rather than
    // in the caller. See the note on CircuitSimulator.executeProtocolCircuit.
    const { witness, duration } = await this.executeProtocolCircuitToWitness(
      abiEncode(artifact.abi, inputs),
      artifact,
      callback,
    );
    return { returnValue: abiDecode(artifact.abi, witness).return_value as ReturnType, duration };
  }

  async executeProtocolCircuitToWitness(
    input: ACVMWitness,
    artifact: NoirCompiledCircuitWithName,
    callback: ForeignCallHandler,
  ): Promise<ACVMSuccess> {
    this.log.debug('init', { hash: artifact.hash });
    await this.init();

    // Decode the bytecode from base64 since the acvm does not know about base64 encoding
    const decodedBytecode = Buffer.from(artifact.bytecode, 'base64');
    //
    // Execute the circuit
    try {
      const timer = new Timer();
      const result = await executeCircuit(
        decodedBytecode,
        input,
        callback, // handle calls to debug_log
      );
      this.log.debug('execution successful', { hash: artifact.hash });
      return { witness: result, duration: timer.ms() } as ACVMSuccess;
    } catch (err) {
      // Typescript types caught errors as unknown or any, so we need to narrow its type to check if it has raw
      // assertion payload.
      if (typeof err === 'object' && err !== null && 'rawAssertionPayload' in err) {
        const parsed = enrichNoirError(artifact, err as ExecutionError);
        this.log.debug('execution failed', {
          hash: artifact.hash,
          error: parsed,
          message: parsed.message,
        });
        throw parsed;
      }
      this.log.debug('execution failed', { hash: artifact.hash, error: err });
      throw new Error(`Circuit execution failed: ${err}`);
    }
  }

  async executeUserCircuit(
    input: ACVMWitness,
    artifact: FunctionArtifactWithContractName,
    callback: ACIRCallback,
  ): Promise<ACIRExecutionResult> {
    await this.init();
    return acvm(artifact.bytecode, input, callback, this.log.createChild('acvm'));
  }
}