#!/usr/bin/env node

const fs = require("fs");
const path = require("path");
const parser = require("@babel/parser");
const traverse = require("@babel/traverse").default;
const recast = require("recast");
const glob = require("glob");
const { v4: uuidv4 } = require("uuid");

// List of standard HTML elements
const HTML_ELEMENTS = new Set([
  "div",
  "span",
  "p",
  "h1",
  "h2",
  "h3",
  "h4",
  "h5",
  "h6",
  "section",
  "article",
  "nav",
  "header",
  "footer",
  "main",
  "aside",
  "form",
  "input",
  "button",
  "label",
  "select",
  "option",
  "textarea",
  "a",
  "img",
  "ul",
  "ol",
  "li",
  "table",
  "tr",
  "td",
  "th",
  "thead",
  "tbody",
  "tfoot",
  "dl",
  "dt",
  "dd",
  "figure",
  "figcaption",
  "blockquote",
  "pre",
  "code",
  "strong",
  "em",
  "mark",
  "small",
  "sub",
  "sup",
  "time",
  "video",
  "audio",
  "canvas",
  "svg",
  "path",
  "circle",
  "rect",
  "line",
  "polygon",
  "text",
  "iframe",
  "object",
  "embed",
  "param",
  "source",
  "track",
  "map",
  "area",
  "col",
  "colgroup",
  "caption",
  "thead",
  "tbody",
  "tfoot",
  "th",
  "td",
  "tr",
  "br",
  "hr",
  "meta",
  "link",
  "style",
  "script",
  "noscript",
  "title",
  "head",
  "body",
  "html",
]);

// Helper function to detect if an element is inside a loop
function isInsideLoop(path: any): boolean {
  let currentPath = path;

  while (currentPath) {
    // Check if we're inside a JSX expression that contains array iteration methods
    if (currentPath.isJSXExpressionContainer()) {
      const expression = currentPath.node.expression;

      // Check if the expression is a call expression (like array.map())
      if (expression && expression.type === "CallExpression") {
        const { callee } = expression;

        // Check for method calls like .map(), .forEach(), .filter().map(), etc.
        if (callee.type === "MemberExpression") {
          const propertyName = callee.property?.name;
          if (propertyName === "map" || propertyName === "forEach") {
            return true;
          }
        }
      }

      // Check for chained methods ending with map (e.g., array.filter().map())
      if (expression && expression.type === "CallExpression") {
        let current = expression;
        while (current && current.type === "CallExpression") {
          if (current.callee?.type === "MemberExpression") {
            const propertyName = current.callee.property?.name;
            if (propertyName === "map" || propertyName === "forEach") {
              return true;
            }
          }
          current = current.callee?.object;
        }
      }
    }

    // Check if we're inside a for loop, while loop, or do-while loop
    if (
      currentPath.isForStatement() ||
      currentPath.isForInStatement() ||
      currentPath.isForOfStatement() ||
      currentPath.isWhileStatement() ||
      currentPath.isDoWhileStatement()
    ) {
      return true;
    }

    currentPath = currentPath.parentPath;
  }

  return false;
}

