import * as fs from "node:fs";
import * as path from "node:path";

interface FlowStep {
	id: string;
	label: string;
	major: number;
	level: number;
	file: string; // New: Track file path
	line: number; // New: Track line number
}

interface FlowMap {
	[flowId: string]: FlowStep[];
}

/** CONFIGURATION **/
const SRC_DIR = "./src";
const OUTPUT_DIR = "./docs/flows";
const FILE_EXTENSIONS = [".ts", ".js", ".tsx", ".jsx"];

// Updated Regex: Handles //FLOWS and // FLOWS
const FLOW_REGEX = /\/\/\s*FLOWS:\s*([^|]+)\s*\|\s*([\w.]+)\s*\|\s*(.*)/g;

/** UTILS **/
function getFiles(dir: string, allFiles: string[] = []): string[] {
	if (!fs.existsSync(dir)) return allFiles;
	const files = fs.readdirSync(dir);
	for (const file of files) {
		const name = path.join(dir, file);
		if (fs.statSync(name).isDirectory()) getFiles(name, allFiles);
		else if (FILE_EXTENSIONS.includes(path.extname(name))) allFiles.push(name);
	}
	return allFiles;
}

function safeId(id: string): string {
	// Mermaid IDs cannot contain dots
	return `step_${id.replace(/\./g, "_")}`;
}

/** CORE LOGIC **/
function main() {
	const flowData: FlowMap = {};
	const files = getFiles(SRC_DIR);

	console.log(`🔍 Scanning ${files.length} files for // FLOWS comments...`);

	// 1. Extraction
	files.forEach((file) => {
		const content = fs.readFileSync(file, "utf-8");
		const lines = content.split("\n");
		lines.forEach((lineText, index) => {
			let match: RegExpExecArray | null;
			FLOW_REGEX.lastIndex = 0;
			match = FLOW_REGEX.exec(lineText);
			while (match !== null) {
				const [_, flowId, id, label] = match.map((s) => s.trim());
				if (!flowId || !id || !label) {
					console.warn(`⚠️ Invalid FLOWS comment in ${file}: ${match[0]}`);
					continue;
				}
				if (!flowData[flowId]) flowData[flowId] = [];

				flowData[flowId].push({
					id,
					label,
					major: parseInt(id.split(".")?.[0] || "0"),
					level: id.split(".").length,
					file: path.relative(process.cwd(), file), // Store relative path
					line: index + 1, // Store line number (1-based)
				});
				match = FLOW_REGEX.exec(lineText);
			}
		});
	});

	if (Object.keys(flowData).length === 0) {
		console.log(
			"⚠️ No flow comments found. Try adding: // FLOWS: myflow | 1 | Start",
		);
		return;
	}

	// 2. Generate Output
	if (!fs.existsSync(OUTPUT_DIR)) fs.mkdirSync(OUTPUT_DIR, { recursive: true });

	for (const [flowId, steps] of Object.entries(flowData)) {
		// Sort logically (1, 1.1, 1.2, 2)
		const sorted = steps.sort((a, b) =>
			a.id.localeCompare(b.id, undefined, { numeric: true }),
		);
		const mermaid = ["graph TD"];
		const connections = new Set<string>();

		sorted.forEach((current) => {
			const nextSteps = sorted.filter((candidate) => {
				const currentIdPrefix = `${current.id}.`;

				// 1. DESCEND: Link to immediate children (3 -> 3.1, 3 -> 3.2)
				// Check if candidate starts with current ID and is exactly one level deeper
				if (
					candidate.id.startsWith(currentIdPrefix) &&
					candidate.level === current.level + 1
				) {
					// Ensure it's a direct child (e.g., 3 -> 3.1, not 3 -> 3.1.1)
					const remaining = candidate.id.replace(currentIdPrefix, "");
					return !remaining.includes(".");
				}

				// 2. MERGE: Link leaf nodes to the next major integer (3.1.2 -> 4, 3.2.1 -> 4)
				const isLeaf = !sorted.some(
					(s) => s.id.startsWith(currentIdPrefix) && s.id !== current.id,
				);
				if (
					isLeaf &&
					candidate.major === current.major + 1 &&
					candidate.level === 1
				) {
					return true;
				}

				// 3. SEQUENCE (Optional): If you want 3.1.1 -> 3.1.2
				// If they share the exact same parent and are consecutive
				const currentParent = current.id.includes(".")
					? current.id.substring(0, current.id.lastIndexOf("."))
					: null;
				const candParent = candidate.id.includes(".")
					? candidate.id.substring(0, candidate.id.lastIndexOf("."))
					: null;

				if (
					currentParent !== null &&
					currentParent === candParent &&
					candidate.level === current.level
				) {
					// Only link if you want branches to be sequential.
					// In your case, you likely want 3.1.1 and 3.1.2 to both go to 4 independently.
					// So we leave this FALSE to keep them as parallel paths.
					return false;
				}

				return false;
			});

			nextSteps.forEach((next) => {
				const link = `    ${safeId(current.id)}["${current.label}"] --> ${safeId(next.id)}["${next.label}"]`;
				connections.add(link);
			});
		});

		// Handle single-node flows (start node with no connections)
		if (connections.size === 0 && sorted.length > 0 && sorted[0]) {
			mermaid.push(`    ${safeId(sorted[0].id)}["${sorted[0].label}"]`);
		} else {
			connections.forEach((c) => mermaid.push(c));
		}

		const outPath = path.join(OUTPUT_DIR, `${flowId}.md`);

		const header = "```mermaid\n\n";
		const footer = "\n```";
		fs.writeFileSync(outPath, header + mermaid.join("\n") + footer);
		console.log(`✅ Generated: ${outPath}`);
	}
}

main();