import { parse } 				from "@babel/parser";
import * as babelTraverse 		from "@babel/traverse";
import generator from "@babel/generator";
const generate = generator.default;


import Shader 					from "../framework/WebGpu.js";

import { readFile } 			from "fs/promises";
import { writeFile } 			from "fs/promises";

const traverse = babelTraverse.default.default;

async function loadFile( pathName ) {

	const content = await readFile( pathName, "utf-8" );

	return content;

}

async function saveToJsonFile( pathName, object ) {

	const jsonContent = JSON.stringify( object, null, 4 );

	await writeFile( pathName, jsonContent, "utf-8" );

}

function parseBindings( wgsl ) {

	var bindings = [];

	const regex = /@group\((\d+)\)\s+@binding\((\d+)\)\s+var<(\w+)(?:,\s*(\w+))?>\s+(\w+)\s*:\s*([^;]+);/g;

	let match;

	while ( ( match = regex.exec( wgsl ) ) !== null ) {

		const group 		= Number( match[1] );
		const binding 		= Number( match[2] );
		const storageClass 	= match[3];
		const access 		= match[4];
		const varName 		= match[5];

		let type;

		if ( storageClass === "storage" ) {

			type = access === "read" ? "read-only-storage" : "storage";

		} else if ( storageClass === "uniform" ) {

			type = "uniform";

		} else {

			type = "storage";

		}

		const varType = match[6];

		bindings.push({ group, binding, type, varName, varType });

	}

	bindings.sort( ( a, b ) => a.binding - b.binding );

	return bindings;

}

function extractShaderEdges(code) {

	const shaderInstances = new Map();
	const edges = [];
	const ifStatements = [];
	const definitions = [];

	const ast = parse(code, {
		sourceType: "module",
		plugins: ["classProperties"]
	});

	traverse(ast, {

		IfStatement(path) {

			const { test, consequent, alternate, start, end } = path.node;

			ifStatements.push({
				test: code.slice(test.start, test.end),
				consequent: code.slice(consequent.start, consequent.end),
				alternate: alternate ? code.slice(alternate.start, alternate.end) : null,
				start,
				end
			});
		},

		AssignmentExpression(path) {

			const node = path.node;

			if (node.left.type === "MemberExpression") {

				const shaderName = node.left.property.name;

				if (node.right.type === "NewExpression" && node.right.callee.name === "Shader") {

					const args = node.right.arguments;

					if (args.length >= 2 && args[1].type === "StringLiteral") {

						const shaderPath = args[1].value;

						// Store more precise location info for matching
						shaderInstances.set(shaderName, {
							name: shaderName,
							path: shaderPath,
							start: path.parentPath.node.start,
							end: path.parentPath.node.end
						});
					}
				}
			}
		},

		CallExpression(path) {

			const node = path.node;

			if (node.callee.type === "MemberExpression") {

				const object = node.callee.object;
				const method = node.callee.property.name;

				if (object.type === "MemberExpression" && object.object.type === "ThisExpression") {

					const toShaderName = object.property.name;

					if (method === "setBuffer" && node.arguments.length === 2) {

						const bufferName = node.arguments[0].type === "StringLiteral"
							? node.arguments[0].value
							: null;

						const arg = node.arguments[1];

						if (
							arg.type === "CallExpression" &&
							arg.callee.type === "MemberExpression" &&
							arg.callee.property.name === "getBuffer" &&
							arg.callee.object.type === "MemberExpression" &&
							arg.callee.object.object.type === "ThisExpression"
						) {

							const fromShaderName = arg.callee.object.property.name;

							const fromBufferName = arg.arguments.length === 1 && arg.arguments[0].type === "StringLiteral"
								? arg.arguments[0].value
								: null;

							edges.push({
								from: fromShaderName,
								to: toShaderName,
								buffer: bufferName,
								fromBuffer: fromBufferName
							});
						}
					}
				}
			}
		},

		VariableDeclarator(path) {
			const name = path.node.id.name;
			const init = path.node.init ? generate(path.node.init).code : null;
			definitions.push({
				name,
				init,
				start: path.node.start,
				end: path.node.end
			});
		},

		ClassProperty(path) {
			const name = path.node.key.name;
			const init = path.node.value ? generate(path.node.value).code : null;
			definitions.push({
				name,
				init,
				start: path.node.start,
				end: path.node.end
			});
		}

	});

	// Attach ifStatements to shaders they are declared within
	for (const shader of shaderInstances.values()) {

		for (const ifStmt of ifStatements) {

			if (shader.start >= ifStmt.start && shader.end <= ifStmt.end) {

				if (!shader.ifStatements) shader.ifStatements = [];

				shader.ifStatements.push(ifStmt);
			}
		}
	}

	return {
		shaders: Array.from(shaderInstances.values()),
		edges: edges,
		definitions: definitions
	};
}

const code 		= await loadFile("../Graphics/demo.js");

const result 	= extractShaderEdges(code);

for (var i = 0; i < result.shaders.length; i++) {

	var shaderInfo 	= result.shaders[i];

	var source 		= await loadFile(shaderInfo.path);

	var bindings	= parseBindings(source);

	shaderInfo.bindings = bindings;
}

await saveToJsonFile("nodes.json", result);

console.log(result);