#!/usr/bin/env node
// SPDX-License-Identifier: MIT

import fs from "node:fs";
import path from "node:path";
import { parseArgs } from "node:util";
import { Node, Project, SyntaxKind } from "ts-morph";

const projectRoot = process.cwd();
const MAX_EDGES = Number(process.env.MAX_GRAPH_EDGES ?? "450");

const {
  values: { source: sourceArgsRaw, output: outputArg },
} = parseArgs({
  options: {
    source: {
      type: "string",
      multiple: true,
    },
    output: {
      type: "string",
    },
  },
});

function collectFunctionEntries(sourceFile) {
  const entries = [];
  const namesSeen = new Set();

  for (const fn of sourceFile.getFunctions()) {
    const name = fn.getName();
    if (!name || namesSeen.has(name)) {
      continue;
    }
    namesSeen.add(name);
    entries.push({ name, node: fn });
  }

  for (const declaration of sourceFile.getVariableDeclarations()) {
    const name = declaration.getName();
    if (namesSeen.has(name)) {
      continue;
    }

    const initializer = declaration.getInitializer();
    if (!initializer) {
      continue;
    }
    if (!Node.isArrowFunction(initializer) && !Node.isFunctionExpression(initializer)) {
      continue;
    }

    namesSeen.add(name);
    entries.push({ name, node: initializer });
  }

  return entries;
}

function nodeId(filePath, functionName) {
  return `${filePath}__${functionName}`.replace(/[^A-Za-z0-9_]/g, "_");
}

function normalizeRelativePath(filePath) {
  return filePath.replaceAll("\\", "/").replace(/^\.\//, "");
}

function normalizeSourceArg(sourceArg) {
  const projectRelative = path.isAbsolute(sourceArg)
    ? path.relative(projectRoot, sourceArg)
    : sourceArg;
  return normalizeRelativePath(projectRelative).replace(/\/+$/, "");
}

const functionsByFile = new Map();
const allFunctionNames = new Set();
const functionEntriesByFile = new Map();

const project = new Project({
  tsConfigFilePath: path.join(projectRoot, "tsconfig.json"),
});
const sourceFiles = project
  .getSourceFiles()
  .filter((sourceFile) => !sourceFile.isDeclarationFile())
  .sort((a, b) => a.getFilePath().localeCompare(b.getFilePath()));

const sourceContext = sourceFiles.map((sourceFile) => ({
  sourceFile,
  relativePath: normalizeRelativePath(path.relative(projectRoot, sourceFile.getFilePath())),
}));

const filteredSourceFiles = (() => {
  if (!sourceArgsRaw || sourceArgsRaw.length === 0) {
    return sourceContext.map((item) => item.sourceFile);
  }

  const sourceArgs = sourceArgsRaw.map((sourceArg) => normalizeSourceArg(sourceArg));
  const matchedSourceFiles = new Set();

  for (const sourceArg of sourceArgs) {
    const hasMatch = sourceContext.some((item) => (
      item.relativePath === sourceArg
      || item.relativePath.startsWith(`${sourceArg}/`)
    ));

    if (!hasMatch) {
      throw new Error(`No source file matched --source=${sourceArg}`);
    }

    for (const item of sourceContext) {
      if (item.relativePath === sourceArg || item.relativePath.startsWith(`${sourceArg}/`)) {
        matchedSourceFiles.add(item.sourceFile);
      }
    }
  }

  return [...matchedSourceFiles];
})();

for (const sourceFile of filteredSourceFiles) {
  const absolutePath = sourceFile.getFilePath();

  const relativePath = path.relative(projectRoot, absolutePath).replaceAll("\\", "/");
  const functionEntries = collectFunctionEntries(sourceFile);
  const functionNames = functionEntries.map((entry) => entry.name);

  functionsByFile.set(relativePath, functionNames);
  functionEntriesByFile.set(relativePath, functionEntries);
  for (const functionName of functionNames) {
    allFunctionNames.add(functionName);
  }
}

const firstNodeForFunction = new Map();
for (const [relativePath, functionNames] of functionsByFile.entries()) {
  for (const functionName of functionNames) {
    if (!firstNodeForFunction.has(functionName)) {
      firstNodeForFunction.set(functionName, nodeId(relativePath, functionName));
    }
  }
}

const edgeSet = new Set();
for (const functionEntries of functionEntriesByFile.values()) {
  for (const { name: sourceName, node } of functionEntries) {
    const body = node.getBody();
    if (!body) {
      continue;
    }

    const sourceNode = firstNodeForFunction.get(sourceName);
    if (!sourceNode) {
      continue;
    }

    for (const call of body.getDescendantsOfKind(SyntaxKind.CallExpression)) {
      const expression = call.getExpression();
      let targetName;

      if (Node.isIdentifier(expression)) {
        targetName = expression.getText();
      } else if (Node.isPropertyAccessExpression(expression)) {
        targetName = expression.getName();
      }

      if (!targetName || targetName === sourceName || !allFunctionNames.has(targetName)) {
        continue;
      }

      const targetNode = firstNodeForFunction.get(targetName);
      if (targetNode) {
        edgeSet.add(`${sourceNode} --> ${targetNode}`);
      }
    }
  }
}

let diagram = "%% Function Dependency Graph (`src/**/*.ts`)\n";
diagram += "%% Generated from current package source files.\n";
diagram += "flowchart LR\n";

for (const [relativePath, functionNames] of functionsByFile.entries()) {
  if (functionNames.length === 0) {
    continue;
  }
  const subgraphId = relativePath.replace(/[^A-Za-z0-9_]/g, "_");
  diagram += `  subgraph ${subgraphId}[\"${relativePath}\"]\n`;
  for (const functionName of functionNames) {
    diagram += `    ${nodeId(relativePath, functionName)}[\"${functionName}\"]\n`;
  }
  diagram += "  end\n";
}

const sortedEdges = [...edgeSet].sort();
const emittedEdges = sortedEdges.slice(0, MAX_EDGES);
for (const edge of emittedEdges) {
  diagram += `  ${edge}\n`;
}

if (sortedEdges.length > MAX_EDGES) {
  diagram += `%% Note: showing ${MAX_EDGES} of ${sortedEdges.length} detected edges to stay within Mermaid default edge limits.\n`;
}

if (outputArg) {
  const outputPath = path.resolve(projectRoot, outputArg);
  fs.mkdirSync(path.dirname(outputPath), { recursive: true });
  fs.writeFileSync(outputPath, diagram, "utf8");
} else {
  process.stdout.write(diagram);
}