import { existsSync, statSync } from "node:fs";
import { readFile } from "node:fs/promises";
import { dirname, isAbsolute, join, resolve } from "node:path";
import { createRequire } from "node:module";
import { compileTemplate } from "./compiler.ts";
import { lex } from "./lexer.ts";
import { collectProps, parse } from "./parser.ts";
import type { Node, PropDeclaration } from "./types.ts";
import {
	isSpugComponentRender,
	type ComponentContext,
	type ComponentRegistry,
	type SpugComponentRender,
} from "../components.ts";
import { compileCss, styleTag } from "../css/index.ts";
import { compileStcssModule } from "../css/modules.ts";
import { isRaw, raw, styleValue } from "../raw.ts";
import type { DeferredChunk } from "../../contracts/typegen/http/response.ts";
import type { MuddleHooks } from "../hooks.ts";

/**
 * Template runtime.
 *
 * Compiles templates, resolves includes and extends and provides the helpers
 * used by the generated code.
 */

export interface TemplateHelpers {
	escapeHtml(value: string): string;
	stringify(value: unknown): string;
	/** Renders a value, raw values are never escaped. */
	render(value: unknown, escaped: boolean): string;
	/** Transforms an object into the value of a style attribute. */
	style(value: unknown): string;
	attr(name: string, value: unknown): string;
	attrBoolean(name: string): string;
	/** Merges the classes of a tag. */
	classes(values: unknown[]): string;
}

export type CompiledFunction = (
	runtime: TemplateRuntime,
	locals: Record<string, unknown>,
) => Promise<string>;

export interface CompiledBlock {
	mode: "replace" | "append" | "prepend";
	fn: CompiledFunction;
}

export interface CompiledTemplate {
	/** Absolute path of the template. */
	path: string;
	template: CompiledFunction;
	blocks: Map<string, CompiledBlock>;
	mixins: Record<string, CompiledFunction>;
	/** Props declared by the template. */
	props: PropDeclaration[];
	/** Frontmatter of the template. */
	frontmatter: string;
	/** Frontmatter of the template, evaluated before the body. */
	frontmatterFn?: CompiledFunction;
	/** Identifiers used by the template. */
	uses: string[];
	/** The assertions of the template, written in a `$$$` block. */
	tests: string[];
	/** Template this template extends, if any. */
	extends?: string;
}

/** Context of a single render. */
export interface RenderContext {
	/** Locals of the render. */
	locals: Record<string, unknown>;
	/** Route parameters of the render. */
	params: Record<string, string>;
	/** Chunks rendered after the shell. */
	deferred: DeferredChunk[];
}

export interface TemplateRuntime {
	helpers: TemplateHelpers;
	/** File of the template, available in the frontmatter as `__muddle`. */
	file: { path: string; dir: string };
	getBlock(name: string, fallback: CompiledFunction): Promise<string>;
	renderComponent(
		name: string,
		props: Record<string, unknown>,
	): Promise<string>;
	renderIsland(name: string, props: Record<string, unknown>): Promise<string>;
	/** Server components registered in the application. */
	components: ComponentRegistry;
	/** Context of the current render. */
	context: RenderContext;
}

