/* eslint-disable camelcase */
import type { Fr } from '@aztec/foundation/curves/bn254';
import { createLogger } from '@aztec/foundation/log';
import { schemas, zodFor } from '@aztec/foundation/schemas';

import { inflate } from 'pako';
import { z } from 'zod';

import { DEV_VERSION } from '../update-checker/dev_version.js';
import { FunctionSelector } from './function_selector.js';

/** A basic value. */
export interface BasicValue<T extends string, V> {
  /** The kind of the value. */
  kind: T;
  value: V;
}

const logger = createLogger('aztec:foundation:abi');

/** An exported value. */
export type AbiValue =
  | BasicValue<'boolean', boolean>
  | BasicValue<'string', string>
  | BasicValue<'array', AbiValue[]>
  | TupleValue
  | IntegerValue
  | StructValue;

export const AbiValueSchema: z.ZodType<AbiValue> = z.discriminatedUnion('kind', [
  z.object({ kind: z.literal('boolean'), value: z.boolean() }),
  z.object({ kind: z.literal('string'), value: z.string() }),
  z.object({ kind: z.literal('array'), value: z.array(z.lazy(() => AbiValueSchema)) }),
  z.object({ kind: z.literal('tuple'), fields: z.array(z.lazy(() => AbiValueSchema)) }),
  z.object({ kind: z.literal('integer'), value: z.string(), sign: z.boolean() }),
  z.object({
    kind: z.literal('struct'),
    fields: z.array(z.object({ name: z.string(), value: z.lazy(() => AbiValueSchema) })),
  }),
]);

export type TypedStructFieldValue<T> = { name: string; value: T };

export interface StructValue {
  kind: 'struct';
  fields: TypedStructFieldValue<AbiValue>[];
}

export interface TupleValue {
  kind: 'tuple';
  fields: AbiValue[];
}

export interface IntegerValue extends BasicValue<'integer', string> {
  sign: boolean;
}

/** Indicates whether a parameter is public or secret/private. */
export const ABIParameterVisibility = ['public', 'private', 'databus'] as const;

/** Indicates whether a parameter is public or secret/private. */
export type ABIParameterVisibility = (typeof ABIParameterVisibility)[number];

/** A basic type. */
export interface BasicType<T extends string> {
  /** The kind of the type. */
  kind: T;
}

/** Sign for numeric types. */
const Sign = ['unsigned', 'signed'] as const;
type Sign = (typeof Sign)[number];

/** A variable type. */
export type AbiType =
  | BasicType<'field'>
  | BasicType<'boolean'>
  | IntegerType
  | ArrayType
  | StringType
  | StructType
  | TupleType;

export const AbiTypeSchema: z.ZodType<AbiType> = z.discriminatedUnion('kind', [
  z.object({ kind: z.literal('field') }),
  z.object({ kind: z.literal('boolean') }),
  z.object({ kind: z.literal('integer'), sign: z.enum(Sign), width: z.number() }),
  z.object({ kind: z.literal('array'), length: z.number(), type: z.lazy(() => AbiTypeSchema) }),
  z.object({ kind: z.literal('string'), length: z.number() }),
  z.object({ kind: z.literal('struct'), fields: z.array(z.lazy(() => ABIVariableSchema)), path: z.string() }),
  z.object({ kind: z.literal('tuple'), fields: z.array(z.lazy(() => AbiTypeSchema)) }),
]);

/** A named type. */
export const ABIVariableSchema = z.object({
  /** The name of the variable. */
  name: z.string(),
  /** The type of the variable. */
  type: AbiTypeSchema,
});

/** A named type. */
export type ABIVariable = z.infer<typeof ABIVariableSchema>;

/** A function parameter. */
export const ABIParameterSchema = ABIVariableSchema.and(
  z.object({
    /** Visibility of the parameter in the function. */
    visibility: z.enum(ABIParameterVisibility),
  }),
);

/** A function parameter. */
export type ABIParameter = z.infer<typeof ABIParameterSchema>;

/** An integer type. */
export interface IntegerType extends BasicType<'integer'> {
  /** The sign of the integer. */
  sign: Sign;
  /** The width of the integer in bits. */
  width: number;
}

/** An array type. */
export interface ArrayType extends BasicType<'array'> {
  /** The length of the array. */
  length: number;
  /** The type of the array elements. */
  type: AbiType;
}

/** A tuple type. */
export interface TupleType extends BasicType<'tuple'> {
  /** The types of the tuple elements. */
  fields: AbiType[];
}

