import type { Node, PropDeclaration, RawAttribute } from "./types.ts";

export interface CompileOptions {
	/** Absolute path of the template, used to resolve frontmatter imports. */
	path: string;
	/** Directory of the template, used to resolve frontmatter imports. */
	dir: string;
	/** Frontmatter of the template. */
	frontmatter?: string;
}

export interface CompiledSource {
	/** Source of the module creating the compiled template. */
	source: string;
	/** Prop declarations found in the template. */
	props: PropDeclaration[];
	/** Names declared by the frontmatter. */
	frontmatterNames: string[];
}

const RESERVED = new Set([
	"await",
	"async",
	"break",
	"case",
	"catch",
	"class",
	"const",
	"continue",
	"debugger",
	"default",
	"delete",
	"do",
	"else",
	"enum",
	"export",
	"extends",
	"false",
	"finally",
	"for",
	"from",
	"function",
	"get",
	"if",
	"implements",
	"import",
	"in",
	"instanceof",
	"interface",
	"let",
	"new",
	"null",
	"of",
	"package",
	"private",
	"protected",
	"public",
	"return",
	"set",
	"static",
	"super",
	"switch",
	"this",
	"throw",
	"true",
	"try",
	"typeof",
	"undefined",
	"var",
	"void",
	"while",
	"with",
	"yield",
]);

/**
 * Identifiers used by an expression.
 */
