import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';

import {
  AZTEC_INITIALIZER_ATTRIBUTE,
  AZTEC_ONLY_SELF_ATTRIBUTE,
  AZTEC_PRIVATE_ATTRIBUTE,
  AZTEC_PUBLIC_ATTRIBUTE,
  AZTEC_UTILITY_ATTRIBUTE,
  AZTEC_VIEW_ATTRIBUTE,
  type NoirCompiledContract,
} from '../noir/index.js';
import {
  type ABIParameter,
  type ABIParameterVisibility,
  ARTIFACT_VERSION_BEFORE_INJECTION,
  type AbiType,
  type BasicValue,
  type ContractArtifact,
  ContractArtifactSchema,
  type FieldLayout,
  type FunctionAbi,
  type FunctionArtifact,
  FunctionType,
  type IntegerValue,
  type StructValue,
  type TypedStructFieldValue,
} from './abi.js';

/**
 * Serializes a contract artifact to a buffer for storage.
 * @param artifact - Artifact to serialize.
 * @returns A buffer.
 */
export function contractArtifactToBuffer(artifact: ContractArtifact): Buffer {
  return Buffer.from(jsonStringify(artifact), 'utf-8');
}

/**
 * Deserializes a contract artifact from storage.
 * @param buffer - Buffer to deserialize.
 * @returns Deserialized artifact.
 */
export function contractArtifactFromBuffer(buffer: Buffer): ContractArtifact {
  return jsonParseWithSchema(buffer.toString('utf-8'), ContractArtifactSchema);
}

/**
 * Gets nargo build output and returns a valid contract artifact instance.
 * Does not include public bytecode, apart from the public_dispatch function.
 * @param input - Input object as generated by nargo compile.
 * @returns A valid contract artifact instance.
 */
export function loadContractArtifact(input: NoirCompiledContract): ContractArtifact {
  if (isContractArtifact(input)) {
    // TODO(F-557): Remove this fallback once pre-version artifacts are no longer tested.
    if (!(input as unknown as Record<string, unknown>).aztecVersion) {
      return { ...input, aztecVersion: ARTIFACT_VERSION_BEFORE_INJECTION };
    }
    return input;
  }
  return generateContractArtifact(input);
}

/**
 * Gets nargo build output and returns a valid contract artifact instance.
 * Differs from loadContractArtifact() by retaining all bytecode.
 * @param input - Input object as generated by nargo compile.
 * @returns A valid contract artifact instance.
 */
export function loadContractArtifactForPublic(input: NoirCompiledContract): ContractArtifact {
  return generateContractArtifactForPublic(input);
}

/**
 * Checks if the given input looks like a valid ContractArtifact. The check is not exhaustive,
 * and it's just meant to differentiate between nargo raw build artifacts and the ones
 * produced by this compiler.
 * @param input - Input object.
 * @returns True if it looks like a ContractArtifact.
 */
function isContractArtifact(input: any): input is ContractArtifact {
  if (typeof input !== 'object') {
    return false;
  }
  const maybeContractArtifact = input as ContractArtifact;
  if (typeof maybeContractArtifact.name !== 'string') {
    return false;
  }
  if (!Array.isArray(maybeContractArtifact.functions)) {
    return false;
  }
  for (const fn of maybeContractArtifact.functions) {
    if (typeof fn.name !== 'string') {
      return false;
    }
    if (typeof fn.functionType !== 'string') {
      return false;
    }
    if (!retainBytecode(fn) && fn.bytecode.length) {
      // We want to remove the bytecode of public fns (apart from the dispatch fn) to save space
      // If the input is private-only, we don't need to use generateContractArtifact() below
      return false;
    }
  }
  if (!Array.isArray(maybeContractArtifact.nonDispatchPublicFunctions)) {
    return false;
  }
  for (const fn of maybeContractArtifact.nonDispatchPublicFunctions) {
    if (typeof fn.name !== 'string') {
      return false;
    }
    if (typeof fn.functionType !== 'string') {
      return false;
    }
  }
  return true;
}