/** A string type. */
export interface StringType extends BasicType<'string'> {
  /** The length of the string. */
  length: number;
}

/** A struct type. */
export interface StructType extends BasicType<'struct'> {
  /** The fields of the struct. */
  fields: ABIVariable[];
  /** Fully qualified name of the struct. */
  path: string;
}

/** An error could be a custom error of any regular type or a string error. */
export type AbiErrorType =
  | { error_kind: 'string'; string: string }
  | { error_kind: 'fmtstring'; length: number; item_types: AbiType[] }
  | ({ error_kind: 'custom' } & AbiType);

const AbiErrorTypeSchema = zodFor<AbiErrorType>()(
  z.union([
    z.object({ error_kind: z.literal('string'), string: z.string() }),
    z.object({ error_kind: z.literal('fmtstring'), length: z.number(), item_types: z.array(AbiTypeSchema) }),
    z.object({ error_kind: z.literal('custom') }).and(AbiTypeSchema),
  ]),
);

/** Aztec.nr function types. */
export enum FunctionType {
  PRIVATE = 'private',
  PUBLIC = 'public',
  UTILITY = 'utility',
}

/** The abi entry of a function. */
export interface FunctionAbi {
  /** The name of the function. */
  name: string;
  /** Whether the function is secret. */
  functionType: FunctionType;
  /** Whether the function is marked as `#[only_self]` and hence callable only from within the contract. */
  isOnlySelf: boolean;
  /** Whether the function can alter state or not */
  isStatic: boolean;
  /** Function parameters. */
  parameters: ABIParameter[];
  /** The types of the return values. */
  returnTypes: AbiType[];
  /** The types of the errors that the function can throw. */
  errorTypes: Partial<Record<string, AbiErrorType>>;
  /** Whether the function is flagged as an initializer. */
  isInitializer: boolean;
}

export const FunctionAbiSchema = z.object({
  name: z.string(),
  functionType: z.nativeEnum(FunctionType),
  isOnlySelf: z.boolean(),
  isStatic: z.boolean(),
  isInitializer: z.boolean(),
  parameters: z.array(z.object({ name: z.string(), type: AbiTypeSchema, visibility: z.enum(ABIParameterVisibility) })),
  returnTypes: z.array(AbiTypeSchema),
  errorTypes: z.record(z.string(), AbiErrorTypeSchema.optional()),
}) satisfies z.ZodType<FunctionAbi>;

/** Debug metadata for a function. */
export interface FunctionDebugMetadata {
  /** Maps opcodes to source code pointers */
  debugSymbols: DebugInfo;
  /** Maps the file IDs to the file contents to resolve pointers */
  files: DebugFileMap;
}

export const FunctionDebugMetadataSchema = z.object({
  debugSymbols: z.object({
    location_tree: z.object({
      locations: z.array(
        z.object({
          parent: z.number().nullable(),
          value: z.object({
            span: z.object({ start: z.number(), end: z.number() }),
            file: z.number(),
          }),
        }),
      ),
    }),
    acir_locations: z.record(z.number()),
    brillig_locations: z.record(z.record(z.number())),
  }),
  // `function_locations` is required in the static `DebugFileMap` type (noir-types compat), but legacy v4-next
  // artifacts were compiled before this field existed — see {@link fillMissingFunctionLocations} for the full
  // rationale. The preprocessor injects `[]` for missing entries so parsing succeeds; the cast restores the
  // output type since `z.preprocess` widens the schema's input type to `unknown`.
  files: z.preprocess(
    fillMissingFunctionLocations,
    z.record(
      z.object({
        source: z.string(),
        path: z.string(),
        function_locations: z.array(z.object({ start: z.number(), name: z.string() })),
      }),
    ),
  ) as z.ZodType<DebugFileMap>,
}) satisfies z.ZodType<FunctionDebugMetadata>;

/** The artifact entry of a function. */
export interface FunctionArtifact extends FunctionAbi {
  /** The ACIR bytecode of the function. */
  bytecode: Buffer;
  /** The verification key of the function, base64 encoded, if it's a private fn. */
  verificationKey?: string;
  /** Maps opcodes to source code pointers */
  debugSymbols: string;
  /** Debug metadata for the function. */
  debug?: FunctionDebugMetadata;
}

export interface FunctionArtifactWithContractName extends FunctionArtifact {
  /** The name of the contract. */
  contractName: string;
}

