/**
 * UDS server for AVM CDB requests.
 *
 * Transport (socket, framing, per-connection response ordering) is handled by the shared
 * `UdsIpcServer` from `@aztec-foundation/ipc-runtime`; message dispatch comes from the generated CDB
 * server. This class only implements the generated handler interface and routes requests to
 * PublicContractsDB instances by fork ID.
 */
import {
  type CdbAddContracts,
  type CdbAddContractsResponse,
  type CdbCommitCheckpoint,
  type CdbCommitCheckpointResponse,
  type CdbCreateCheckpoint,
  type CdbCreateCheckpointResponse,
  type CdbGetBytecodeCommitment,
  type CdbGetBytecodeCommitmentResponse,
  type CdbGetContractClass,
  type CdbGetContractClassResponse,
  type CdbGetContractInstance,
  type CdbGetContractInstanceResponse,
  type CdbGetDebugFunctionName,
  type CdbGetDebugFunctionNameResponse,
  type Handler as CdbHandler,
  type CdbRevertCheckpoint,
  type CdbRevertCheckpointResponse,
  handleRequest,
} from '@aztec-foundation/cdb';
import { UdsIpcServer } from '@aztec-foundation/ipc-runtime';

import { Fr } from '@aztec-labs/foundation/curves/bn254';
import { type Logger, createLogger } from '@aztec-labs/foundation/log';
import { FunctionSelector } from '@aztec-labs/stdlib/abi';
import { AztecAddress } from '@aztec-labs/stdlib/aztec-address';
import { ContractDeploymentData, type ContractInstanceWithAddress } from '@aztec-labs/stdlib/contract';
import { Buffer } from 'buffer';
import { Decoder, Encoder } from 'msgpackr';
import * as os from 'node:os';
import * as path from 'node:path';
import { threadId } from 'node:worker_threads';

import type { PublicContractsDB } from './public_db_sources.js';

const encoder = new Encoder({ useRecords: false });
const decoder = new Decoder({ useRecords: false });

let instanceCounter = 0;

function toFieldBuffer(field: Fr | AztecAddress): Buffer {
  return field.toBuffer();
}

/** Serializes contract instances to the AVM CDB wire shape. */
export function serializeContractInstance(instance: ContractInstanceWithAddress): Record<string, unknown> {
  return {
    salt: toFieldBuffer(instance.salt),
    deployer: toFieldBuffer(instance.deployer),
    currentContractClassId: toFieldBuffer(instance.currentContractClassId),
    originalContractClassId: toFieldBuffer(instance.originalContractClassId),
    initializationHash: toFieldBuffer(instance.initializationHash),
    immutablesHash: toFieldBuffer(instance.immutablesHash),
    publicKeys: {
      npkMHash: toFieldBuffer(instance.publicKeys.npkMHash),
      ivpkM: {
        x: toFieldBuffer(instance.publicKeys.ivpkM.x),
        y: toFieldBuffer(instance.publicKeys.ivpkM.y),
      },
      ovpkMHash: toFieldBuffer(instance.publicKeys.ovpkMHash),
      tpkMHash: toFieldBuffer(instance.publicKeys.tpkMHash),
      mspkMHash: toFieldBuffer(instance.publicKeys.mspkMHash),
      fbpkMHash: toFieldBuffer(instance.publicKeys.fbpkMHash),
    },
  };
}

/** Serializes contract classes to the AVM CDB wire shape. */
function serializeContractClass(contractClass: {
  id: Fr;
  artifactHash: Fr;
  privateFunctionsRoot: Fr;
  packedBytecode: Buffer;
}): Record<string, unknown> {
  return {
    id: toFieldBuffer(contractClass.id),
    artifactHash: toFieldBuffer(contractClass.artifactHash),
    privateFunctionsRoot: toFieldBuffer(contractClass.privateFunctionsRoot),
    packedBytecode: contractClass.packedBytecode,
  };
}

/** Routes AVM CDB IPC requests to registered PublicContractsDB forks. */
export class CdbIpcServer implements CdbHandler {
  public readonly ipcPath: string;
  private server: Promise<UdsIpcServer>;
  private log: Logger;
  /** Maps WSDB fork IDs to contracts DB/timestamp pairs for concurrent simulations. */
  private forks = new Map<number, { db: PublicContractsDB; timestamp: bigint }>();

