/**
 * @for p5
 */
import * as constants from '../core/constants.js';
import * as z from 'zod/mini';
import dataDoc from '../../docs/parameterData.json';
import { FES } from './fes.js';
import { errorStackParser, processStack, getFriendlyStack } from './stacktrace.js';

z.config(z.locales.en());
let documentationData = dataDoc;

function validateParams(p5, fn, lifecycles) {
  // Cache for Zod schemas
  let schemaRegistry = new Map();

  // Mapping names of p5 types to their constructor functions.
  // p5Constructors:
  //   - Color: f()
  //   - Graphics: f()
  //   - Vector: f()
  // and so on.
  // const p5Constructors = {};
  // NOTE: This is a tempt fix for unit test but is not correct
  // Attaced constructors are `undefined`
  const p5Constructors = Object.keys(p5).reduce((acc, val) => {
    if (
      val.match(/^[A-Z]/) && // Starts with a capital
      !val.match(/^[A-Z][A-Z0-9]*$/) && // Is not an all caps constant
      p5[val] instanceof Function // Is a function
    ) {
      acc[val] = p5[val];
    }
    return acc;
  }, {});

  function loadP5Constructors() {
    // Make a list of all p5 classes to be used for argument validation
    // This must be done only when everything has loaded otherwise we get
    // an empty array
    for (let key of Object.keys(p5)) {
      // Get a list of all constructors in p5. They are functions whose names
      // start with a capital letter
      if (typeof p5[key] === 'function' && key[0] !== key[0].toLowerCase()) {
        p5Constructors[key] = p5[key];
      }
    }
  }

  // `constantsMap` maps constants to their values, e.g.
  // {
  //   ADD: 'lighter',
  //   ALT: 18,
  //   ARROW: 'default',
  //   AUTO: 'auto',
  //   ...
  // }
  const constantsMap = {};
  for (const [key, value] of Object.entries(constants)) {
    constantsMap[key] = value;
  }

  // Start initializing `schemaMap` with primitive types. `schemaMap` will
  // eventually contain both primitive types and web API objects.
  const schemaMap = {
    'Any': z.any(),
    'Array': z.array(z.any()),
    'Boolean': z.boolean(),
    'Function': z.function(),
    'Integer': z.number().check(z.int()),
    'Number': z.union([z.number(), z.literal(Infinity), z.literal(-Infinity)]),
    'Object': z.object({}),
    'String': z.string()
  };

  const webAPIObjects = [
    'AudioNode',
    'HTMLCanvasElement',
    'HTMLElement',
    'KeyboardEvent',
    'MouseEvent',
    'RegExp',
    'TouchEvent',
    'UIEvent',
    'WheelEvent'
  ];

  function generateWebAPISchemas(apiObjects) {
    return apiObjects.reduce((acc, obj) => {
      acc[obj] = z.custom(data => data instanceof globalThis[obj], {
        message: `Expected a ${obj}`
      });
      return acc;
    }, {});
  }

  const webAPISchemas = generateWebAPISchemas(webAPIObjects);
  // Add web API schemas to the schema map.
  Object.assign(schemaMap, webAPISchemas);

  // For mapping 0-indexed parameters to their ordinal representation, e.g.
  // "first" for 0, "second" for 1, "third" for 2, etc.
  const ordinals = ['first', 'second', 'third', 'fourth', 'fifth', 'sixth', 'seventh', 'eighth', 'ninth', 'tenth'];

  function extractFuncNameAndClass(func) {
    const ichDot = func.lastIndexOf('.');
    const funcName = func.slice(ichDot + 1);
    const funcClass = func.slice(0, ichDot !== -1 ? ichDot : 0) || 'p5';
    return { funcName, funcClass };
  }

  function validBracketNesting(type) {
    let level = 0;
    for (let i = 0; i < type.length; i++) {
      if (type[i] === '[') {
        level++;
      } else if (type[i] === ']') {
        level--;
        if (level < 0) return false;
      }
    }
    return level === 0;
  }

  /**
   * This is a helper function that generates Zod schemas for a function based on
   * the parameter data from `docs/parameterData.json`.
   *
   * Example parameter data for function `background`:
   * "background": {
   *    "overloads": [
   *      ["p5.Color"],
   *      ["String", "Number?"],
   *      ["Number", "Number?"],
   *      ["Number", "Number", "Number", "Number?"],
   *      ["Number[]"],
   *      ["p5.Image", "Number?"]
   *    ]
   *  }
   * Where each array in `overloads` represents a set of valid overloaded
   * parameters, and `?` is a shorthand for `Optional`.
   *
   * @method generateZodSchemasForFunc
   * @private
   * @param {String} func - Name of the function. Expect global functions like `sin` and class methods like `p5.Vector.add`
   * @returns {z.ZodSchema} Zod schema
   */
  const generateZodSchemasForFunc = function (func) {
    const { funcName, funcClass } = extractFuncNameAndClass(func);
    let funcInfo = documentationData[funcClass][funcName];

    if(!funcInfo) return;

    let overloads = [];
    if (funcInfo.hasOwnProperty('overloads')) {
      overloads = funcInfo.overloads;
    }

    // Returns a schema for a single type, i.e. z.boolean() for `boolean`.
    const generateTypeSchema = baseType => {
      if (!baseType) return z.any();

      let typeSchema;

      // Check for constants. Note that because we're ultimately interested in the value of
      // the constant, mapping constants to their values via `constantsMap` is
      // necessary.
      if (baseType in constantsMap) {
        typeSchema = z.literal(constantsMap[baseType]);
      }
      // Some more constants are attached directly to p5.prototype, e.g. by addons:
      else if (baseType.match(/^[A-Z][A-Z0-9]*$/) && baseType in fn) {
        typeSchema = z.literal(fn[baseType]);
      }
      // Function types
      else if (baseType.startsWith('function')) {
        typeSchema = z.function();
      }
      // All p5 objects start with `p5` in the documentation, i.e. `p5.Camera`.
      else if (/^p5\.[a-zA-Z0-9]+$/.exec(baseType) || baseType === 'p5') {
        const className = baseType.substring(baseType.indexOf('.') + 1);
        typeSchema = z.instanceof(p5Constructors[className]);
      }
      // For primitive types and web API objects.
      else if (schemaMap[baseType]) {
        typeSchema = schemaMap[baseType];
      }
      // Tuple types
      else if (
        baseType.startsWith('[') &&
        baseType.endsWith(']') &&
        validBracketNesting(baseType.slice(1, -1))
      ) {
        typeSchema = z.tuple(
          baseType
            .slice(1, -1)
            .split(/, */g)
            .map(entry => generateTypeSchema(entry))
        );
      }
      // JavaScript classes, e.g. Request
      else if (baseType.match(/^[A-Z]/) && baseType in window) {
        typeSchema = z.instanceof(window[baseType]);
      }
      // Generate a schema for a single parameter that can be of multiple
      // types / constants, i.e. `String|Number|Array`.
      //
      // Here, z.union() is used over z.enum() (which seems more intuitive) for
      // constants for the following reasons:
      // 1) z.enum() only allows a fixed set of allowable string values. However,
      // our constants sometimes have numeric or non-primitive values.
      // 2) In some cases, the type can be constants or strings, making z.enum()
      // insufficient for the use case.
      else if (baseType.includes('|') && baseType.split('|').every(t => validBracketNesting(t))) {
        const types = baseType.split('|');
        typeSchema = z.union(types
          .map(t => generateTypeSchema(t))
          .filter(s => s !== undefined));
      } else if (baseType.endsWith('[]')) {
        typeSchema = z.array(generateTypeSchema(baseType.slice(0, -2)));
      } else {
        throw new Error(`Unsupported type '${baseType}' in parameter validation. Please report this issue.`);
      }

      return typeSchema;
    };

    // Generate a schema for a single parameter. In the case where a parameter can
    // be of multiple types, `generateTypeSchema` is called for each type.
    const generateParamSchema = param => {
      const isOptional = param?.endsWith('?');
      param = param?.replace(/\?$/, '');

      const isRest = param?.startsWith('...') && param?.endsWith('[]');
      param = param?.replace(/^\.\.\.(.+)\[\]$/, '$1');

      let schema = generateTypeSchema(param);
      // Fallback to z.custom() because function types are no longer
      // returns a Zod schema.
      if (schema.def.type === 'function') {
        schema = z.custom(val => val instanceof Function);
      }

      if (isOptional) {
        schema = z.optional(schema);
      }
      return { schema, rest: isRest };
    };

    // Note that in Zod, `optional()` only checks for undefined, not the absence
    // of value.
    //
    // Let's say we have a function with 3 parameters, and the last one is
    // optional, i.e. func(a, b, c?). If we only have a z.tuple() for the
    // parameters, where the third schema is optional, then we will only be able
    // to validate func(10, 10, undefined), but not func(10, 10), which is
    // a completely valid call.
    //
    // Therefore, on top of using `optional()`, we also have to generate parameter
    // combinations that are valid for all numbers of parameters.
    const generateOverloadCombinations = params => {
      // No optional parameters, return the original parameter list right away.
      if (!params.some(p => p?.endsWith('?'))) {
        return [params];
      }

      const requiredParamsCount = params.filter(p => p === null || !p.endsWith('?')).length;
      const result = [];

      for (let i = requiredParamsCount; i <= params.length; i++) {
        result.push(params.slice(0, i));
      }

      return result;
    };

    // Generate schemas for each function overload and merge them
    const overloadSchemas = overloads.flatMap(overload => {
      const combinations = generateOverloadCombinations(overload);

      return combinations.map(combo => {
        const params = combo
          .map(p => generateParamSchema(p))
          .filter(s => s.schema !== undefined);

        let rest;
        if (params.at(-1)?.rest) {
          rest = params.pop();
        }

        let combined = z.tuple(params.map(s => s.schema), rest?.schema);
        // if (rest) {
        //   combined = combined.rest(rest.schema);
        // }
        return combined;
      });
    });

    return overloadSchemas.length === 1
      ? overloadSchemas[0]
      : z.union(overloadSchemas);
  };

  /**
   * Finds the closest schema to the input arguments.
   *
   * This is a helper function that identifies the closest schema to the input
   * arguments, in the case of an initial validation error. We will then use the
   * closest schema to generate a friendly error message.
   *
   * @private
   * @param {z.ZodSchema} schema - Zod schema.
   * @param {Array} args - User input arguments.
   * @returns {z.ZodSchema} Closest schema matching the input arguments.
   */
  const findClosestSchema = function (schema, args) {
    if (!(schema instanceof z.ZodMiniUnion)) {
      return schema;
    }

    // Helper function that scores how close the input arguments are to a schema.
    // Lower score means closer match.
    const scoreSchema = schema => {
      let score = Infinity;
      if (!(schema instanceof z.ZodMiniTuple)) {
        console.warn('Schema below is not a tuple: ');
        printZodSchema(schema);
        return score;
      }

      const numArgs = args.length;
      const schemaItems = schema.def.items;
      const numSchemaItems = schemaItems.length;
      const numRequiredSchemaItems = schemaItems
        .filter(item => !item.safeParse(undefined).success)
        .length;

      if (numArgs >= numRequiredSchemaItems && numArgs <= numSchemaItems) {
        score = 0;
      }
      // Here, give more weight to mismatch in number of arguments.
      //
      // For example, color() can either take [Number, Number?] or
      // [Number, Number, Number, Number?] as list of parameters.
      // If the user passed in 3 arguments, [10, undefined, undefined], it's
      // more than likely that they intended to pass in 3 arguments, but the
      // last two arguments are invalid.
      //
      // If there's no bias towards matching the number of arguments, the error
      // message will show that we're expecting at most 2 arguments, but more
      // are received.
      else {
        score = Math.abs(
          numArgs < numRequiredSchemaItems ?
            numRequiredSchemaItems - numArgs :
            numArgs - numSchemaItems
        ) * 4;
      }

      for (let i = 0; i < Math.min(schemaItems.length, args.length); i++) {
        const paramSchema = schemaItems[i];
        const arg = args[i];

        if (!paramSchema.safeParse(arg).success) score++;
      }

      return score;
    };

    // Default to the first schema, so that we are guaranteed to return a result.
    let closestSchema = schema.def.options[0];
    // We want to return the schema with the lowest score.
    let bestScore = Infinity;

    const schemaUnion = schema.def.options;
    schemaUnion.forEach(schema => {
      const score = scoreSchema(schema);
      if (score < bestScore) {
        closestSchema = schema;
        bestScore = score;
      }
    });

    return closestSchema;
  };

  /**
   * Prints a friendly error message after parameter validation, if validation
   * has failed.
   *
   * @method _friendlyParamError
   * @private
   * @param {z.ZodError} zodErrorObj - The Zod error object containing validation errors.
   * @param {String} func - Name of the function. Expect global functions like `sin` and class methods like `p5.Vector.add`
   * @returns {String} The friendly error message.
   */
  const friendlyParamError = function (zodErrorObj, func, args) {
    let message = '';

    let isVersionError = false;
    // The `zodErrorObj` might contain multiple errors of equal importance
    // (after scoring the schema closeness in `findClosestSchema`). Here, we
    // always print the first error so that user can work through the errors
    // one by one.
    let currentError = zodErrorObj.issues[0];

    // Helper function to build a type mismatch message.
    const buildTypeMismatchMessage =
      (actualType, expectedTypeStr, position) => {
        const positionStr = position ? FES.log`at the ${ordinals[position]} parameter` : '';
        const actualTypeStr = actualType ? FES.log`, but received ${actualType}` : '';
        return FES.log`Expected ${expectedTypeStr} ${positionStr}${actualTypeStr}`;
      };

    // Union errors occur when a parameter can be of multiple types but is not
    // of any of them. In this case, aggregate all possible types and print
    // a friendly error message that indicates what the expected types are at
    // which position (position is not 0-indexed, for accessibility reasons).

    const processUnionError = error => {
      const expectedTypes = new Set();
      let actualType, message = '';

      const collectIssue = issue => {
        if (!issue) return;
        if (!actualType) {
          actualType = issue.message;
        }

        if (issue.code === 'invalid_type') {
          actualType = issue.message.split(', received ')[1];
          expectedTypes.add(issue.expected);
        }
        // The case for constants. Since we don't want to print out the actual
        // constant values in the error message, the error message will
        // direct users to the documentation.
        else if (issue.code === 'invalid_value') {
          if (
            Array.isArray(issue.values) &&
            issue.values.every(v => v === Infinity || v === -Infinity)
          ) {
            expectedTypes.add('number');
          } else {
            expectedTypes.add('constant (please refer to documentation for allowed values)');
            actualType = args[error.path[0]];
          }
        } else if (issue.code === 'custom') {
          const match = issue.message.match(/Input not instance of (\w+)/);
          if (match) expectedTypes.add(match[1]);
          actualType = undefined;
        } else if (issue.code === 'invalid_union') {
          issue.errors.forEach(nestedErr => {
            nestedErr.forEach(nestedIssue => collectIssue(nestedIssue));
          });
        }
      };

      error.errors.forEach(err => {
        err.forEach(issue => collectIssue(issue));
      });

      if (expectedTypes.size > 0) {
        if (error.path?.length > 0 && args[error.path[0]] instanceof Promise)  {
          message = FES.log`Did you mean to put \`await\` before a loading function? An unexpected Promise was found. `;
        }

        const expectedTypesStr = Array.from(expectedTypes).join(' or ');
        const position = error.path.join('.');

        message = FES.log`${message}${buildTypeMismatchMessage(
          actualType, expectedTypesStr, position
        )} in ${func + '()'}.`;
      }

      return message;
    };

    // Generates a link to the documentation based on the given function name.
    const generateDocumentationLink = func => {
      const { funcName, funcClass } = extractFuncNameAndClass(func);
      const p5BaseUrl = 'https://p5js.org/reference';
      const url = `${p5BaseUrl}/${funcClass}/${funcName}`;

      return url;
    };

    switch (currentError.code) {
      case 'invalid_union': {
        message = processUnionError(currentError);
        break;
      }
      case 'too_small': {
        const minArgs = currentError.minimum;
        const documentationLink = generateDocumentationLink(func);
        const referenceLink = FES.log`For more information, see ${documentationLink}.`;
        message = FES.log`Expected at least ${minArgs} argument, but received fewer in ${func}(). ${referenceLink}`;
        break;
      }
      case 'invalid_type': {
        const position = FES.premade.ordinals[currentError.path.join('.')];
        const expectedType = FES.premade.types[currentError.expected];
        // NOTE: need type info
        message = FES.log`Expected ${expectedType} at the ${position} parameter in ${func + '()'}.`
        break;
      }
      case 'too_big': {
        const maxArgs = currentError.maximum;
        const documentationLink = generateDocumentationLink(func);
        const referenceLink = FES.log`For more information, see ${documentationLink}.`;
        message = FES.log`Expected at most ${maxArgs} arguments, but received more in ${func}(). ${referenceLink}`;
        break;
      }
      case 'custom': {
        const position = FES.premade.ordinals[currentError.path.join('.')];
        const match = currentError.message.match(/Input not instance of (\w+)/);
        if (match) {
          message = FES.log`Expected ${match[1]} at the ${position} parameter in ${func + '()'}.`
          break;
        }
      }
      default: {
        console.log('Zod error object', currentError);
      }
    }

    if (isVersionError) {
      FES.log`${message}`();
    } else {
      const [_null, stacktrace] = processStack(
        null,
        errorStackParser.parse(Error()).slice(3)
      );
      const locationStr = getFriendlyStack(stacktrace, true);
      FES.log`${locationStr} ${message}`();
    }
    return message;
  };

  /**
   * Runs parameter validation by matching the input parameters to Zod schemas
   * generated from the parameter data from `docs/parameterData.json`.
   *
   * @private
   * @param {String} func - Name of the function.
   * @param {Array} args - User input arguments.
   * @returns {Object} The validation result.
   * @returns {Boolean} result.success - Whether the validation was successful.
   * @returns {any} [result.data] - The parsed data if validation was successful.
   * @returns {String} [result.error] - The validation error message if validation has failed.
   */
  const validate = function (func, args) {
    if (p5.disableFriendlyErrors) {
      return; // skip FES
    }

    if (!Array.isArray(args)) {
      args = Array.from(args);
    }

    // An edge case: even when all arguments are optional and therefore,
    // theoretically allowed to stay undefined and valid, it is likely that the
    // user intended to call the function with non-undefined arguments. Skip
    // regular workflow and return a friendly error message right away.
    if (
      Array.isArray(args) &&
      args.length > 0 &&
      args.every(arg => arg === undefined)
    ) {
      const undefinedErrorMessage = `All arguments for ${func}() are undefined. There is likely an error in the code.`;

      return {
        success: false,
        error: undefinedErrorMessage
      };
    }

    let funcSchemas = schemaRegistry.get(func);
    if (!funcSchemas) {
      funcSchemas = generateZodSchemasForFunc(func);
      if (!funcSchemas) return;
      schemaRegistry.set(func, funcSchemas);
    }

    try {
      return {
        success: true,
        data: funcSchemas.parse(args)
      };
    } catch (error) {
      const closestSchema = findClosestSchema(funcSchemas, args);
      const zodError = closestSchema.safeParse(args).error;
      const errorMessage = friendlyParamError(zodError, func, args);

      return {
        success: false,
        error: errorMessage
      };
    }
  };

  fn._validate = validate; // TEMP: For unit tests

  // Suppress FES param checking for the duration of a callback.
  // Use this to wrap internal p5 calls that happen after an await.
  // NOTE: shares the same _isUserCall flag logic as the decorator below.
  fn._internal = function(callback) {
    const wasInternalCall = this._isUserCall;
    this._isUserCall = true;
    try {
      return callback();
    } finally {
      this._isUserCall = wasInternalCall;
    }
  };

  // Skip FES validation for nested (internal) calls.
  // NOTE: shares the same _isUserCall flag logic as _internal() above.
  p5.registerDecorator(
    ({ path }) => {
      return path.startsWith('p5.prototype');
    },
    function(target, { kind, name }){
      if(kind === 'method'){
        return function(...args){
          if (p5.disableFriendlyErrors) {
            return target.apply(this, args);
          }
          const wasInternalCall = this._isUserCall;
          this._isUserCall = true;
          try {
            if (!wasInternalCall && !p5.disableFriendlyErrors && !p5.disableParameterValidator) {
              validate(name, args);
            }
            return target.apply(this, args);
          } finally {
            this._isUserCall = wasInternalCall;
          }
        };
      }
    }
  );

  p5.extendParameterValidation = (data) => {
    documentationData = mergeDeep(documentationData, data);
  };

  lifecycles.presetup = function(){
    loadP5Constructors();
  };
}

function isObject(item) {
  return (item && typeof item === 'object' && !Array.isArray(item));
}

// Deep merge implementation: https://stackoverflow.com/questions/27936772/how-to-deep-merge-instead-of-shallow-merge
function mergeDeep(target, ...sources) {
  if (!sources.length) return target;
  const source = sources.shift();

  if (isObject(target) && isObject(source)) {
    for (const key in source) {
      if (isObject(source[key])) {
        if (!target[key]) Object.assign(target, { [key]: {} });
        mergeDeep(target[key], source[key]);
      } else {
        Object.assign(target, { [key]: source[key] });
      }
    }
  }

  return mergeDeep(target, ...sources);
}

export default validateParams;

if (typeof p5 !== 'undefined') {
  p5.registerAddon(validateParams);
}