export const FunctionArtifactSchema = zodFor<FunctionArtifact>()(
  FunctionAbiSchema.and(
    z.object({
      bytecode: schemas.Buffer,
      verificationKey: z.string().optional(),
      debugSymbols: z.string(),
      debug: FunctionDebugMetadataSchema.optional(),
    }),
  ),
);

/** A file ID. It's assigned during compilation. */
type FileId = number;

/** A pointer to a specific section of the source code. */
interface SourceCodeLocation {
  /** The section of the source code. */
  span: {
    /** The byte where the section starts. */
    start: number;
    /** The byte where the section ends. */
    end: number;
  };
  /** The source code file pointed to. */
  file: FileId;
}

/**
 * The location of an opcode in the bytecode.
 * It's a string of the form `{acirIndex}` or `{acirIndex}:{brilligIndex}`.
 */
export type OpcodeLocation = string;

export type BrilligFunctionId = number;

export type OpcodeToLocationsMap = Record<OpcodeLocation, number>;

export type LocationNodeDebugInfo = {
  parent: number | null;
  value: SourceCodeLocation;
};

export type LocationTree = {
  locations: LocationNodeDebugInfo[];
};

/** The debug information for a given function. */
export interface DebugInfo {
  /** A map of the opcode location to the source code location. */
  location_tree: LocationTree;
  acir_locations: OpcodeToLocationsMap;
  /** For each Brillig function, we have a map of the opcode location to the source code location. */
  brillig_locations: Record<BrilligFunctionId, OpcodeToLocationsMap>;
}

/** The debug information for a given program (a collection of functions) */
export interface ProgramDebugInfo {
  /** A list of debug information that matches with each function in a program */
  debug_infos: Array<DebugInfo>;
}

/** The range a function occupies in a file. */
export type FunctionLocation = {
  /** The byte where the function starts. */
  start: number;
  /** The name of the function. */
  name: string;
};

/** Maps a file ID to its metadata for debugging purposes. */
export type DebugFileMap = Record<
  FileId,
  {
    /** The source code of the file. */
    source: string;
    /** The path of the file. */
    path: string;
    /** The range each function occupies in the file. */
    function_locations: FunctionLocation[];
  }
>;

/**
 * Preprocessor for {@link FunctionDebugMetadataSchema} / {@link ContractArtifactSchema} that injects an empty
 * `function_locations` array on any file map entry missing it. Invoked via `z.preprocess(...)` before the strict
 * object schema runs, so legacy inputs do not fail validation.
 *
 * Why this exists on v4-next (and not on `next` / v5):
 *
 *   - The Noir submodule on `next` was bumped past nightly-2026-03-19, which added `function_locations` to the
 *     debug file map in the Noir `@aztec/noir-types` package and to every freshly-compiled artifact's `file_map`
 *     entries. The matching stdlib change (making `DebugFileMap.function_locations` required) was done in one
 *     go on `next`, so there `function_locations` is always present both at compile time and at runtime.
 *
 *   - On v4-next we had to cherry-pick the Noir bump (chore: Update Noir to nightly-2026-04-10 (#22393)) without
 *     re-running the full noir-projects compilation pipeline, so the on-disk contract and protocol-circuit
 *     artifacts (e.g. `yarn-project/protocol-contracts/artifacts/AuthRegistry.json`) still carry the pre-update
 *     `file_map` shape with only `source` + `path`. Separately, `ivc-integration` compares our `DebugFileMap`
 *     against the Noir-types `DebugFileMap`, which *does* require `function_locations`, so we cannot simply make
 *     our own field optional — that would break the structural compatibility check.
 *
 *   - To satisfy both constraints at once we keep `function_locations` required in the static TypeScript type
 *     (so `ivc-integration` keeps compiling) but relax the runtime input shape via this preprocessor (so legacy
 *     artifacts keep parsing). Once all v4-next artifacts get regenerated with the new Noir compiler, this
 *     helper — and both `z.preprocess(...)` wrappers — can be removed.
 */
function fillMissingFunctionLocations(val: unknown): unknown {
  if (val && typeof val === 'object') {
    for (const entry of Object.values(val as Record<string, unknown>)) {
      if (
        entry &&
        typeof entry === 'object' &&
        (entry as { function_locations?: unknown }).function_locations === undefined
      ) {
        (entry as { function_locations: FunctionLocation[] }).function_locations = [];
      }
    }
  }
  return val;
}

/** Type representing a field layout in the storage of a contract. */
export type FieldLayout = {
  /** Slot in which the field is stored. */
  slot: Fr;
};