export function collectIdentifiers(code: string): string[] {
	const withoutStrings = code
		.replace(/"(?:[^"\\]|\\.)*"/g, '""')
		.replace(/'(?:[^'\\]|\\.)*'/g, "''")
		.replace(/`(?:[^`\\]|\\.)*`/g, "``");
	const names = new Set<string>();
	const pattern = /(?<![.\w$])([A-Za-z_$][\w$]*)/g;
	let match = pattern.exec(withoutStrings);
	while (match) {
		const name = match[1] as string;
		if (!RESERVED.has(name)) names.add(name);
		match = pattern.exec(withoutStrings);
	}
	return [...names];
}

/**
 * Names declared by the frontmatter.
 */
export function extractDeclarations(code: string): string[] {
	const names: string[] = [];
	const pattern = /^\s*(?:const|let|var)\s+([A-Za-z_$][\w$]*)\s*=/gm;
	let match = pattern.exec(code);
	while (match) {
		names.push(match[1] as string);
		match = pattern.exec(code);
	}
	return names;
}

const VOID_ELEMENTS = new Set([
	"area",
	"base",
	"br",
	"col",
	"embed",
	"hr",
	"img",
	"input",
	"link",
	"meta",
	"param",
	"source",
	"track",
	"wbr",
]);

/** Escapes a string for use inside a JavaScript string literal. */
function jsString(value: string): string {
	return JSON.stringify(value);
}

/** Escapes HTML entities of static text. */
export function escapeHtmlText(value: string): string {
	return value
		.replace(/&/g, "&amp;")
		.replace(/</g, "&lt;")
		.replace(/>/g, "&gt;");
}

/** Resolves an import specifier of the frontmatter. */
function resolveSpecifier(specifier: string, dir: string): string {
	if (!specifier.startsWith(".")) return specifier;
	return `${dir}/${specifier}`.replace(/\/{2,}/g, "/");
}

/**
 * Rewrites the `import` statements of the frontmatter into dynamic imports,
 * because the frontmatter is compiled into a function body.
 */
export function transformFrontmatter(code: string, dir: string): string {
	const lines = code.split("\n");
	const output: string[] = [];

	for (const line of lines) {
		const named = /^\s*import\s+([\s\S]+?)\s+from\s+['"]([^'"]+)['"]\s*$/.exec(
			line,
		);
		if (named) {
			const clause = (named[1] as string).trim();
			const specifier = resolveSpecifier(named[2] as string, dir);
			const defaultMatch = /^([A-Za-z_$][\w$]*)\s*,\s*([\s\S]+)$/.exec(clause);
			if (defaultMatch) {
				const alias = defaultMatch[1] as string;
				const rest = (defaultMatch[2] as string).replace(/^\{|\}$/g, "").trim();
				output.push(`const __module = await import(${jsString(specifier)});`);
				output.push(`const ${alias} = __module.default;`);
				output.push(`const { ${rest} } = __module;`);
				continue;
			}
			if (/^\*\s*as\s+/.test(clause)) {
				const alias = clause.replace(/^\*\s*as\s+/, "").trim();
				output.push(`const ${alias} = await import(${jsString(specifier)});`);
				continue;
			}
			if (/^\{[\s\S]*\}$/.test(clause)) {
				const names = clause.replace(/^\{|\}$/g, "").trim();
				output.push(
					`const { ${names} } = await import(${jsString(specifier)});`,
				);
				continue;
			}
			if (/^[A-Za-z_$][\w$]*$/.test(clause)) {
				output.push(
					`const ${clause} = (await import(${jsString(specifier)})).default;`,
				);
				continue;
			}
			output.push(`await import(${jsString(specifier)});`);
			continue;
		}

		const sideEffect = /^\s*import\s+['"]([^'"]+)['"]\s*$/.exec(line);
		if (sideEffect) {
			output.push(
				`await import(${jsString(resolveSpecifier(sideEffect[1] as string, dir))});`,
			);
			continue;
		}

		output.push(line);
	}

	return output.join("\n");
}

/** Collects every mixin definition of a template. */
export function collectMixinNodes(
	nodes: Node[],
): Extract<Node, { kind: "mixin" }>[] {
	const result: Extract<Node, { kind: "mixin" }>[] = [];
	for (const node of nodes) {
		if (node.kind === "mixin") {
			result.push(node);
			continue;
		}
		const children = childrenOf(node);
		if (children) result.push(...collectMixinNodes(children));
	}
	return result;
}

/** Collects every named block of a template. */
export function collectBlockNodes(nodes: Node[]): {
	name: string;
	mode: "replace" | "append" | "prepend";
	children: Node[];
}[] {
	const result: {
		name: string;
		mode: "replace" | "append" | "prepend";
		children: Node[];
	}[] = [];
	for (const node of nodes) {
		if (node.kind === "block") {
			result.push({
				name: node.name,
				mode: node.mode,
				children: node.children,
			});
			continue;
		}
		const children = childrenOf(node);
		if (children) result.push(...collectBlockNodes(children));
	}
	return result;
}

class Compiler {
	private readonly buffer: string[] = [];

	compile(nodes: Node[]): string {
		this.emitBody(nodes);
		return this.buffer.join("\n");
	}

	private emit(line: string): void {
		this.buffer.push(line);
	}

	private emitBody(nodes: Node[]): void {
		for (const node of nodes) {
			this.emitNode(node);
		}
	}

	private emitNode(node: Node): void {
		switch (node.kind) {
			case "element":
				this.emitElement(node);
				break;
			case "text":
				this.emitText(node.parts);
				break;
			case "rawtext":
				this.emit(`__push(${jsString(escapeHtmlText(node.value))});`);
				break;
			case "code":
				this.emit(node.code);
				break;
			case "out":
				this.emit(
					`__push(__helpers.render(await (${node.code}), ${node.escaped !== false}));`,
				);
				break;
			case "comment":
				if (!node.silent) {
					this.emit(`__push(${jsString(`<!--${node.value}-->`)});`);
				}
				break;
			case "doctype":
				this.emit(`__push(${jsString(`<!DOCTYPE ${node.value || "html"}>`)});`);
				break;
			case "if":
				this.emitIf(node.branches);
				break;
			case "each":
				this.emitEach(node);
				break;
			case "while":
				this.emit(`while (${node.test}) {`);
				this.emitBody(node.children);
				this.emit("}");
				break;
			case "case":
				this.emitCase(node);
				break;
			case "mixin":
				break;
			case "mixin-call":
				this.emit(
					`__push(await __mixins[${jsString(node.name)}](__runtime, __ctx, ${node.args || ""}));`,
				);
				break;
			case "include":
				this.emitBody(node.children);
				break;
			case "extends":
				break;
			case "block":
			case "slot":
				this.emitBlock(node.name, node.children);
				break;
			case "yield":
				this.emitBlock(node.name, []);
				break;
			case "portal":
				this.emit(`__push(${jsString(`<div data-muddle-portal="${node.name}" style="view-transition-name:${node.name}">`)});`);
				this.emitBody(node.children);
				this.emit(`__push("</div>");`);
				break;
			case "props":
			case "test":
				break;
			case "component":
				this.emitComponent(node.name, node.props, node.renders, false);
				break;
			case "island":
				this.emitComponent(node.name, node.props, node.renders, true);
				break;
		}
	}

	private emitElement(node: Extract<Node, { kind: "element" }>): void {
		const anchor = node.anchor
			? node.anchor.expression
				? ` + __helpers.attr("data-muddle-anchor", ${jsString(node.anchor.name)} + String(${node.anchor.expression}))`
				: ` + __helpers.attr("data-muddle-anchor", ${jsString(node.anchor.name)})`
			: "";
		const attributes = this.emitAttributes(node.attributes) + anchor;
		this.emit(
			`__push(${jsString(`<${node.name}`)}${attributes} + ${jsString(">")});`,
		);

		if (VOID_ELEMENTS.has(node.name)) return;

		this.emitBody(node.children);
		this.emit(`__push(${jsString(`</${node.name}>`)});`);
	}

	/**
	 * Renders the attributes of a tag.
	 */
	private emitAttributes(attributes: RawAttribute[]): string {
		const classes: RawAttribute[] = [];
		const others: RawAttribute[] = [];
		for (const attribute of attributes) {
			if (attribute.name === "class") classes.push(attribute);
			else others.push(attribute);
		}

		let rendered = others
			.map((attribute) => this.emitAttribute(attribute))
			.join("");

		if (classes.length === 1) {
			rendered += this.emitAttribute(classes[0] as RawAttribute);
			return rendered;
		}
		if (classes.length > 1) {
			const values = classes
				.map((attribute) =>
					attribute.kind === "string"
						? jsString(attribute.value ?? "")
						: `(${attribute.value ?? ""})`,
				)
				.join(", ");
			rendered += ` + __helpers.attr("class", __helpers.classes([${values}]))`;
		}
		return rendered;
	}

	private emitAttribute(attribute: RawAttribute): string {
		if (attribute.kind === "boolean") {
			return ` + __helpers.attrBoolean(${jsString(attribute.name)})`;
		}
		if (attribute.kind === "string") {
			return ` + __helpers.attr(${jsString(attribute.name)}, ${jsString(attribute.value ?? "")})`;
		}
		if (attribute.name === "style") {
			return ` + __helpers.attr("style", __helpers.style(${attribute.value ?? ""}))`;
		}
		return ` + __helpers.attr(${jsString(attribute.name)}, (${attribute.value ?? ""}))`;
	}

	private emitText(
		parts: {
			type: "literal" | "expression";
			value?: string;
			code?: string;
			escaped?: boolean;
		}[],
	): void {
		for (const part of parts) {
			if (part.type === "literal") {
				this.emit(`__push(${jsString(escapeHtmlText(part.value ?? ""))});`);
				continue;
			}
			const code = part.code ?? "";
			this.emit(
				`__push(__helpers.render(await (${code}), ${part.escaped !== false}));`,
			);
		}
	}

	private emitIf(
		branches: { test: string | null; negate: boolean; children: Node[] }[],
	): void {
		branches.forEach((branch, index) => {
			if (index > 0) {
				this.emit("}");
			}
			if (branch.test === null) {
				this.emit("else {");
			} else {
				const test = branch.negate ? `!(${branch.test})` : branch.test;
				this.emit(`${index === 0 ? "if" : "else if"} (${test}) {`);
			}
			this.emitBody(branch.children);
		});
		this.emit("}");
	}

	private emitEach(node: Extract<Node, { kind: "each" }>): void {
		this.emit("{");
		this.emit(`const __list = (${node.list});`);
		this.emit("if (Array.isArray(__list)) {");
		this.emit("for (let __index = 0; __index < __list.length; __index++) {");
		this.emit(`const ${node.value} = __list[__index];`);
		if (node.key) {
			this.emit(`const ${node.key} = __index;`);
		}
		this.emitBody(node.children);
		this.emit("}");
		this.emit("} else {");
		this.emit("for (const __key of Object.keys(__list ?? {})) {");
		this.emit(`const ${node.value} = __list[__key];`);
		if (node.key) {
			this.emit(`const ${node.key} = __key;`);
		}
		this.emitBody(node.children);
		this.emit("}");
		this.emit("}");
		this.emit("}");
	}

	private emitCase(node: Extract<Node, { kind: "case" }>): void {
		this.emit(`switch (${node.subject}) {`);
		for (const entry of node.cases) {
			if (entry.test === undefined) {
				this.emit("default: {");
			} else {
				this.emit(`case (${entry.test}): {`);
			}
			this.emitBody(entry.children);
			this.emit("break;");
			this.emit("}");
		}
		this.emit("}");
	}

	private emitBlock(name: string, children: Node[]): void {
		this.emit(
			`__push(await __runtime.getBlock(${jsString(name)}, async (__runtime, __ctx) => {`,
		);
		this.emit("const __out = [];");
		this.emit("const __push = (__value) => { __out.push(__value); };");
		this.emit("with (__ctx) {");
		this.emitBody(children);
		this.emit("}");
		this.emit('return __out.join("");');
		this.emit("}));");
	}

	private emitComponent(
		name: string,
		props: Record<string, string>,
		renders: { name: string; children: Node[] }[] = [],
		island = false,
	): void {
		const entries = Object.entries(props).map(
			([key, value]) => `${jsString(key)}: (${value || "undefined"})`,
		);

		for (const render of renders) {
			const renderer = new Compiler();
			const body = renderer.compile(render.children);
			entries.push(
				`${jsString(render.name)}: async () => { const __out = []; const __push = (__value) => { __out.push(__value); }; const __helpers = __runtime.helpers; with (__ctx) { ${body} } return __out.join(""); }`,
			);
		}

		const method = island ? "renderIsland" : "renderComponent";
		this.emit(
			`__push(await __runtime.${method}(${jsString(name)}, { ${entries.join(", ")} }));`,
		);
	}
}


export function compileTemplate(
	nodes: Node[],
	options: CompileOptions,
): CompiledSource {
	const bodyCompiler = new Compiler();
	const body = bodyCompiler.compile(nodes);
	const uses = collectUses(nodes);
	const tests = collectTests(nodes);

	const blockSource: string[] = [];
	for (const block of collectBlockNodes(nodes)) {
		const blockCompiler = new Compiler();
		const blockBody = blockCompiler.compile(block.children);
		blockSource.push(
			`${jsString(block.name)}: { mode: ${jsString(block.mode)}, fn: async (__runtime, __ctx) => { const __out = []; const __push = (__value) => { __out.push(__value); }; const __helpers = __runtime.helpers; with (__ctx) { ${blockBody} } return __out.join(""); } }`,
		);
	}

	const mixinSource: string[] = [];
	for (const mixin of collectMixinNodes(nodes)) {
		const mixinCompiler = new Compiler();
		const mixinBody = mixinCompiler.compile(mixin.children);
		const params = mixin.params.length > 0 ? mixin.params : "";
		mixinSource.push(
			`${jsString(mixin.name)}: async (__runtime, __ctx, ${params}) => { const __out = []; const __push = (__value) => { __out.push(__value); }; const __helpers = __runtime.helpers; with (__ctx) { ${mixinBody} } return __out.join(""); }`,
		);
	}

	const frontmatter =
		options.frontmatter && options.frontmatter.trim().length > 0
			? transformFrontmatter(options.frontmatter, options.dir)
			: "";
	const frontmatterNames =
		frontmatter.length > 0 ? extractDeclarations(frontmatter) : [];
	const frontmatterAssignments = frontmatterNames
		.map((name) => `    __values.${name} = ${name};`)
		.join("\n");

	const source = `
const __blocks = { ${blockSource.join(", ")} };
const __mixins = { ${mixinSource.join(", ")} };
const __frontmatter = async (__runtime, __ctx) => {
  const __values = {};
  const __helpers = __runtime.helpers;
  const __muddle = __runtime.file;
  with (__ctx) {
${frontmatter}
${frontmatterAssignments}
  }
  return __values;
};
const __template = async (__runtime, __ctx) => {
  const __out = [];
  const __push = (__value) => { __out.push(__value); };
  const __helpers = __runtime.helpers;
  const __muddle = __runtime.file;
  with (__ctx) {
    ${body}
  }
  return __out.join("");
};
return { template: __template, frontmatter: __frontmatter, blocks: __blocks, mixins: __mixins, uses: ${JSON.stringify(uses)}, tests: ${JSON.stringify(tests)} };
`;

	return { source, props: [], frontmatterNames };
}

/**
 * The assertions of a template.
 */
export function collectTests(nodes: Node[]): string[] {
	const tests: string[] = [];
	const visit = (nodes: Node[]): void => {
		for (const node of nodes) {
			if (node.kind === "test") {
				tests.push(node.code);
				continue;
			}
			const children = childrenOf(node);
			if (children) visit(children);
		}
	};
	visit(nodes);
	return tests;
}

/**
 * Every identifier used by a template.
 */
export function collectUses(nodes: Node[]): string[] {
	const names = new Set<string>();
	const declared = new Set<string>();

	const visit = (node: Node): void => {
		if (node.kind === "mixin") {
			for (const name of mixinParamNames(node.params)) declared.add(name);
		}
		switch (node.kind) {
			case "out":
				for (const name of collectIdentifiers(node.code)) names.add(name);
				break;
			case "code":
				for (const name of collectIdentifiers(node.code)) names.add(name);
				break;
			case "element":
				for (const attribute of node.attributes) {
					if (attribute.kind !== "expression") continue;
					for (const name of collectIdentifiers(attribute.value ?? "")) {
						names.add(name);
					}
				}
				break;
			case "text":
				for (const part of node.parts) {
					if (part.type !== "expression") continue;
					for (const name of collectIdentifiers(part.code ?? "")) {
						names.add(name);
					}
				}
				break;
			case "if":
				for (const branch of node.branches) {
					if (branch.test) {
						for (const name of collectIdentifiers(branch.test)) names.add(name);
					}
				}
				break;
			case "each":
				for (const name of collectIdentifiers(node.list)) names.add(name);
				break;
			case "while":
				for (const name of collectIdentifiers(node.test)) names.add(name);
				break;
			case "case":
				for (const name of collectIdentifiers(node.subject)) names.add(name);
				break;
			case "mixin-call":
				for (const name of collectIdentifiers(node.args)) names.add(name);
				break;
			case "component":
			case "island":
				for (const value of Object.values(node.props)) {
					for (const name of collectIdentifiers(value)) names.add(name);
				}
				break;
		}

		const children = childrenOf(node);
		if (children) {
			for (const child of children) visit(child);
		}
	};

	for (const node of nodes) visit(node);
	return [...names].filter((name) => !declared.has(name));
}

/** Names of the parameters of a mixin. */
export function mixinParamNames(params: string): string[] {
	const names: string[] = [];
	const braced = /[{[(]([\s\S]*)[}\])]/.exec(params);
	const inner = braced ? (braced[1] as string) : params;
	for (const part of inner.split(",")) {
		const cleaned = part.trim().replace(/^[.]{3}/, "");
		if (!cleaned) continue;
		const name = /^[A-Za-z_$][\w$]*/.exec(cleaned);
		if (name) names.push(name[0]);
	}
	return names;
}

function childrenOf(node: Node): Node[] | undefined {
	switch (node.kind) {
		case "element":
		case "each":
		case "while":
		case "mixin":
		case "mixin-call":
		case "include":
		case "block":
		case "slot":
		case "comment":
			return node.children;
		case "if":
			return node.branches.flatMap((branch) => branch.children);
		default:
			return undefined;
	}
}