import { Ast, ValueWithMode, Value } from './tokenToAstToken.js';
import { inspect as nodeInspect } from 'node:util';
import { capitalize } from 'lodash-es';

const ALLOWED_FIRST_CHARS = /^[A-z_$]/g;
const ALLOWED_CHARS = /[^A-z0-9_$]/g;
const GET_TOKEN_BY_MODE_FN_NAME = 'getTokenByMode';
const GET_TOKEN_WITH_MODES_FN_NAME = 'getTokenWithModes';
const TOKENS_VAR_NAME = 'tokens';
const TOKENS_TYPE_NAME = capitalize(TOKENS_VAR_NAME);
const ALL_PATHS_TYPE_NAME = 'AllPath';
const ALL_MODES_TYPE_NAME = 'AllMode';
const PATHS_BY_TYPE_VAR_NAME = 'pathsByType';
const PATHS_BY_TYPE_TYPE_NAME = 'PathsByType';
const GET_TOKENS_BY_TYPE_FN_NAME = 'getTokensByType';
const WITH_MODE_FN_NAME = 'withMode';

type commonParams = {
  ident: string;
  isTypescript: boolean;
  moduleExport: 'es6' | 'commonjs';
};

type Primitive =
  | string
  | number
  | Array<string | number | Record<string, number | string | Array<string | number>>>
  | Record<string, string | number | Array<string | number>>;

type PrimitiveWithMode = {
  [mode: string]: Primitive;
};

const inspect = (str: Parameters<typeof nodeInspect>[0]) =>
  nodeInspect(str, {
    depth: Infinity,
    maxArrayLength: Infinity,
    maxStringLength: Infinity,
  });

export function sanitizeIdent(ident: string) {
  return `${!ident.match(ALLOWED_FIRST_CHARS) ? '_' : ''}${ident.replaceAll(ALLOWED_CHARS, '_')}`;
}

export function generateTokens({
  tokens,
  isTypescript,
  moduleExport,
  tokensToUnquote,
}: Omit<commonParams, 'ident'> & {
  tokens: {
    [path: string]: PrimitiveWithMode | Primitive;
  };
  tokensToUnquote: Array<{ path: string; key?: string }>;
}) {
  let code = `/**
* @typedef {typeof ${TOKENS_VAR_NAME}} ${capitalize(TOKENS_VAR_NAME)} - All the tokens.
* Use \`${GET_TOKEN_BY_MODE_FN_NAME}\` to retrieve one.
*/\n`;

  code += `${
    moduleExport === 'es6' ? 'export ' : ''
  }const ${TOKENS_VAR_NAME} = /** @type {const} */ (${inspect(tokens)})${
    isTypescript ? ' as const' : ''
  };${
    moduleExport === 'commonjs'
      ? '\nmodule.exports.' + TOKENS_VAR_NAME + ' = ' + TOKENS_VAR_NAME + ';'
      : ''
  }`;

  if (isTypescript) {
    code += `\ntype ${TOKENS_TYPE_NAME} = typeof ${TOKENS_VAR_NAME};`;
  }

  tokensToUnquote.forEach(({ path, key }) => {
    code = unquoteToken(code, path, key);
  });

  return code;
}

function unquoteToken(code: string, path: string, key: string | undefined): string {
  return code.replace(
    key
      ? new RegExp(
          `(['"]?${path.replaceAll('.', '\\.')}['"]?: {[^]*['"]?${key}['"]?: )["'](.*)['"]`,
        )
      : new RegExp(`(['"]?${path.replaceAll('.', '\\.')}['"]?: )["'](.*)['"]`),
    `$1$2`,
  );
}

function makePathsIdent(ident: string, isType: boolean = false) {
  return `${isType ? capitalize(ident) : ident}Paths`;
}