/** Placeholder version injected into artifacts compiled before aztecVersion was added. TODO(F-557): Remove. */
export const ARTIFACT_VERSION_BEFORE_INJECTION = 'FROM_RELEASE_BEFORE_VERSION_INJECTION';

/** Defines artifact of a contract. */
export interface ContractArtifact {
  /** The name of the contract. */
  name: string;

  /** The version of the Aztec stack that compiled this artifact. */
  aztecVersion: string;

  /** The functions of the contract. Includes private and utility functions, plus the public dispatch function. */
  functions: FunctionArtifact[];

  /** The public functions of the contract, excluding dispatch. */
  nonDispatchPublicFunctions: FunctionAbi[];

  /** The outputs of the contract. */
  outputs: {
    structs: Record<string, AbiType[]>;
    globals: Record<string, AbiValue[]>;
  };

  /** Storage layout */
  storageLayout: Record<string, FieldLayout>;

  /** The map of file ID to the source code and path of the file. */
  fileMap: DebugFileMap;
}

export const ContractArtifactSchema = zodFor<ContractArtifact>()(
  z.object({
    name: z.string(),
    aztecVersion: z.string().default(ARTIFACT_VERSION_BEFORE_INJECTION), // TODO(F-557): Remove default.
    functions: z.array(FunctionArtifactSchema),
    nonDispatchPublicFunctions: z.array(FunctionAbiSchema),
    outputs: z.object({
      structs: z.record(z.array(AbiTypeSchema)).transform(structs => {
        for (const [key, value] of Object.entries(structs)) {
          // We are manually ordering events and functions in the abi by path.
          // The path ordering is arbitrary, and only needed to ensure deterministic order.
          // These are the only arrays in the artifact with arbitrary order, and hence the only ones
          // we need to sort.
          if (key === 'events' || key === 'functions') {
            structs[key] = (value as StructType[]).sort((a, b) => (a.path > b.path ? -1 : 1));
          }
        }
        return structs;
      }),
      globals: z.record(z.array(AbiValueSchema)),
    }),
    storageLayout: z.record(z.object({ slot: schemas.Fr })),
    // Legacy v4-next contract artifacts (e.g. `protocol-contracts/artifacts/AuthRegistry.json`) were emitted
    // before Noir started including `function_locations` in their `file_map` entries — see
    // {@link fillMissingFunctionLocations} for the full rationale. The preprocessor injects `[]` for missing
    // entries so parsing succeeds without weakening the static `DebugFileMap` type.
    fileMap: z.preprocess(
      fillMissingFunctionLocations,
      z.record(
        z.coerce.number(),
        z.object({
          source: z.string(),
          path: z.string(),
          function_locations: z.array(z.object({ start: z.number(), name: z.string() })),
        }),
      ),
    ),
  }),
);

export function getFunctionArtifactByName(artifact: ContractArtifact, functionName: string): FunctionArtifact {
  const functionArtifact = artifact.functions.find(f => f.name === functionName);

  if (!functionArtifact) {
    throw new Error(`Unknown function ${functionName}`);
  }

  const debugMetadata = getFunctionDebugMetadata(artifact, functionArtifact);
  return { ...functionArtifact, debug: debugMetadata };
}

/** Gets a function artifact including debug metadata given its name or selector. */
export async function getFunctionArtifact(
  artifact: ContractArtifact,
  functionNameOrSelector: string | FunctionSelector,
): Promise<FunctionArtifactWithContractName> {
  let functionArtifact;
  if (typeof functionNameOrSelector === 'string') {
    functionArtifact = artifact.functions.find(f => f.name === functionNameOrSelector);
  } else {
    const functionsAndSelectors = await Promise.all(
      artifact.functions.map(async fn => ({
        fn,
        selector: await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters),
      })),
    );
    functionArtifact = functionsAndSelectors.find(fnAndSelector =>
      functionNameOrSelector.equals(fnAndSelector.selector),
    )?.fn;
  }
  if (!functionArtifact) {
    throw new Error(`Unknown function ${functionNameOrSelector}`);
  }

  const debugMetadata = getFunctionDebugMetadata(artifact, functionArtifact);

  return { ...functionArtifact, debug: debugMetadata, contractName: artifact.name };
}

/** Gets all function abis */
export function getAllFunctionAbis(artifact: ContractArtifact): FunctionAbi[] {
  return artifact.functions.map(f => f as FunctionAbi).concat(artifact.nonDispatchPublicFunctions || []);
}

export function parseDebugSymbols(debugSymbols: string): DebugInfo[] {
  return JSON.parse(inflate(Buffer.from(debugSymbols, 'base64'), { to: 'string', raw: true })).debug_infos;
}