// Simplified function to check if element should get auto-id-list prefix
function shouldUseListPrefix(path: any): boolean {
  let currentPath = path;
  let jsxDepth = 0;

  while (currentPath) {
    // Count JSX elements we traverse
    if (currentPath.isJSXElement() || currentPath.isJSXFragment()) {
      jsxDepth++;
    }

    // Stop counting after the first JSX element (the element itself)
    if (jsxDepth > 1) {
      return false; // We're nested inside another JSX element
    }

    // Check for direct map/forEach in JSX expression
    if (currentPath.isJSXExpressionContainer()) {
      const expression = currentPath.node.expression;
      if (
        expression?.type === "CallExpression" &&
        expression.callee?.type === "MemberExpression"
      ) {
        const methodName = expression.callee.property?.name;
        if (methodName === "map" || methodName === "forEach") {
          return true; // Direct child of map/forEach
        }
      }

      // Check for ternary operators that contain map/forEach calls
      if (expression?.type === "ConditionalExpression") {
        // Check the consequent (true branch) for map/forEach
        const consequent = expression.consequent;
        if (
          consequent?.type === "CallExpression" &&
          consequent.callee?.type === "MemberExpression"
        ) {
          const methodName = consequent.callee.property?.name;
          if (methodName === "map" || methodName === "forEach") {
            return true; // Direct child of map/forEach in ternary
          }
        }

        // Check for nested ternary operators
        if (consequent?.type === "ConditionalExpression") {
          // Recursively check nested ternary
          const nestedConsequent = consequent.consequent;
          if (
            nestedConsequent?.type === "CallExpression" &&
            nestedConsequent.callee?.type === "MemberExpression"
          ) {
            const methodName = nestedConsequent.callee.property?.name;
            if (methodName === "map" || methodName === "forEach") {
              return true; // Direct child of map/forEach in nested ternary
            }
          }
        }
      }

      // Check for logical expressions (&&, ||, ??) that contain map/forEach calls
      if (expression?.type === "LogicalExpression") {
        // Check the right side of logical expressions (&&, ||, ??)
        const right = expression.right;
        if (
          right?.type === "CallExpression" &&
          right.callee?.type === "MemberExpression"
        ) {
          const methodName = right.callee.property?.name;
          if (methodName === "map" || methodName === "forEach") {
            return true; // Direct child of map/forEach in logical expression
          }
        }

        // Check for nested logical expressions
        if (right?.type === "LogicalExpression") {
          // Recursively check nested logical expression
          const nestedRight = right.right;
          if (
            nestedRight?.type === "CallExpression" &&
            nestedRight.callee?.type === "MemberExpression"
          ) {
            const methodName = nestedRight.callee.property?.name;
            if (methodName === "map" || methodName === "forEach") {
              return true; // Direct child of map/forEach in nested logical expression
            }
          }
        }
      }
    }

    // Check for push to array inside forEach
    if (
      currentPath.node?.type === "CallExpression" &&
      currentPath.node.callee?.type === "MemberExpression" &&
      currentPath.node.callee.property?.name === "push"
    ) {
      // Look for forEach in the ancestors
      let pushParent = currentPath.parentPath;
      while (pushParent) {
        if (
          pushParent.node?.type === "ArrowFunctionExpression" ||
          pushParent.node?.type === "FunctionExpression"
        ) {
          // Check if this function is inside a forEach
          let funcParent = pushParent.parentPath;
          while (funcParent) {
            if (
              funcParent.node?.type === "CallExpression" &&
              funcParent.node.callee?.type === "MemberExpression" &&
              funcParent.node.callee.property?.name === "forEach"
            ) {
              return true;
            }
            funcParent = funcParent.parentPath;
          }
        }
        pushParent = pushParent.parentPath;
      }
    }

    // Check for traditional loops
    if (
      currentPath.isForStatement() ||
      currentPath.isForInStatement() ||
      currentPath.isForOfStatement() ||
      currentPath.isWhileStatement() ||
      currentPath.isDoWhileStatement()
    ) {
      return true;
    }

    currentPath = currentPath.parentPath;
  }

  return false;
}

