// eslint-plugin-no-multiple-returns.ts
import { Rule } from "eslint";
import * as ESTree from "estree";

interface NoMultipleReturnsVisitor extends Rule.RuleListener {
  FunctionDeclaration(node: ESTree.FunctionDeclaration): void;
  FunctionExpression(node: ESTree.FunctionExpression): void;
  ArrowFunctionExpression(node: ESTree.ArrowFunctionExpression): void;
}

function isUseEffectCall(node: ESTree.CallExpression): boolean {
  if (node.callee.type === "Identifier" && node.callee.name === "useEffect") {
    return true;
  }
  return false;
}

function isReactComponent(
  node:
    | ESTree.FunctionExpression
    | ESTree.ArrowFunctionExpression
    | ESTree.FunctionDeclaration
): boolean {
  // Check if the function returns JSX (React.createElement call or JSX syntax)
  // For arrow functions with expression body
  if (
    node.type === "ArrowFunctionExpression" &&
    node.body.type !== "BlockStatement"
  ) {
    return isJSXExpression(node.body);
  }

  // For functions with block statement body, check if they have return statements with JSX
  if (node.body && node.body.type === "BlockStatement") {
    return hasJSXReturn(node.body);
  }

  return false;
}

function isJSXExpression(node: ESTree.Node): boolean {
  if (node.type === "CallExpression") {
    // Check for React.createElement
    if (
      node.callee.type === "MemberExpression" &&
      node.callee.object.type === "Identifier" &&
      node.callee.object.name === "React" &&
      node.callee.property.type === "Identifier" &&
      node.callee.property.name === "createElement"
    ) {
      return true;
    }
  }
  return false;
}

function hasJSXReturn(node: ESTree.BlockStatement): boolean {
  for (const stmt of node.body) {
    if (
      stmt.type === "ReturnStatement" &&
      stmt.argument &&
      isJSXExpression(stmt.argument)
    ) {
      return true;
    }
  }
  return false;
}

function countReturns(node: ESTree.Node): number {
  let count = 0;
  const visited = new Set<ESTree.Node>();
  const stack: Array<ESTree.Node> = [node];

  while (stack.length > 0) {
    const current = stack.pop();

    if (!current || visited.has(current)) {
      continue;
    }

    visited.add(current);

    if (current.type === "ReturnStatement") {
      count += 1;
      continue;
    }

    // Only traverse specific AST child properties to avoid infinite loops
    const childProperties: Array<string> = [];

    switch (current.type) {
      case "BlockStatement":
        childProperties.push("body");
        break;
      case "IfStatement":
        childProperties.push("test", "consequent", "alternate");
        break;
      case "WhileStatement":
      case "DoWhileStatement":
        childProperties.push("test", "body");
        break;
      case "ForStatement":
        childProperties.push("init", "test", "update", "body");
        break;
      case "ForInStatement":
      case "ForOfStatement":
        childProperties.push("left", "right", "body");
        break;
      case "SwitchStatement":
        childProperties.push("discriminant", "cases");
        break;
      case "SwitchCase":
        childProperties.push("test", "consequent");
        break;
      case "TryStatement":
        childProperties.push("block", "handler", "finalizer");
        break;
      case "CatchClause":
        childProperties.push("param", "body");
        break;
      case "ExpressionStatement":
        childProperties.push("expression");
        break;
      case "ConditionalExpression":
        childProperties.push("test", "consequent", "alternate");
        break;
      case "LogicalExpression":
      case "BinaryExpression":
        childProperties.push("left", "right");
        break;
      case "UnaryExpression":
      case "UpdateExpression":
        childProperties.push("argument");
        break;
      case "CallExpression":
      case "NewExpression":
        // Skip traversing into useEffect calls to exclude their cleanup function returns
        if (current.type === "CallExpression" && isUseEffectCall(current)) {
          // Only traverse the callee, not the arguments which contain the effect function
          childProperties.push("callee");
        } else {
          childProperties.push("callee", "arguments");
        }
        break;
      case "MemberExpression":
        childProperties.push("object", "property");
        break;
      case "ArrayExpression":
        childProperties.push("elements");
        break;
      case "ObjectExpression":
        childProperties.push("properties");
        break;
      case "Property":
        childProperties.push("key", "value");
        break;
      case "FunctionExpression":
      case "ArrowFunctionExpression":
        // Skip traversing into React components to exclude their return statements
        if (
          isReactComponent(
            current as
              | ESTree.FunctionExpression
              | ESTree.ArrowFunctionExpression
          )
        ) {
          // Only traverse params, not the body which contains the component's return
          childProperties.push("params");
        } else {
          childProperties.push("params", "body");
        }
        break;
      case "VariableDeclaration":
        childProperties.push("declarations");
        break;
      case "VariableDeclarator":
        childProperties.push("id", "init");
        break;
      case "AssignmentExpression":
        childProperties.push("left", "right");
        break;
      case "SequenceExpression":
        childProperties.push("expressions");
        break;
      // Add more cases as needed
    }

    for (const prop of childProperties) {
      const value = (current as any)[prop];
      if (Array.isArray(value)) {
        for (const item of value) {
          if (item && typeof item === "object" && item.type) {
            stack.push(item);
          }
        }
      } else if (value && typeof value === "object" && value.type) {
        stack.push(value);
      }
    }
  }

  return count;
}

const rule: Rule.RuleModule = {
  meta: {
    type: "suggestion",
    docs: {
      description: "disallow multiple return statements in a function",
      recommended: false,
    },
    schema: [], // no options
  },

  create(context: Rule.RuleContext): NoMultipleReturnsVisitor {
    return {
      FunctionDeclaration(node: ESTree.FunctionDeclaration): void {
        const returnCount = countReturns(node.body);
        if (returnCount > 1) {
          context.report({
            node,
            message: `Function '${
              (node.id && node.id.name) || "<anonymous>"
            }' has ${returnCount} return statements.`,
          });
        }
      },
      FunctionExpression(node: ESTree.FunctionExpression): void {
        const returnCount = countReturns(node.body);
        if (returnCount > 1) {
          context.report({
            node,
            message: `Function expression has ${returnCount} return statements.`,
          });
        }
      },
      ArrowFunctionExpression(node: ESTree.ArrowFunctionExpression): void {
        if (node.body.type === "BlockStatement") {
          const returnCount: number = countReturns(node.body);
          if (returnCount > 1) {
            context.report({
              node,
              message: `Arrow function has ${returnCount} return statements.`,
            });
          }
        }
      },
    };
  },
};

export default rule;