/**
 * Gets the debug metadata of a given function from the contract artifact
 * @param artifact - The contract build artifact
 * @param functionName - The name of the function
 * @returns The debug metadata of the function
 */
export function getFunctionDebugMetadata(
  contractArtifact: ContractArtifact,
  functionArtifact: FunctionArtifact,
): FunctionDebugMetadata | undefined {
  try {
    if (functionArtifact.debugSymbols && contractArtifact.fileMap) {
      // TODO(https://github.com/AztecProtocol/aztec-packages/issues/10546) investigate why debugMetadata is so big for some tests.
      const programDebugSymbols = parseDebugSymbols(functionArtifact.debugSymbols);
      // TODO(https://github.com/AztecProtocol/aztec-packages/issues/5813)
      // We only support handling debug info for the contract function entry point.
      // So for now we simply index into the first debug info.
      return {
        debugSymbols: programDebugSymbols[0],
        files: contractArtifact.fileMap,
      };
    }
  } catch (err: any) {
    if (err instanceof RangeError && err.message.includes('Invalid string length')) {
      logger.warn(
        `Caught RangeError: Invalid string length. This suggests the debug_symbols field of the contract ${contractArtifact.name} and function ${functionArtifact.name} is huge; too big to parse. We'll skip returning this info until this issue is resolved. Here's the error:\n${err.message}`,
      );
      // We'll return undefined.
    } else {
      // Rethrow unexpected errors
      throw err;
    }
  }

  return undefined;
}

/**
 * Returns an initializer from the contract, assuming there is at least one. If there are multiple initializers,
 * it returns the one named "constructor" or "initializer"; if there is none with that name, it returns the first
 * initializer it finds, prioritizing initializers with no arguments and then private ones.
 * @param contractArtifact - The contract artifact.
 * @returns An initializer function, or none if there are no functions flagged as initializers in the contract.
 */
export function getDefaultInitializer(contractArtifact: ContractArtifact): FunctionAbi | undefined {
  const functionAbis = getAllFunctionAbis(contractArtifact);
  const initializers = functionAbis.filter(f => f.isInitializer);
  return initializers.length > 1
    ? (initializers.find(f => f.name === 'constructor') ??
        initializers.find(f => f.name === 'initializer') ??
        initializers.find(f => f.parameters?.length === 0) ??
        initializers.find(f => f.functionType === FunctionType.PRIVATE) ??
        initializers[0])
    : initializers[0];
}

/**
 * Returns an initializer from the contract.
 * @param initializerNameOrArtifact - The name of the constructor, or the artifact of the constructor, or undefined
 * to pick the default initializer.
 */
export function getInitializer(
  contract: ContractArtifact,
  initializerNameOrArtifact: string | undefined | FunctionArtifact,
): FunctionAbi | undefined {
  if (typeof initializerNameOrArtifact === 'string') {
    const functionAbis = getAllFunctionAbis(contract);
    const found = functionAbis.find(f => f.name === initializerNameOrArtifact);
    if (!found) {
      throw new Error(`Constructor method ${initializerNameOrArtifact} not found in contract artifact`);
    } else if (!found.isInitializer) {
      throw new Error(`Method ${initializerNameOrArtifact} is not an initializer`);
    }
    return found;
  } else if (initializerNameOrArtifact === undefined) {
    return getDefaultInitializer(contract);
  } else {
    if (!initializerNameOrArtifact.isInitializer) {
      throw new Error(`Method ${initializerNameOrArtifact.name} is not an initializer`);
    }
    return initializerNameOrArtifact;
  }
}

export function emptyFunctionAbi(): FunctionAbi {
  return {
    name: '',
    functionType: FunctionType.PRIVATE,
    isOnlySelf: false,
    isStatic: false,
    parameters: [],
    returnTypes: [],
    errorTypes: {},
    isInitializer: false,
  };
}

export function emptyFunctionArtifact(): FunctionArtifact {
  const abi = emptyFunctionAbi();
  return {
    ...abi,
    bytecode: Buffer.from([]),
    debugSymbols: '',
  };
}

export function emptyContractArtifact(): ContractArtifact {
  return {
    name: '',
    aztecVersion: DEV_VERSION,
    functions: [emptyFunctionArtifact()],
    nonDispatchPublicFunctions: [emptyFunctionAbi()],
    outputs: {
      structs: {},
      globals: {},
    },
    storageLayout: {},
    fileMap: {},
  };
}