function injectIdsInFile(filePath: string) {
  try {
    const code = fs.readFileSync(filePath, "utf-8");
    const ast = recast.parse(code, {
      parser: {
        parse(source: string) {
          return parser.parse(source, {
            sourceType: "module",
            plugins: [
              "jsx",
              "typescript",
              "classProperties",
              "objectRestSpread",
            ],
            tokens: true,
          });
        },
      },
    });

    let changed = false;
    traverse(ast, {
      JSXOpeningElement(path: any) {
        // Get the element name
        const elementName = path.node.name.name;
        if (!elementName) {
          return; // Skip if no element name
        }

        // Check if this is a standard HTML element or a custom element
        const isHtmlElement =
          !IGNORE_DEFAULTS && HTML_ELEMENTS.has(elementName.toLowerCase());
        const isCustomElement = CUSTOM_ELEMENTS.has(elementName);

        if (!isHtmlElement && !isCustomElement) {
          return; // Skip non-HTML and non-custom elements
        }

        if (REMOVE_TESTIDS) {
          // Remove data-testid attributes
          const testIdIndex = path.node.attributes.findIndex(
            (attr: any) => attr.name && attr.name.name === "data-testid"
          );

          if (testIdIndex !== -1) {
            // Remove the attribute
            path.node.attributes.splice(testIdIndex, 1);
            changed = true;
          }
        } else {
          // Add or update data-testid attributes
          const existingTestIdAttr = path.node.attributes.find(
            (attr: any) => attr.name && attr.name.name === "data-testid"
          );

          // Check if element should get auto-id-list prefix
          const useListPrefix = shouldUseListPrefix(path);
          const expectedPrefix = useListPrefix ? "auto-id-list" : "auto-id";

          if (existingTestIdAttr) {
            // Check if existing data-testid has the wrong prefix
            const currentValue = existingTestIdAttr.value?.value || "";
            const hasAutoIdPrefix =
              currentValue.startsWith("auto-id-list-") ||
              currentValue.startsWith("auto-id-");

            if (hasAutoIdPrefix) {
              const currentPrefix = currentValue.startsWith("auto-id-list-")
                ? "auto-id-list"
                : "auto-id";

              // Only update if the prefix is incorrect
              if (currentPrefix !== expectedPrefix) {
                // Extract the UUID part and replace the prefix
                const uuidPart = currentValue.replace(/^auto-id(-list)?-/, "");
                existingTestIdAttr.value.value = `${expectedPrefix}-${uuidPart}`;
                changed = true;
              }
            }
          } else {
            // No data-testid exists, add one
            const randomId = uuidv4();

            path.node.attributes.push({
              type: "JSXAttribute",
              name: { type: "JSXIdentifier", name: "data-testid" },
              value: {
                type: "StringLiteral",
                value: `${expectedPrefix}-${randomId}`,
              },
            });
            changed = true;
          }
        }
      },
    });

    if (changed) {
      const output = recast.print(ast).code;
      fs.writeFileSync(filePath, output, "utf-8");
      const action = REMOVE_TESTIDS
        ? "Removed testids from"
        : "Added testids to";
      console.log(`${action}: ${filePath}`);
      return true;
    }
    return false;
  } catch (error: any) {
    console.error(`Error processing file ${filePath}:`, error.message);
    return false;
  }
}

async function getStagedFiles() {
  const { execSync } = require("child_process");
  try {
    const output = execSync(
      "git diff --cached --name-only --diff-filter=ACMR",
      { encoding: "utf-8" }
    );
    return output
      .split("\n")
      .filter((file: string) => /\.(js|jsx|ts|tsx)$/.test(file));
  } catch (error: any) {
    console.error("Error getting staged files:", error.message);
    return [];
  }
}

// Custom elements to process alongside HTML elements
// Define CUSTOM_ELEMENTS in module scope
const CUSTOM_ELEMENTS: Set<string> = new Set();

// Flag to determine whether to ignore default HTML elements
let IGNORE_DEFAULTS = false;

// Flag to determine whether to remove data-testid attributes instead of adding them
let REMOVE_TESTIDS = false;

// Array of paths to ignore
let IGNORE_PATHS: string[] = [];