export function generateAllPaths({
  ast,
  isTypescript,
  moduleExport,
}: Omit<commonParams, 'ident'> & { ast: Ast }) {
  let result = '';
  const allPaths: Array<string> = [];

  const typeIdents: Array<string> = [];

  Object.entries(ast.dataByType).forEach(
    ([
      ident,
      {
        metadata: { paths },
      },
    ]) => {
      if (paths.size === 0) return;

      const arrPaths = Array.from(paths);

      const varIdent = makePathsIdent(ident);
      const typeIdent = makePathsIdent(ident, true);

      let js = `/**
* @typedef {typeof ${varIdent}} ${typeIdent} - All the valid paths for the tokens of type ${ident}.
* To use this type you can do: \`@type {import('path/to/myTokensFile').${typeIdent}}\`
*/\n`;

      allPaths.push(typeIdent);

      typeIdents.push(ident);

      js += `${
        moduleExport === 'es6' ? 'export ' : ''
      }const ${varIdent} = /** @type {const} */ (${inspect(arrPaths)}${
        isTypescript ? ' as const' : ''
      });\n`;

      if (moduleExport === 'commonjs') {
        js += `module.exports.${varIdent} = ${varIdent};\n\n`;
      }

      if (isTypescript) {
        js += `export type ${typeIdent} = typeof ${varIdent};\n\n`;
      }

      result += js;
    },
  );

  Object.entries(ast.collectionPaths).forEach(([ident, paths]) => {
    if (paths.length === 0) return;

    let typeIdent = makePathsIdent(sanitizeIdent(ident), true);
    let fixedIdent = sanitizeIdent(ident);

    // Collection can have the same name than a type, and we can also have multiple collection
    // with the same name, but in different scopes
    while (allPaths.includes(typeIdent)) {
      fixedIdent += '_';
      typeIdent = makePathsIdent(fixedIdent, true);
    }

    const varIdent = makePathsIdent(fixedIdent, false);

    let js = `/**
* @typedef {typeof ${varIdent}} ${typeIdent} - All the valid paths for the collection ${ident}.
* To use this type you can do: \`@type {import('path/to/myTokensFile').${typeIdent}}\`
*/\n`;

    allPaths.push(typeIdent);

    js += `${
      moduleExport === 'es6' ? 'export ' : ''
    }const ${varIdent} = /** @type {const} */ (${inspect(paths)}${
      isTypescript ? ' as const' : ''
    });\n`;

    if (moduleExport === 'commonjs') {
      js += `module.exports.${varIdent} = ${varIdent};\n\n`;
    }

    if (isTypescript) {
      js += `export type ${typeIdent} = typeof ${varIdent};\n\n`;
    }

    result += js;
  });

  if (result === '') return '';

  const allPathsCode = `/**
* @typedef {${allPaths.join(' | ')}} ${ALL_PATHS_TYPE_NAME} - All possible paths
*/
${
  isTypescript ? 'export type ' + ALL_PATHS_TYPE_NAME + ' = ' + allPaths.join(' | ') + ';\n\n' : ''
}`;

  const pathsByTypeCode = `/**
* @typedef {typeof ${PATHS_BY_TYPE_VAR_NAME}} ${PATHS_BY_TYPE_TYPE_NAME} - All the paths for a given token type. Needed for \`${GET_TOKENS_BY_TYPE_FN_NAME}\`
*/
const ${PATHS_BY_TYPE_VAR_NAME} = /** @type {const} */ ({\n${typeIdents
    .map(ident => `  ${ident}: ${makePathsIdent(ident)},\n`)
    .join('')}})${isTypescript ? ' as const' : ''};${
    isTypescript
      ? '\ntype ' + PATHS_BY_TYPE_TYPE_NAME + ' = typeof ' + PATHS_BY_TYPE_VAR_NAME + ';'
      : ''
  }`;

  return result + allPathsCode + pathsByTypeCode;
}

function makeModesIdent(baseIdent: string, isType: boolean) {
  return `${isType ? capitalize(baseIdent) : baseIdent}Modes`;
}