export const helpers: TemplateHelpers = {
	escapeHtml(value: string): string {
		return String(value)
			.replace(/&/g, "&amp;")
			.replace(/</g, "&lt;")
			.replace(/>/g, "&gt;")
			.replace(/"/g, "&quot;")
			.replace(/'/g, "&#39;");
	},
	stringify(value: unknown): string {
		if (value === null || value === undefined) return "";
		if (typeof value === "string") return value;
		if (typeof value === "number" || typeof value === "boolean") {
			return String(value);
		}
		if (value instanceof Date) return value.toISOString();
		if (typeof value === "object") {
			try {
				return JSON.stringify(value);
			} catch {
				return String(value);
			}
		}
		return String(value);
	},
	attr(name: string, value: unknown): string {
		if (value === false || value === null || value === undefined) return "";
		if (value === true) return ` ${name}="${name}"`;
		return ` ${name}="${helpers.escapeHtml(helpers.stringify(value))}"`;
	},
	attrBoolean(name: string): string {
		return ` ${name}="${name}"`;
	},
	render(value: unknown, escaped: boolean): string {
		if (isRaw(value)) return value.html;
		const text = helpers.stringify(value);
		return escaped ? helpers.escapeHtml(text) : text;
	},
	style(value: unknown): string {
		return styleValue(value);
	},
	classes(values: unknown[]): string {
		return values
			.filter(
				(value) =>
					value !== undefined &&
					value !== null &&
					value !== false &&
					value !== "",
			)
			.map((value) => String(value))
			.join(" ");
	},
};

interface CompiledArtifact {
	template: CompiledFunction;
	frontmatter?: CompiledFunction;
	blocks: Record<string, CompiledBlock>;
	mixins: Record<string, CompiledFunction>;
	/** Identifiers used by the template. */
	uses: string[];
	/** The assertions of the template, written in a `$$$` block. */
	tests: string[];
}

export interface TemplateEngineOptions {
	/** Directory holding the templates. */
	viewsDir: string;
	/** Root directory of the project, template paths are resolved against it. */
	root?: string;
	/** Directory holding the components, defaults to `viewsDir/components`. */
	componentsDir?: string;
	/** Server components available to the templates. */
	components?: ComponentRegistry;
	/** Layout templates wrapped around every page render. */
	globalLayouts?: string[];
	/** Contextual render expansions applied after Spug composition. */
	hooks?: MuddleHooks;
}

/** Merges the blocks of a child template into the blocks of its parent. */
export function mergeBlocks(
	parent: Map<string, CompiledBlock>,
	child: Map<string, CompiledBlock>,
): Map<string, CompiledBlock> {
	const merged = new Map(parent);
	for (const [name, block] of child) {
		const existing = merged.get(name);
		if (!existing) {
			merged.set(name, block);
			continue;
		}
		if (block.mode === "append") {
			const parentFn = existing.fn;
			const childFn = block.fn;
			merged.set(name, {
				mode: "replace",
				fn: async (runtime, locals) =>
					`${await parentFn(runtime, locals)}${await childFn(runtime, locals)}`,
			});
			continue;
		}
		if (block.mode === "prepend") {
			const parentFn = existing.fn;
			const childFn = block.fn;
			merged.set(name, {
				mode: "replace",
				fn: async (runtime, locals) =>
					`${await childFn(runtime, locals)}${await parentFn(runtime, locals)}`,
			});
			continue;
		}
		merged.set(name, block);
	}
	return merged;
}

export class TemplateEngine {
	private readonly cache = new Map<
		string,
		{ mtime: number; compiled: CompiledTemplate }
	>();

	constructor(private readonly options: TemplateEngineOptions) {}

	/** Directory holding the templates. */
	get viewsDir(): string {
		return this.options.viewsDir;
	}

	/** Directory holding the components. */
	get componentsDir(): string {
		return (
			this.options.componentsDir ?? join(this.options.viewsDir, "components")
		);
	}

	/** Server components available to the templates. */
	get components(): ComponentRegistry {
		return this.options.components ?? new Map();
	}

	/** Resolves a template path relative to the views directory. */
	resolvePath(templatePath: string, from?: string): string {
		const candidates: string[] = [];
		if (isAbsolute(templatePath)) {
			candidates.push(templatePath);
		} else {
			if (from) candidates.push(resolve(dirname(from), templatePath));
			candidates.push(resolve(this.options.viewsDir, templatePath));
			if (this.options.root) {
				candidates.push(resolve(this.options.root, templatePath));
			}
		}
		for (const candidate of candidates) {
			if (existsSync(candidate)) return candidate;
			if (existsSync(`${candidate}.spug`)) return `${candidate}.spug`;
		}
		throw new Error(`Template not found: ${templatePath}`);
	}

	/** Loads and compiles a template, using the cache when possible. */
	async load(templatePath: string, from?: string): Promise<CompiledTemplate> {
		const path = this.resolvePath(templatePath, from);
		const mtime = statSync(path).mtimeMs;
		const cached = this.cache.get(path);
		if (cached && cached.mtime === mtime) return cached.compiled;

		const compiled = await this.compile(path);
		this.cache.set(path, { mtime, compiled });
		return compiled;
	}

	/** Compiles a template without touching the cache. */
	async compile(path: string): Promise<CompiledTemplate> {
		const source = await readFile(path, "utf8");
		const tokens = lex(source);
		const parsed = parse(tokens);

		const frontmatter = await this.linkIncludes(
			parsed.nodes,
			parsed.frontmatter,
			path,
		);
		const extendsPath = parsed.nodes.find((node) => node.kind === "extends");
		const { source: compiledSource } = compileTemplate(parsed.nodes, {
			path,
			dir: dirname(path),
			frontmatter,
		});

		const factory = new Function(compiledSource) as () => CompiledArtifact;
		const artifact = factory();

		const blocks = new Map<string, CompiledBlock>(
			Object.entries(artifact.blocks),
		);
		const props = collectProps(parsed.nodes);

		return {
			path,
			template: artifact.template,
			frontmatterFn: artifact.frontmatter,
			uses: artifact.uses ?? [],
			tests: artifact.tests ?? [],
			blocks,
			mixins: artifact.mixins,
			props,
			frontmatter,
			extends: extendsPath?.kind === "extends" ? extendsPath.path : undefined,
		};
	}

	/** Renders a template with the given locals. */
	async render(
		templatePath: string,
		initialLocals: Record<string, unknown> = {},
		context?: Partial<RenderContext>,
	): Promise<string> {
		const loaded = await this.load(templatePath);
		const locals: Record<string, unknown> = { ...initialLocals };
		const compiled = loaded;
		this.fillLocals(compiled, locals);
		const renderContext: RenderContext = {
			locals,
			params: context?.params ?? {},
			deferred: context?.deferred ?? [],
		};

		let html: string;
		if (compiled.extends) {
			const parent = await this.load(compiled.extends, compiled.path);
			const runtime = this.createRuntime(
				compiled,
				locals,
				compiled.blocks,
				renderContext,
			);
			await this.runFrontmatter(compiled, locals, runtime);
			this.fillLocals(parent, locals);
			const blocks = mergeBlocks(parent.blocks, compiled.blocks);
			const parentRuntime = this.createRuntime(
				parent,
				locals,
				blocks,
				renderContext,
			);
			await this.runFrontmatter(parent, locals, parentRuntime);
			html = await parent.template(parentRuntime, locals);
		} else {
			const blocks = mergeBlocks(compiled.blocks, new Map());
			const runtime = this.createRuntime(
				compiled,
				locals,
				blocks,
				renderContext,
			);
			await this.runFrontmatter(compiled, locals, runtime);
			html = await compiled.template(runtime, locals);
		}
		for (const layout of this.options.globalLayouts ?? []) {
			html = await this.renderLayout(
				layout,
				{ ...locals, content: raw(html) },
				renderContext,
			);
		}
		return this.options.hooks
			? this.options.hooks.renderSpug({ html, path: loaded.path, locals })
			: html;
	}

	private async renderLayout(
		path: string,
		locals: Record<string, unknown>,
		context: RenderContext,
	): Promise<string> {
		const compiled = await this.load(path);
		const runtime = this.createRuntime(
			compiled,
			locals,
			compiled.blocks,
			context,
		);
		await this.runFrontmatter(compiled, locals, runtime);
		return compiled.template(runtime, locals);
	}

	/**
	 * Fills the missing identifiers of a template with `undefined`.
	 */
	private fillLocals(
		compiled: CompiledTemplate,
		locals: Record<string, unknown>,
	): void {
		for (const name of compiled.uses ?? []) {
			if (!(name in locals)) locals[name] = undefined;
		}
	}

	/**
	 * Evaluates the frontmatter of a template.
	 */
	private async runFrontmatter(
		compiled: CompiledTemplate,
		locals: Record<string, unknown>,
		runtime: TemplateRuntime,
	): Promise<void> {
		if (!compiled.frontmatterFn) return;
		const values = await compiled.frontmatterFn(runtime, locals);
		Object.assign(locals, values);
	}

	/** Renders a component by name. */
	async renderComponent(
		name: string,
		props: Record<string, unknown>,
		context?: Partial<RenderContext>,
		from?: string,
	): Promise<string> {
		const path = this.resolveComponent(name, from);
		const compiled = await this.load(path);
		props = { ...props };
		this.fillLocals(compiled, props);
		const runtime = this.createRuntime(
			compiled,
			props,
			compiled.blocks,
			context,
		);
		await this.runFrontmatter(compiled, props, runtime);
		const html = await compiled.template(runtime, props);
		return this.attachComponentModule(path, html);
	}

	/** Renders a component and marks it for client side hydration. */
	async renderIsland(
		name: string,
		props: Record<string, unknown>,
		context?: Partial<RenderContext>,
	): Promise<string> {
		const html = await this.renderComponent(name, props, context);
		const payload = JSON.stringify(props).replace(/</g, "\\u003c");
		return `${html}<script type="application/json" data-muddle-island="${name}">${payload}</script>`;
	}

	private resolveComponent(name: string, from?: string): string {
		if (isAbsolute(name) && existsSync(name)) return name;
		const candidates = [
			...(from && name.startsWith(".")
				? [resolve(dirname(from), name), resolve(dirname(from), `${name}.spug`)]
				: []),
			join(this.componentsDir, `${name}.spug`),
			join(this.componentsDir, name),
			join(this.viewsDir, `${name}.spug`),
			join(this.viewsDir, name),
		];
		for (const candidate of candidates) {
			if (existsSync(candidate) && statSync(candidate).isFile()) {
				return candidate;
			}
		}
		const packagePath = this.resolvePackageComponent(name, from);
		if (packagePath) return packagePath;
		throw new Error(`Component not found: ${name}`);
	}

	private resolvePackageComponent(
		name: string,
		from?: string,
	): string | undefined {
		if (name.startsWith(".")) return undefined;
		const packageName = name.startsWith("@")
			? name.split("/").slice(0, 2).join("/")
			: (name.split("/")[0] as string);
		const subpath = name.slice(packageName.length).replace(/^\//, "");
		const require = createRequire(
			from ?? join(this.options.root ?? this.viewsDir, "package.json"),
		);
		let packageJson: string;
		try {
			packageJson = require.resolve(`${packageName}/package.json`);
		} catch {
			return undefined;
		}
		const root = dirname(packageJson);
		const candidates = [
			subpath && join(root, subpath),
			subpath && join(root, `${subpath}.spug`),
			join(root, "component.spug"),
			join(root, "index.spug"),
		].filter(Boolean) as string[];
		return candidates.find(
			(candidate) => existsSync(candidate) && statSync(candidate).isFile(),
		);
	}

	private async attachComponentModule(
		path: string,
		html: string,
	): Promise<string> {
		const modulePath = path.replace(/\.spug$/, ".module.stcss");
		if (!existsSync(modulePath)) return html;
		const module = await compileStcssModule(
			await readFile(modulePath, "utf8"),
			{
				filename: modulePath,
				root: this.options.root,
				hooks: this.options.hooks,
			},
		);
		const scopedHtml = html.replace(
			/\bclass=(['"])([^'"]*)\1/g,
			(_all, quote: string, value: string) => {
				const classes = value
					.split(/\s+/)
					.map((name) => module.classes[name] ?? name)
					.join(" ");
				return `class=${quote}${classes}${quote}`;
			},
		);
		return `${styleTag(module.css)}${scopedHtml}`;
	}

	private createRuntime(
		compiled: CompiledTemplate,
		locals: Record<string, unknown>,
		blocks: Map<string, CompiledBlock>,
		context?: Partial<RenderContext>,
	): TemplateRuntime {
		const renderContext: RenderContext = {
			locals,
			params: context?.params ?? {},
			deferred: context?.deferred ?? [],
		};

		const componentContext: ComponentContext = {
			defer: (id, content) => {
				renderContext.deferred.push({
					id,
					content:
						typeof content === "string" ? Promise.resolve(content) : content,
				});
			},
			locals,
			params: renderContext.params,
		};

		let deferredCount = 0;

		const runtime: TemplateRuntime = {
			helpers,
			file: { path: compiled.path, dir: dirname(compiled.path) },
			getBlock: async (name, fallback) => {
				const block = blocks.get(name);
				if (!block) return fallback(runtime, locals);
				return block.fn(runtime, locals);
			},
			renderComponent: async (name, props) => {
				const definition = this.components.get(name);
				if (!definition) {
					return this.renderComponent(
						name,
						props,
						renderContext,
						compiled.path,
					);
				}
				if (definition.deferred) {
					deferredCount += 1;
					const id = `${name}-${deferredCount}`;
					componentContext.defer(
						id,
						Promise.resolve(
							definition.component.render(props, componentContext),
						).then((output) =>
							this.renderComponentOutput(
								output,
								props,
								renderContext,
								definition.requiresSpug,
							),
						),
					);
					return `<div data-muddle-deferred="${id}"></div>`;
				}
				return this.renderComponentOutput(
					await definition.component.render(props, componentContext),
					props,
					renderContext,
					definition.requiresSpug,
				);
			},
			renderIsland: (name, props) =>
				this.renderIsland(name, props, renderContext),
			components: this.components,
			context: renderContext,
		};

		return runtime;
	}

	private async renderComponentOutput(
		output: string | SpugComponentRender,
		props: Record<string, unknown>,
		context: Partial<RenderContext>,
		requiresSpug = false,
	): Promise<string> {
		if (!isSpugComponentRender(output)) {
			if (requiresSpug)
				throw new Error(
					"@Component renderers must return { spug, stcss } rather than raw HTML",
				);
			return output;
		}
		const html = await this.renderComponent(
			output.spug,
			{ ...props, ...(output.props ?? {}) },
			context,
		);
		const styles = Array.isArray(output.stcss)
			? output.stcss
			: output.stcss
				? [output.stcss]
				: [];
		const css = await Promise.all(
			styles.map(async (style) => {
				const path = resolve(this.componentsDir, style);
				if (!existsSync(path)) return compileCss(style);
				const source = await readFile(path, "utf8");
				return path.endsWith(".module.stcss")
					? (
							await compileStcssModule(source, {
								filename: path,
								root: this.options.root,
								hooks: this.options.hooks,
							})
						).css
					: compileCss(source, { hooks: this.options.hooks });
			}),
		);
		return `${css.filter(Boolean).map(styleTag).join("")}${html}`;
	}

	/** Inlines the included files and collects their frontmatter. */
	private async linkIncludes(
		nodes: Node[],
		frontmatter: string,
		from: string,
	): Promise<string> {
		let collected = frontmatter;

		for (const node of nodes) {
			const children = childrenOf(node);
			if (!children) continue;

			if (node.kind === "include") {
				const path = this.resolvePath(node.path, from);
				const source = await readFile(path, "utf8");
				const parsed = parse(lex(source));
				node.children = parsed.nodes;
				const nested = await this.linkIncludes(
					parsed.nodes,
					parsed.frontmatter,
					path,
				);
				collected = `${collected}\n${nested}`;
				continue;
			}

			collected = await this.linkIncludes(children, collected, from);
		}

		return collected;
	}
}

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;
	}
}