  constructor() {
    this.log = createLogger('cdb-ipc-server');
    this.ipcPath = path.join(os.tmpdir(), `cdb-ts-${process.pid}-${threadId}-${instanceCounter++}.sock`);

    // Listen asynchronously; the C++ AVM retries connecting until the socket is up.
    this.server = UdsIpcServer.listen(this.ipcPath, (_clientId, request) => handleRequest(this, request));
    this.server.then(
      () => this.log.debug(`CDB IPC server listening on ${this.ipcPath}`),
      err => this.log.error(`CDB IPC server failed to listen on ${this.ipcPath}`, { err }),
    );
  }

  /** Register a PublicContractsDB for a given WSDB fork ID. */
  registerFork(forkId: number, contractsDB: PublicContractsDB, timestamp: bigint): void {
    this.forks.set(forkId, { db: contractsDB, timestamp });
  }

  /** Unregister a fork's contracts DB (call after simulation completes). */
  unregisterFork(forkId: number): void {
    this.forks.delete(forkId);
  }

  /** Close the server and all active connections. */
  async close(): Promise<void> {
    const server = await this.server.catch(() => undefined);
    await server?.close();
  }

  async [Symbol.asyncDispose](): Promise<void> {
    await this.close();
  }

  /** Look up the contracts DB for a given fork ID, throwing if not registered. */
  private getFork(forkId: number): { db: PublicContractsDB; timestamp: bigint } {
    const fork = this.forks.get(forkId);
    if (!fork) {
      const registered = Array.from(this.forks.keys()).sort((a, b) => a - b);
      throw new Error(
        `CDB server: no contracts DB registered for forkId ${forkId} (registered=[${registered.join(',')}])`,
      );
    }
    return fork;
  }

  async getContractInstance(command: CdbGetContractInstance): Promise<CdbGetContractInstanceResponse> {
    const { db, timestamp } = this.getFork(command.forkId);
    const address = AztecAddress.fromBuffer(Buffer.from(command.address));
    const instance = await db.getContractInstance(address, timestamp);
    return { instance: instance ? encoder.encode(serializeContractInstance(instance)) : null };
  }

  async getContractClass(command: CdbGetContractClass): Promise<CdbGetContractClassResponse> {
    const { db } = this.getFork(command.forkId);
    const classId = Fr.fromBuffer(Buffer.from(command.classId));
    const contractClass = await db.getContractClass(classId);
    return { contractClass: contractClass ? encoder.encode(serializeContractClass(contractClass)) : null };
  }

  async getBytecodeCommitment(command: CdbGetBytecodeCommitment): Promise<CdbGetBytecodeCommitmentResponse> {
    const { db } = this.getFork(command.forkId);
    const classId = Fr.fromBuffer(Buffer.from(command.classId));
    const commitment = await db.getBytecodeCommitment(classId);
    return { commitment: commitment ? toFieldBuffer(commitment) : null };
  }

  async getDebugFunctionName(command: CdbGetDebugFunctionName): Promise<CdbGetDebugFunctionNameResponse> {
    const { db } = this.getFork(command.forkId);
    const address = AztecAddress.fromBuffer(Buffer.from(command.address));
    const selectorField = Fr.fromBuffer(Buffer.from(command.selector));
    const selector = FunctionSelector.fromFieldOrUndefined(selectorField);
    const name = selector ? await db.getDebugFunctionName(address, selector) : undefined;
    return { name: name ?? null };
  }

  addContracts(command: CdbAddContracts): Promise<CdbAddContractsResponse> {
    const { db } = this.getFork(command.forkId);
    const contractDeploymentData = ContractDeploymentData.fromPlainObject(
      decoder.decode(command.contractDeploymentData),
    );
    db.addContractsFromLogs(contractDeploymentData);
    return Promise.resolve({});
  }

  createCheckpoint(command: CdbCreateCheckpoint): Promise<CdbCreateCheckpointResponse> {
    const { db } = this.getFork(command.forkId);
    db.createCheckpoint();
    return Promise.resolve({});
  }

  commitCheckpoint(command: CdbCommitCheckpoint): Promise<CdbCommitCheckpointResponse> {
    const { db } = this.getFork(command.forkId);
    db.commitCheckpoint();
    return Promise.resolve({});
  }

  revertCheckpoint(command: CdbRevertCheckpoint): Promise<CdbRevertCheckpointResponse> {
    const { db } = this.getFork(command.forkId);
    db.revertCheckpoint();
    return Promise.resolve({});
  }
}