function generateModes({
  ident,
  modes,
  isTypescript,
  moduleExport,
}: commonParams & {
  modes: Array<string>;
}) {
  if (modes.length === 0) return '';

  const varIdent = makeModesIdent(ident, false);
  const typeIdent = makeModesIdent(ident, true);

  const code = `/**
* @typedef {typeof ${varIdent}[number]} ${typeIdent} - All the valid modes of ${ident}.
* To use this type you can do: \`@type {import('path/to/myTokensFile').${typeIdent}}\`
*/
${moduleExport === 'es6' ? 'export ' : ''}const ${makeModesIdent(
    ident,
    false,
  )} = /** @type {const} */ (${inspect(modes)})${isTypescript ? ' as const' : ''};${
    moduleExport === 'commonjs' ? '\nmodule.exports.' + varIdent + ' = ' + varIdent + ';' : ''
  }`;

  if (!isTypescript) {
    return code;
  }

  return code + `\nexport type ${typeIdent} = typeof ${varIdent}[number];`;
}

export function generateAllModes({
  ast,
  isTypescript,
  moduleExport,
}: Omit<commonParams, 'ident'> & {
  ast: Ast;
}) {
  let result = '';
  let allModeUnion: Array<string> = [];

  Object.entries(ast.dataByType).forEach(([ident, astData]) => {
    const output = generateModes({
      ident,
      modes: Array.from(astData.metadata.modes),
      isTypescript,
      moduleExport,
    });

    if (output === '') return;

    allModeUnion.push(makeModesIdent(ident, true));

    result += output + '\n\n';
  });

  Object.entries(ast.collectionModes).forEach(([ident, modes]) => {
    let modeIdent = makeModesIdent(sanitizeIdent(ident), true);
    let fixedIdent = sanitizeIdent(ident);

    // Collection can have the same name than a type, and we can also have multiple collection
    // with the same name, but in different scopes
    while (allModeUnion.includes(modeIdent)) {
      fixedIdent += '_';
      modeIdent = makeModesIdent(fixedIdent, true);
    }

    const output = generateModes({ ident: fixedIdent, modes, isTypescript, moduleExport });

    if (output === '') return;

    allModeUnion.push(modeIdent);

    result += output + '\n\n';
  });

  if (allModeUnion.length === 0) return result;

  const union = allModeUnion.join(' | ');
  return (
    result +
    `/**
* @typedef {${union}} ${ALL_MODES_TYPE_NAME} - All the available modes
*/${isTypescript ? '\nexport type ' + ALL_MODES_TYPE_NAME + ' = ' + union + ';' : ''}`
  );
}