/** Parameter in a function from a noir contract compilation artifact */
type NoirCompiledContractFunctionParameter = NoirCompiledContractFunction['abi']['parameters'][number];

/**
 * Generates a function parameter out of one generated by a nargo build.
 * @param param - Noir parameter.
 * @returns A function parameter.
 */
function generateFunctionParameter(param: NoirCompiledContractFunctionParameter): ABIParameter {
  const { visibility } = param;
  if ((visibility as string) === 'databus') {
    throw new Error(`Unsupported visibility ${param.visibility} for noir contract function parameter ${param.name}.`);
  }
  return { ...param, visibility: visibility as ABIParameterVisibility };
}

/** Function from a noir contract compilation artifact */
type NoirCompiledContractFunction = NoirCompiledContract['functions'][number];

/**
 * Returns true if we should retain bytecode
 */
export function retainBytecode(input: NoirCompiledContractFunction | FunctionArtifact): boolean {
  const functionType =
    (input as FunctionArtifact).functionType ?? getFunctionType(input as NoirCompiledContractFunction);
  return functionType !== FunctionType.PUBLIC || input.name == 'public_dispatch';
}

/**
 * Generates a function abi.
 * @param fn - Noir function entry.
 * @param contract - Parent contract.
 * @returns Function abi.
 */
function generateFunctionAbi(fn: NoirCompiledContractFunction, contract: NoirCompiledContract): FunctionAbi {
  if (fn.custom_attributes === undefined) {
    throw new Error(
      `No custom attributes found for contract function ${fn.name}. Try rebuilding the contract with the latest nargo version.`,
    );
  }
  const functionType = getFunctionType(fn);
  const isOnlySelf = fn.custom_attributes.includes(AZTEC_ONLY_SELF_ATTRIBUTE);
  const isStatic = fn.custom_attributes.includes(AZTEC_VIEW_ATTRIBUTE);

  // If the function is not a utility function, the first item is inputs or CallContext which we should omit
  let parameters = fn.abi.parameters.map(generateFunctionParameter);
  if (hasKernelFunctionInputs(parameters)) {
    parameters = parameters.slice(1);
  }

  let returnTypes: AbiType[] = [];
  if (functionType === FunctionType.UTILITY) {
    returnTypes = fn.abi.return_type ? [fn.abi.return_type.abi_type] : returnTypes;
  } else {
    const pathToFind = `${contract.name}::${fn.name}_abi`;
    const abiStructs: AbiType[] = contract.outputs.structs['functions'];

    const returnStruct = abiStructs.find(a => a.kind === 'struct' && a.path === pathToFind);

    if (returnStruct) {
      if (returnStruct.kind !== 'struct') {
        throw new Error('Could not generate contract function artifact');
      }

      const returnTypeField = returnStruct.fields.find(field => field.name === 'return_type');
      if (returnTypeField) {
        returnTypes = [returnTypeField.type];
      }
    }
  }

  return {
    name: fn.name,
    functionType,
    isOnlySelf,
    isStatic,
    isInitializer: fn.custom_attributes.includes(AZTEC_INITIALIZER_ATTRIBUTE),
    parameters,
    returnTypes,
    errorTypes: fn.abi.error_types,
    ...(fn.assert_messages ? { assertMessages: fn.assert_messages } : undefined),
  };
}

/**
 * Generates a function build artifact.
 * @param fn - Noir function entry.
 * @param contract - Parent contract.
 * @returns Function artifact.
 */
function generateFunctionArtifact(
  fn: NoirCompiledContractFunction,
  contract: NoirCompiledContract,
): Omit<FunctionArtifact, 'bytecode'> & { bytecode: string } {
  const abi = generateFunctionAbi(fn, contract);
  return {
    ...abi,
    bytecode: fn.bytecode,
    debugSymbols: fn.debug_symbols,
    ...(fn.verification_key ? { verificationKey: fn.verification_key } : undefined),
  };
}

