import { Fr } from '@aztec-labs/foundation/curves/bn254';
import {
  AvmCircuitPublicInputs,
  type AvmTxHint,
  PublicTxEffect,
  deserializeFromMessagePack,
  serializeWithMessagePack,
} from '@aztec-labs/stdlib/avm';
import { GlobalVariables, TreeSnapshots } from '@aztec-labs/stdlib/tx';
import { NativeWorldStateService } from '@aztec-labs/world-state';
import { Buffer } from 'buffer';
import { createInterface } from 'readline';

import { AvmFuzzerSimulator, FuzzerSimulationRequest } from './avm_fuzzer_simulator.js';

/** Write data to stdout, letting Node handle buffering. */
function writeOutput(data: string): Promise<void> {
  return new Promise<void>((resolve, reject) => {
    process.stdout.write(data, err => {
      if (err) {
        reject(err);
      } else {
        resolve();
      }
    });
  });
}

// The fuzzer self-bootstraps its own world state instead of reading a shared on-disk database written by
// the C++ side. A fresh NativeWorldStateService produces a genesis identical to the C++ MemoryMerkleDB by
// construction (same 128 nullifier/public-data tree prefill and header-generator point), and the dynamic
// state is applied per-input below. The instance is opened once and reused across inputs; each input forks
// from it (in AvmFuzzerSimulator.create).
let worldStatePromise: Promise<NativeWorldStateService> | undefined;

function getWorldState(): Promise<NativeWorldStateService> {
  return (worldStatePromise ??= NativeWorldStateService.tmp());
}

async function simulateWithFuzzer(
  txHint: AvmTxHint,
  globals: GlobalVariables,
  rawContractClasses: any[], // Replace these when we are moving contract classes to TS
  rawContractInstances: [any, any][], // Replace these when we are moving contract instances to TS
  rawPublicDataWrites: any[], // Public data tree writes to apply before simulation
  rawNoteHashes: any[], // Note hashes to apply before simulation
): Promise<{
  reverted: boolean;
  output: Fr[];
  revertReason?: string;
  publicInputs: AvmCircuitPublicInputs;
  publicTxEffect: PublicTxEffect;
}> {
  const worldStateService = await getWorldState();

  const simulator = await AvmFuzzerSimulator.create(worldStateService, globals);

  await simulator.applyNoteHashes(rawNoteHashes);

  // Register contract classes from C++ (must happen before public data writes to match C++ order)
  for (const rawClass of rawContractClasses) {
    await simulator.addContractClassFromCpp(rawClass);
  }

  // Register contract instances from C++
  for (const [rawAddress, rawInstance] of rawContractInstances) {
    await simulator.addContractInstanceFromCpp(rawAddress, rawInstance);
  }

  // Apply public data writes after contract registration (e.g., for bytecode upgrades)
  // This must happen last to match C++ setup_fuzzer_state ordering
  await simulator.applyPublicDataWrites(rawPublicDataWrites);

  const result = await simulator.simulate(txHint);

  const output = result
    .getAppLogicReturnValues()
    .flatMap((rv: { values?: Fr[] } | undefined) => rv?.values?.filter((v: Fr | null | undefined) => v != null) ?? []);

  return {
    reverted: !result.revertCode.isOK(),
    output,
    revertReason: result.findRevertReason()?.message,
    publicInputs: result.publicInputs!,
    publicTxEffect: result.publicTxEffect,
  };
}

async function execute(base64Line: string): Promise<void> {
  try {
    // Decode base64 and deserialize the entire request from msgpack
    const buffer = Buffer.from(base64Line.trim(), 'base64');
    const rawRequest = deserializeFromMessagePack(buffer);
    const request = FuzzerSimulationRequest.fromPlainObject(rawRequest);

    // Run the simulation on the C++ AVM simulator
    const result = await simulateWithFuzzer(
      request.tx,
      request.globals,
      request.contractClasses,
      request.contractInstances,
      request.publicDataWrites,
      request.noteHashes,
    );

    // Serialize the result to msgpack and encode it in base64 for output
    const resultBuffer = serializeWithMessagePack({
      reverted: result.reverted,
      output: result.output,
      revertReason: result.revertReason ?? '',
      endTreeSnapshots: result.publicInputs.endTreeSnapshots,
      publicTxEffect: result.publicTxEffect,
    });
    const base64Response = resultBuffer.toString('base64') + '\n';
    await writeOutput(base64Response);
  } catch (error: any) {
    // If we error, treat as reverted
    const errorResult = serializeWithMessagePack({
      reverted: true,
      output: [] as Fr[],
      revertReason: `Unexpected Error ${error.message}`,
      endTreeSnapshots: TreeSnapshots.empty(),
      publicTxEffect: PublicTxEffect.empty(),
    });
    await writeOutput(errorResult.toString('base64') + '\n');
  }
}

function mainLoop() {
  const rl = createInterface({ input: process.stdin, terminal: false });

  // Process lines sequentially to avoid race conditions in responses
  const lineQueue: string[] = [];
  let processing = false;

  async function processQueue() {
    if (processing || lineQueue.length === 0) {
      return;
    }
    processing = true;
    while (lineQueue.length > 0) {
      const line = lineQueue.shift()!;
      await execute(line);
    }
    processing = false;
  }

  rl.on('line', (line: string) => {
    if (line.trim()) {
      lineQueue.push(line);
      void processQueue();
    }
  });
  rl.on('close', () => process.exit(0));
}

void mainLoop();