export function generateGetTokenByModeFunction({
  isTypescript,
  moduleExport,
}: Omit<commonParams, 'ident'>) {
  const modeExtends = `${TOKENS_TYPE_NAME}[Path] extends Record<string, any>
    ? keyof ${TOKENS_TYPE_NAME}[Path]
    : undefined`;
  const returnExtends = `${TOKENS_TYPE_NAME}[Path] extends Record<string, any>
    ? ${TOKENS_TYPE_NAME}[Path][Mode extends undefined ? never : Mode]
    : ${TOKENS_TYPE_NAME}[Path]`;

  const jsDoc = `/**
* Retrieve any token for a given mode.
* @template {${ALL_PATHS_TYPE_NAME}} Path - A generic extending all the possible paths
* @template {${modeExtends}} Mode - A generic representing all the valid modes for a given path
* @template {${returnExtends}} Return - The return type
* @param {Path} path - The path to the token
* @param {Mode} mode - The mode of the token you want to retrieve
* @returns {Return} - The value of a token for a given mode
*/\n`;

  const pathError = `"Path: '" + path + "' doesn't exist. Here are all the valid paths:\\n- " + Object.keys(${TOKENS_VAR_NAME}).join('\\n- ')`;

  const modeError = `"Invalid mode '" + mode.toString() + "' at path '" + path + "', here are all the valid modes:\\n- " + Object.keys(${TOKENS_VAR_NAME}[path]).join('\\n- ')`;

  const typescriptGeneric = `<
  Path extends keyof ${TOKENS_TYPE_NAME},
  Mode extends ${modeExtends},
  Return extends ${returnExtends}>`;

  return (
    jsDoc +
    `${moduleExport === 'es6' ? 'export ' : ''}function ${GET_TOKEN_BY_MODE_FN_NAME}${
      isTypescript ? typescriptGeneric : ''
    }(path${isTypescript ? ': Path' : ''}, mode${isTypescript ? ': Mode' : ''})${
      isTypescript ? ': Return' : ''
    } {
  if (!${TOKENS_VAR_NAME}[path]) {
    throw new Error(${pathError})
  }

  if (typeof ${TOKENS_VAR_NAME}[path] !== 'object') {
    return ${TOKENS_VAR_NAME}[path] ${isTypescript ? 'as Return' : ''};
  }

  if (!mode) throw new Error('Mode is undefined but it should be one of ' + Object.keys(${TOKENS_VAR_NAME}[path]).join(', ') + ' for path: ' + path);

  if (!${TOKENS_VAR_NAME}[path][mode]) {
    throw new Error(${modeError})
  }

  return ${TOKENS_VAR_NAME}[path][mode]${isTypescript ? ' as Return' : ''};
}${
      moduleExport === 'commonjs'
        ? '\nmodule.exports.' + GET_TOKEN_BY_MODE_FN_NAME + ' = ' + GET_TOKEN_BY_MODE_FN_NAME + ';'
        : ''
    }`
  );
}

export function generateGetTokenWithModesFunction({
  isTypescript,
  moduleExport,
}: Omit<commonParams, 'ident'>) {
  const jsDoc = `/**
* Retrieve any token with its modes.
* @template {${ALL_PATHS_TYPE_NAME}} Path - A generic extending all the possible paths
* @param {Path} path - The path to the token
* @returns {${TOKENS_TYPE_NAME}[Path]} - The value of a token with its modes
*/\n`;

  const pathError = `"Path: '" + path + "' doesn't exist. Here are all the valid paths:\\n- " + Object.keys(${TOKENS_VAR_NAME}).join('\\n- ')`;

  const typescriptGeneric = `<Path extends keyof ${TOKENS_TYPE_NAME}>`;

  return (
    jsDoc +
    `${moduleExport === 'es6' ? 'export ' : ''}function ${GET_TOKEN_WITH_MODES_FN_NAME}${
      isTypescript ? typescriptGeneric : ''
    }(path${isTypescript ? ': Path' : ''})${
      isTypescript ? ': ' + TOKENS_TYPE_NAME + '[Path]' : ''
    } {
  if (!${TOKENS_VAR_NAME}[path]) {
    throw new Error(${pathError})
  }

  return ${TOKENS_VAR_NAME}[path];
}${
      moduleExport === 'commonjs'
        ? '\nmodule.exports.' +
          GET_TOKEN_WITH_MODES_FN_NAME +
          ' = ' +
          GET_TOKEN_WITH_MODES_FN_NAME +
          ';'
        : ''
    }`
  );
}