async function run() {
  try {
    let files = [];
    const args = process.argv.slice(2);
    const isStagedOnly = args.includes("--staged");

    // Extract path option
    let searchPath = "**";
    const pathFlagIndex = args.findIndex(
      (arg) => arg === "--path" || arg === "-p"
    );
    if (pathFlagIndex !== -1 && args.length > pathFlagIndex + 1) {
      searchPath = args[pathFlagIndex + 1];
    }

    // Set up default ignore patterns
    const ignorePatterns = ["**/node_modules/**", "**/dist/**", "**/build/**"];

    // Add any programmatically set ignore paths
    ignorePatterns.push(...IGNORE_PATHS);

    // Extract ignore path option
    const ignorePathIndex = args.findIndex(
      (arg) => arg === "--ignore-path" || arg === "-ip"
    );
    if (ignorePathIndex !== -1 && args.length > ignorePathIndex + 1) {
      const ignorePath = args[ignorePathIndex + 1];
      if (ignorePath) {
        // Format the ignore path for glob
        const formattedIgnorePath = ignorePath.endsWith("/")
          ? `${ignorePath}**`
          : `${ignorePath}/**`;
        ignorePatterns.push(formattedIgnorePath);
      }
    }

    // Extract custom elements option
    const customElementsIndex = args.findIndex(
      (arg) => arg === "--custom-elements" || arg === "-c"
    );
    if (customElementsIndex !== -1 && args.length > customElementsIndex + 1) {
      try {
        const customElementsJson = args[customElementsIndex + 1];
        const customElements = JSON.parse(customElementsJson);
        if (Array.isArray(customElements)) {
          // Clear the existing set and add new elements
          CUSTOM_ELEMENTS.clear();
          customElements.forEach((el) => CUSTOM_ELEMENTS.add(el));
        }
      } catch (error) {
        console.error("Error parsing custom elements:", error);
      }
    }

    // Check for --ignore-defaults flag
    const ignoreDefaults =
      args.includes("--ignore-defaults") || args.includes("-i");
    if (ignoreDefaults) {
      // Ensure that custom elements are provided when ignore-defaults is used
      if (customElementsIndex === -1 || CUSTOM_ELEMENTS.size === 0) {
        console.error(
          "Error: --ignore-defaults requires --custom-elements to be provided"
        );
        process.exit(1);
      }
      IGNORE_DEFAULTS = true;
    } else {
      IGNORE_DEFAULTS = false;
    }

    // Check for --remove flag
    const removeTestIds = args.includes("--remove") || args.includes("-r");
    REMOVE_TESTIDS = removeTestIds;

    if (isStagedOnly) {
      files = await getStagedFiles();
      // For staged files, we need to filter manually if ignore paths were provided
      if (ignorePathIndex !== -1) {
        files = files.filter((file: string) => {
          return !ignorePatterns.some((pattern) => {
            // Convert glob pattern to regex pattern
            const regexPattern = pattern
              .replace(/\*\*/g, ".*")
              .replace(/\*/g, "[^/]*")
              .replace(/\?/g, ".");
            return new RegExp(regexPattern).test(file);
          });
        });
      }
    } else {
      // Use the configured search path
      const pattern = searchPath.endsWith("/*.{js,jsx,ts,tsx}")
        ? searchPath
        : `${searchPath}/**/*.{js,jsx,ts,tsx}`;

      files = glob.sync(pattern, {
        absolute: true,
        ignore: ignorePatterns,
      });
    }

    if (files.length === 0) {
      console.log("No files found to process");
      return;
    }

    console.log("Processing files:", files);
    const results = files.map(injectIdsInFile);
    const changedFiles = results.filter(Boolean).length;

    if (changedFiles > 0) {
      const action = REMOVE_TESTIDS ? "removed testids from" : "updated";
      console.log(`\nSuccessfully ${action} ${changedFiles} files`);
      if (isStagedOnly) {
        console.log("Please stage the changes and commit again");
        process.exit(1);
      }
    }
  } catch (error: any) {
    console.error("Error running the script:", error.message);
    process.exit(1);
  }
}

// Only run if called directly (not required as a module)
if (require.main === module) {
  run();
}

// Export functions and provide access to custom elements
module.exports = {
  injectIdsInFile,
  run,
  // Allow programmatic setting of custom elements
  setCustomElements: (elements: string[]) => {
    CUSTOM_ELEMENTS.clear();
    elements.forEach((el) => CUSTOM_ELEMENTS.add(el));
  },
  // Allow programmatic setting of ignoreDefaults
  setIgnoreDefaults: (ignore: boolean) => {
    IGNORE_DEFAULTS = ignore;
  },
  // Allow programmatic setting of paths to ignore
  setIgnorePaths: (paths: string[]) => {
    IGNORE_PATHS = paths.map((p) => {
      // Format the ignore path for glob
      return p.endsWith("/") ? `${p}**` : `${p}/**`;
    });
  },
  // Allow programmatic setting of removeTestIds
  setRemoveTestIds: (remove: boolean) => {
    REMOVE_TESTIDS = remove;
  },
};