import { makeTuple } from '@aztec/foundation/array';
import { Fr } from '@aztec/foundation/curves/bn254';
import { schemas } from '@aztec/foundation/schemas';
import { BufferReader, serializeToBuffer } from '@aztec/foundation/serialize';
import { bufferToHex, hexToBuffer } from '@aztec/foundation/string';

import { Proof, makeEmptyProof } from './proof.js';

/**
 * The Recursive proof class is a wrapper around the circuit's proof.
 * We store the proof in 2 forms for convenience. The first is in the 'fields' format.
 * This is a list of fields, for which there are distinct lengths based on the level of recursion.
 * This 'fields' version does not contain the circuits public inputs
 * We also store the raw binary proof which van be directly verified.
 *
 * The 'fieldsValid' member is set to false in the case where this object is constructed solely from the 'binary' proof
 * This is usually when the proof has been received from clients and signals to provers that the 'fields' version needs to be generated
 */
export class RecursiveProof<N extends number> {
  constructor(
    /**
     * Holds the serialized proof data in an array of fields, this is without the public inputs
     */
    public proof: Fr[],

    /**
     * Holds the serialized proof data in a binary buffer, this contains the public inputs
     */
    public binaryProof: Proof,
    /**
     * This flag determines if the 'proof' member is valid, or if we need to generate it from the 'binaryProof' first
     */
    public fieldsValid: boolean,
    public proofLength: N,
  ) {
    if (proof.length !== proofLength) {
      throw new Error(`Proof length ${proof.length} does not match expected length ${proofLength}.`);
    }
  }

  /**
   * Create a Proof from a Buffer or BufferReader.
   * Expects a length-encoding.
   *
   * @param buffer - A Buffer or BufferReader containing the length-encoded proof data.
   * @returns A Proof instance containing the decoded proof data.
   */
  static fromBuffer<N extends number>(buffer: Buffer | BufferReader, expectedSize?: N): RecursiveProof<N> {
    const reader = BufferReader.asReader(buffer);
    const size = reader.readNumber();
    if (typeof expectedSize === 'number' && expectedSize !== size) {
      throw new Error(`Expected proof length ${expectedSize}, got ${size}`);
    }
    return new RecursiveProof(
      reader.readArray(size, Fr) as any,
      Proof.fromBuffer(reader),
      reader.readBoolean(),
      size as N,
    );
  }

  /**
   * Convert the Proof instance to a custom Buffer format.
   * This function serializes the Proof's buffer length and data sequentially into a new Buffer.
   *
   * @returns A Buffer containing the serialized proof data in custom format.
   */
  public toBuffer() {
    return serializeToBuffer(this.proof.length, this.proof, this.binaryProof, this.fieldsValid);
  }

  /**
   * Serialize the Proof instance to a hex string.
   * @returns The hex string representation of the proof data.
   */
  public toString() {
    return bufferToHex(this.toBuffer());
  }

  /**
   * Deserialize a Proof instance from a hex string.
   * @param str - A hex string to deserialize from.
   * @returns - A new Proof instance.
   */
  static fromString<N extends number>(str: string, expectedSize?: N): RecursiveProof<N> {
    return RecursiveProof.fromBuffer(hexToBuffer(str), expectedSize);
  }

  /** Returns a buffer representation for JSON serialization. */
  toJSON() {
    return this.toBuffer();
  }

  /** Creates an instance from a hex string with expected size. */
  static schemaFor<N extends number>(expectedSize?: N) {
    return schemas.Buffer.transform(b => RecursiveProof.fromBuffer(b, expectedSize));
  }
}

/**
 * Makes an empty proof.
 * Note: Used for local devnet milestone where we are not proving anything yet.
 * @returns The empty "proof".
 */
export function makeEmptyRecursiveProof<N extends number>(size: N) {
  return new RecursiveProof(makeTuple<Fr, N>(size, Fr.zero), makeEmptyProof(), true, size);
}

export function makeRecursiveProof<PROOF_LENGTH extends number>(size: PROOF_LENGTH, seed = 1) {
  return new RecursiveProof<PROOF_LENGTH>(
    makeTuple<Fr, PROOF_LENGTH>(size, (i: number) => new Fr(i), seed),
    makeEmptyProof(),
    true,
    size,
  );
}

/**
 * Makes an instance of the recursive proof from a binary only proof
 * @returns The proof object
 */
export function makeRecursiveProofFromBinary<PROOF_LENGTH extends number>(proof: Proof, size: PROOF_LENGTH) {
  return new RecursiveProof<PROOF_LENGTH>(makeTuple<Fr, PROOF_LENGTH>(size, Fr.zero), proof, false, size);
}