export function generateGetTokensByTypeFunction({
  isTypescript,
  moduleExport,
}: Omit<commonParams, 'ident'>) {
  const jsDoc = `/**
* Retrieve all the tokens for a specific type (color, dimension, etc...).
* Note that the value will either be a string or an object if the token has modes
* @template {keyof ${PATHS_BY_TYPE_TYPE_NAME}} Type - A generic extending all the possible types
* @template {${TOKENS_TYPE_NAME}[${PATHS_BY_TYPE_TYPE_NAME}[Type][number]]} Token - A generic representing a union of all the outputs
* @param {Type} type - The path to the token
* @returns {{ [Path in PathsByType[Type][number]]: Tokens[Path] }} - An array with all the values
*/\n`;

  return (
    jsDoc +
    `${moduleExport === 'es6' ? 'export ' : ''}function ${GET_TOKENS_BY_TYPE_FN_NAME}${
      isTypescript ? '<Type extends keyof ' + PATHS_BY_TYPE_TYPE_NAME + '>' : ''
    }(type${isTypescript ? ': Type' : ''}) {
  if (!${PATHS_BY_TYPE_VAR_NAME}[type]) {
    throw new Error('The type: \\'' + type + '\\' does not exist')
  }

  return pathsByType[type].reduce(
    (acc, path) => {
      // @ts-expect-error - Can't cast \`path\` to \`Path\`
      acc[path] = ${TOKENS_VAR_NAME}[path];
      return acc;
    },
    {}${
      isTypescript
        ? ' as { [Path in ' +
          PATHS_BY_TYPE_TYPE_NAME +
          '[Type][number]]: ' +
          TOKENS_TYPE_NAME +
          '[Path] }'
        : ''
    },
  );
}${
      moduleExport === 'commonjs'
        ? '\nmodule.exports.' + GET_TOKENS_BY_TYPE_FN_NAME + ' = ' + GET_TOKEN_BY_MODE_FN_NAME
        : ''
    }`
  );
}

export function generateWithModeFunction({
  isTypescript,
  moduleExport,
}: Omit<commonParams, 'ident'> & { ast: Ast }) {
  return `/**
 * @typedef {T extends T ? keyof T : never} KeysOfUnion<T>
 * @template T
 */
${isTypescript ? 'type KeysOfUnion<T> = T extends T ? keyof T : never;\n' : ''}
/**
 * @typedef {T[keyof T]} IndexSelf<T>
 * @template T
 */
${isTypescript ? 'type IndexSelf<T> = T[keyof T];\n' : ''}
/**
 * @typedef {IndexSelf<{ [Path in keyof Tokens]: Tokens[Path] extends { [key in Mode]: any } ? Path : never; }>} ValidPathsFromMode<Mode>
 * @template {string} Mode
 */
${
  isTypescript
    ? 'export type ValidPathsFromMode<Mode extends string> = IndexSelf<{ [Path in keyof Tokens]: Tokens[Path] extends { [key in Mode]: any } ? Path : never; }>;\n'
    : ''
}
/**
 * @template {KeysOfUnion<Tokens[keyof Tokens]>} Mode
 * @param {Mode} mode - Any valid mode
 * @returns - A function that takes a token path which has the given mode
 */
export function ${WITH_MODE_FN_NAME}${
    isTypescript ? '<Mode extends KeysOfUnion<Tokens[keyof Tokens]>>' : ''
  }(mode${isTypescript ? ': Mode' : ''}) {
  /**
   * @template {ValidPathsFromMode<Mode>} Path
   * @template {Extract<keyof Tokens[Path], Mode>} ValidMode
   * @param {Path} path - A valid path for the given mode
   * @returns {Tokens[Path][ValidMode]}
   */
  return ${
    isTypescript
      ? '<Path extends ValidPathsFromMode<Mode>, ValidMode extends Extract<keyof Tokens[Path], Mode>>'
      : ''
  }(
    path${isTypescript ? ': Path' : ''}
  ) => {
    if (!${TOKENS_VAR_NAME}[path]) {
      throw new Error("Invalid path: '" + path + "'")
    }

    if (!${TOKENS_VAR_NAME}[path][mode${isTypescript ? ' as unknown as ValidMode' : ''}]) {
      throw new Error("Invalid mode: '" + mode + "' for path: '" + path + "'")
    }

    return ${TOKENS_VAR_NAME}[path][mode${isTypescript ? ' as unknown as ValidMode' : ''}];
  }
}${
    moduleExport === 'commonjs'
      ? '\nmodule.exports.' + WITH_MODE_FN_NAME + ' = ' + WITH_MODE_FN_NAME + ';'
      : ''
  }
`;
}