function getFunctionType(fn: NoirCompiledContractFunction): FunctionType {
  if (fn.custom_attributes.includes(AZTEC_PRIVATE_ATTRIBUTE)) {
    return FunctionType.PRIVATE;
  } else if (fn.custom_attributes.includes(AZTEC_PUBLIC_ATTRIBUTE)) {
    return FunctionType.PUBLIC;
  } else if (fn.custom_attributes.includes(AZTEC_UTILITY_ATTRIBUTE)) {
    return FunctionType.UTILITY;
  } else {
    throw new Error(`Invalid function type for a noir contract function ${fn.name}`);
  }
}

/**
 * Returns true if the first parameter is kernel function inputs.
 *
 * Noir macros #[aztec(private|public)] inject the following code
 * fn <name>(inputs: <Public|Private>ContextInputs, ...otherparams) {}
 *
 * Return true if this injected parameter is found
 */
function hasKernelFunctionInputs(params: ABIParameter[]): boolean {
  const firstParam = params[0];
  return firstParam?.type.kind === 'struct' && firstParam.type.path.includes('ContextInputs');
}

/**
 * Generates a storage layout for the contract artifact.
 * @param input - The compiled noir contract to get storage layout for
 * @returns A storage layout for the contract.
 */
function getStorageLayout(input: NoirCompiledContract) {
  // If another contract is imported by the main contract, its storage layout its going to also show up here.
  // The layout export includes the contract name, so here we can find the one that belongs to the current one and
  // ignore the rest.
  const storageExports = input.outputs.globals.storage ? (input.outputs.globals.storage as StructValue[]) : [];
  const storageForContract = storageExports.find(storageExport => {
    const contractNameField = storageExport.fields.find(field => field.name === 'contract_name')?.value as BasicValue<
      'string',
      string
    >;
    return contractNameField.value === input.name;
  });
  const storageFields = storageForContract
    ? ((storageForContract.fields.find(field => field.name == 'fields') as TypedStructFieldValue<StructValue>).value
        .fields as TypedStructFieldValue<StructValue>[])
    : [];

  if (storageFields.length === 0) {
    return {};
  }

  return storageFields.reduce((acc: Record<string, Omit<FieldLayout, 'slot'> & { slot: string }>, field) => {
    const name = field.name;
    const slot = field.value.fields[0].value as IntegerValue;
    acc[name] = {
      slot: `0x${slot.value}`,
    };
    return acc;
  }, {});
}

/**
 * Given a post-processed Nargo output defined as `contract` generates an Aztec-compatible contract artifact.
 *
 * Does not include public bytecode, apart from the public_dispatch function.
 */
function generateContractArtifact(contract: NoirCompiledContract): ContractArtifact {
  try {
    if (!contract.transpiled) {
      throw new Error("Contract's public bytecode has not been transpiled");
    }
    return ContractArtifactSchema.parse({
      name: contract.name,
      aztecVersion: contract.aztec_version,
      functions: contract.functions.filter(f => retainBytecode(f)).map(f => generateFunctionArtifact(f, contract)),
      nonDispatchPublicFunctions: contract.functions
        .filter(f => !retainBytecode(f))
        .map(f => generateFunctionAbi(f, contract)),
      outputs: contract.outputs,
      storageLayout: getStorageLayout(contract),
      fileMap: contract.file_map,
    });
  } catch (err) {
    throw new Error(`Could not generate contract artifact for ${contract.name}: ${err}`);
  }
}

/**
 * Given a post-processed Nargo output defined as `contract` generates an Aztec-compatible contract artifact.
 *
 * Retains all public bytecode.
 */
function generateContractArtifactForPublic(contract: NoirCompiledContract): ContractArtifact {
  try {
    return ContractArtifactSchema.parse({
      name: contract.name,
      aztecVersion: contract.aztec_version,
      functions: contract.functions.map(f => generateFunctionArtifact(f, contract)),
      nonDispatchPublicFunctions: contract.functions
        .filter(f => !retainBytecode(f))
        .map(f => generateFunctionAbi(f, contract)),
      outputs: contract.outputs,
      storageLayout: getStorageLayout(contract),
      fileMap: contract.file_map,
    });
  } catch (err) {
    throw new Error(`Could not generate contract artifact for ${contract.name}: ${err}`);
  }
}