function primitiveOfValue(
  { value, isCode }: Value,
  mode?: string,
): [Primitive, Array<string | undefined>] {
  if (!Array.isArray(value) && typeof value === 'object') {
    const toUnquoteTokens: Array<string> = [];
    return [
      Object.entries(value).reduce((acc, [key, value]) => {
        acc[key] = value.value;

        if (value.isCode) {
          toUnquoteTokens.push(key);
        }

        return acc;
      }, {} as { [key: string]: string | number | Array<string | number> }),
      toUnquoteTokens,
    ];
  } else {
    // undefined mode is in the case of tokens with only 1 mode
    // We want to know that we have to unquote the value, but we don't need a key to do so
    return [value, isCode ? [mode] : []];
  }
}

function getPrimitiveValueOfAstValue(
  valueWithMode: ValueWithMode,
): [jsValue: Primitive | PrimitiveWithMode, toUnquoteKeys: Array<string | undefined>] {
  const modes = Object.keys(valueWithMode);

  let toUnquoteKeys: Array<string | undefined> = [];

  return [
    modes.reduce((acc, mode) => {
      const [value, keys] = primitiveOfValue(valueWithMode[mode], mode);
      toUnquoteKeys = toUnquoteKeys.concat(keys);

      acc[mode] = value;

      return acc;
    }, {} as PrimitiveWithMode),
    toUnquoteKeys,
  ];
}

export function generateAllParts({
  ast,
  isTypescript,
  moduleExport,
}: Omit<commonParams, 'ident'> & { ast: Ast }) {
  let result = '';

  result += generateAllPaths({
    ast,
    isTypescript,
    moduleExport,
  });
  result += '\n\n';

  result += generateAllModes({
    ast,
    isTypescript,
    moduleExport,
  });
  result += '\n\n';

  const tokensToUnquote: Array<{ path: string; key?: string }> = [];

  result += generateTokens({
    tokens: Object.values(ast.dataByType).reduce((acc, astData) => {
      Array.from(astData.tokens.entries()).forEach(([path, value]) => {
        const [jsValue, toUnquoteKeys] = getPrimitiveValueOfAstValue(value);
        acc[path] = jsValue;

        toUnquoteKeys.forEach(key => {
          tokensToUnquote.push({ path, key });
        });
      });

      return acc;
    }, {} as { [path: string]: PrimitiveWithMode | Primitive }),
    tokensToUnquote,
    isTypescript,
    moduleExport,
  });
  result += '\n\n';

  result += generateGetTokenByModeFunction({
    isTypescript,
    moduleExport,
  });
  result += '\n\n';

  result += generateGetTokenWithModesFunction({
    isTypescript,
    moduleExport,
  });
  result += '\n\n';

  result += generateGetTokensByTypeFunction({
    isTypescript,
    moduleExport,
  });
  result += '\n\n';

  result += generateWithModeFunction({
    ast,
    isTypescript,
    moduleExport,
  });
  result += '\n\n';

  return result;
}

type options = {
  isTypescript: boolean;
  exportStyle: commonParams['moduleExport'];
  codeToPrepend?: string;
  codeToAppend?: string;
};

export function generateCode(
  ast: Ast,
  { isTypescript, exportStyle, codeToAppend, codeToPrepend }: options,
): string {
  const moduleExport = isTypescript ? 'es6' : exportStyle;

  let result = !!codeToPrepend ? codeToPrepend + '\n\n' : '';

  result += moduleExport === 'commonjs' ? 'module.exports = {};\n\n' : '';

  result += generateAllParts({ ast, isTypescript, moduleExport });

  result += !!codeToAppend ? '\n\n' + codeToAppend : '';

  return result;
}