import { fileURLToPath } from "node:url";
import * as ai from "@earendil-works/pi-ai";
import type {
	ExtensionAPI,
	ToolDefinition,
} from "@earendil-works/pi-coding-agent";
import * as sdk from "@earendil-works/pi-coding-agent";
import * as tui from "@earendil-works/pi-tui";
import * as schema from "typebox";
import type { EphemletRuntime } from "./runtime.ts";
import { closeDebug, dbg } from "./src/debug.ts";
import { ephemletRenderers } from "./src/tool-renderer.ts";

export default function registerEphemlets(pi: ExtensionAPI): void {
	let runtime: Promise<EphemletRuntime> | undefined;
	let generation = 0;
	let activeBeforeHide = false;
	const get = () =>
		(runtime ??= import("jiti")
			.then(async ({ createJiti }) => {
				// Pi aliases are unavailable to native lazy chunks; preserve them in this lazy loader.
				const loader = createJiti(import.meta.url, {
					fsCache: false,
					moduleCache: false,
					tryNative: false,
					virtualModules: {
						"@earendil-works/pi-ai": ai,
						"@earendil-works/pi-coding-agent": sdk,
						"@earendil-works/pi-tui": tui,
						typebox: schema,
					},
				});
				const path = fileURLToPath(
					new URL(
						import.meta.url.endsWith(".ts") ? "./runtime.ts" : "./runtime.js",
						import.meta.url,
					),
				);
				const { EphemletRuntime: Runtime } =
					await loader.import<typeof import("./runtime.ts")>(path);
				return new Runtime(pi, dbg);
			})
			.catch((error: unknown) => {
				runtime = undefined;
				throw error;
			}));
	pi.on("session_start", (_event, ctx) => {
		const current = ++generation;
		const active = pi.getActiveTools();
		tool.exposure = ctx.mode === "tui" ? "model-only" : "hidden";
		pi.registerTool(tool);
		if (ctx.mode === "tui" && activeBeforeHide && !active.includes("ephemlet"))
			pi.setActiveTools([...active, "ephemlet"]);
		if (ctx.mode !== "tui") {
			activeBeforeHide ||= active.includes("ephemlet");
			if (active.includes("ephemlet"))
				pi.setActiveTools(active.filter((name) => name !== "ephemlet"));
			if (runtime) void runtime.then((host) => host.dispose()).catch(() => {});
			return;
		}
		activeBeforeHide = false;
		dbg?.("session.start");
		const hasEphemlets = ctx.sessionManager
			.getBranch()
			.some(
				(entry) => entry.type === "custom" && entry.customType === "ephemlets",
			);
		if (runtime || hasEphemlets) {
			ctx.ui.setStatus("ephemlets", "Restoring ephemlets…");
			void get()
				.then(async (host) => {
					if (current === generation) await host.restore(ctx);
				})
				.catch((error: unknown) => {
					if (current === generation)
						ctx.ui.notify(`Ephemlet restore failed: ${String(error)}`, "error");
				})
				.finally(() => {
					if (current === generation) ctx.ui.setStatus("ephemlets", undefined);
				});
		}
	});
	pi.on("session_tree", async (_event, ctx) => {
		++generation;
		if (ctx.mode === "tui") ctx.ui.setStatus("ephemlets", undefined);
		if (ctx.mode === "tui" && runtime) await (await runtime).restore(ctx);
	});
	pi.on("before_agent_start", async (_event, ctx) => {
		if (ctx.mode === "tui" && runtime) (await runtime).setContext(ctx);
	});
	pi.on("session_shutdown", async () => {
		++generation;
		if (runtime) (await runtime).dispose();
		dbg?.("session.shutdown");
		closeDebug();
	});
	pi.registerShortcut("alt+e", {
		description: "Hide/show ephemlet panel",
		handler: async (ctx) => {
			if (ctx.mode === "tui") await (await get()).toggle(ctx);
		},
	});
	pi.registerShortcut("ctrl+alt+e", {
		description: "Focus ephemlet panel",
		handler: async (ctx) => {
			if (ctx.mode === "tui") await (await get()).show(ctx);
		},
	});
	pi.registerCommand("ephemlets", {
		description: "Toggle ephemlets, or show/hide/focus panel",
		getArgumentCompletions(prefix) {
			return ["show", "hide", "focus"]
				.filter((value) => value.startsWith(prefix))
				.map((value) => ({ value, label: value }));
		},
		handler: async (args, ctx) => {
			if (ctx.mode !== "tui") {
				ctx.ui.notify("Ephemlets require interactive TUI mode", "warning");
				return;
			}
			const host = await get();
			if (args === "show" || args === "focus") await host.show(ctx);
			else if (args === "hide") await host.hide(ctx);
			else if (!args) await host.toggle(ctx);
			else ctx.ui.notify("Use /ephemlets [show|hide|focus]", "warning");
		},
	});
	// Declare capabilities eagerly; load domain implementation/compiler only on demand.
	const tool: ToolDefinition<ReturnType<typeof parameters>, unknown> = {
		name: "ephemlet",
		annotations: { readOnlyHint: false },
		exposure: "model-only",
		label: "Ephemlet",
		description:
			"Create and update session-persistent Markdown reports or native interactive Pi UI without /reload. Use list to discover IDs and authoring guidance; get returns bounded metadata unless include requests source/state. Put writes source directly; build delegates interactive construction using API tokens. State replaces JSON state; delete removes an ephemlet from the current branch; show/hide controls visibility. Available only in the interactive TUI.",
		promptSnippet:
			"Present living task reports and purpose-built native interactive UI",
		promptGuidelines: [
			"Use ephemlet for substantial task reports (subgoals, decisions, forks, failed attempts, verified and omitted work) or useful interactions; the main agent owns intent and authoritative facts. Markdown reports need no builder.",
			'Discover existing ephemlets with list on resume or before creating duplicates. Reuse IDs; read get for the current revision before put/build/state updates, which require expectedRevision. Request include: ["source"] or ["state"] only when needed; metadata is the default.',
			"Before directly authoring interactive source, use the guidance path returned by list and read it. Components are trusted native code with Pi's permissions; never embed credentials in source/state. Build uses this guidance automatically.",
			'The agent cannot see the rendered panel. Read include: ["state"] when task decisions need current user interaction data; local state changes do not automatically send conversation messages.',
			"Build is asynchronous by default and returns ID/jobId, not completed interaction. Ready/failed notices do not start a new turn; inspect get before claiming success. Use wait: true only when the result is needed before continuing.",
			"Hide preserves an ephemlet; Close/delete removes it from the active branch. Save queues an exact snapshot as a user follow-up for the main agent to preserve as a durable deliverable; it does not export a file itself.",
		],
		executionMode: "sequential",
		constrainedSampling: { type: "json_schema", strict: "prefer" },
		parameters: parameters(),
		outputSchema: outputSchema(),
		...ephemletRenderers,
		async execute(_id, params, signal, _onUpdate, ctx) {
			if (ctx.mode !== "tui")
				throw new Error("Ephemlets require interactive TUI mode");
			return (await get()).execute(params, ctx, signal);
		},
	};
	pi.registerTool(tool);
}

import { Type } from "typebox";

function jsonSchema() {
	return Type.Union([
		Type.Null(),
		Type.Boolean(),
		Type.Number(),
		Type.String(),
		Type.Array(Type.Ref("#/$defs/json")),
		Type.Object({}, { additionalProperties: Type.Ref("#/$defs/json") }),
	]);
}

function parameters() {
	return Type.Object(
		{
			action: Type.Union(
				(
					[
						"put",
						"build",
						"list",
						"get",
						"state",
						"delete",
						"show",
						"hide",
					] as const
				).map((value) => Type.Literal(value)),
				{
					description:
						"Operation; get reads metadata, state replaces durable JSON.",
				},
			),
			id: Type.Optional(
				Type.String({
					minLength: 1,
					maxLength: 80,
					description:
						"Stable ID; required for get/state/delete and existing updates.",
				}),
			),
			expectedRevision: Type.Optional(
				Type.Integer({
					minimum: 1,
					description:
						"Current revision from get; required for existing put/build/state updates.",
				}),
			),
			title: Type.Optional(Type.String({ minLength: 1, maxLength: 256 })),
			purpose: Type.Optional(Type.String({ minLength: 1, maxLength: 4096 })),
			kind: Type.Optional(
				Type.Union([Type.Literal("markdown"), Type.Literal("interactive")]),
			),
			source: Type.Optional(
				Type.String({
					maxLength: 131072,
					description: "Markdown or complete native TypeScript module source",
				}),
			),
			state: Type.Optional(
				Type.Ref("#/$defs/json", {
					description:
						"JSON-serializable durable data/state; required for state action.",
				}),
			),
			include: Type.Optional(
				Type.Array(
					Type.Union([Type.Literal("source"), Type.Literal("state")]),
					{
						minItems: 1,
						maxItems: 2,
						uniqueItems: true,
						description:
							"Get only: explicitly include source and/or state; rollback internals are never returned.",
					},
				),
			),
			brief: Type.Optional(
				Type.String({
					maxLength: 32768,
					description: "Authoritative interaction intent/facts for builder",
				}),
			),
			wait: Type.Optional(
				Type.Boolean({
					description: "Await builder; default false (background)",
				}),
			),
			model: Type.Optional(
				Type.String({
					maxLength: 256,
					description: "Optional builder provider/model; default current model",
				}),
			),
		},
		{ $defs: { json: jsonSchema() } },
	);
}

function outputSchema() {
	const ephemlet = Type.Object({
		id: Type.String(),
		title: Type.String(),
		purpose: Type.String(),
		kind: Type.Union([Type.Literal("markdown"), Type.Literal("interactive")]),
		revision: Type.Integer({ minimum: 1 }),
		status: Type.Union([
			Type.Literal("ready"),
			Type.Literal("building"),
			Type.Literal("failed"),
		]),
		error: Type.Union([Type.String(), Type.Null()]),
		jobId: Type.Union([Type.String(), Type.Null()]),
		source: Type.Optional(Type.String()),
		state: Type.Optional(Type.Ref("#/$defs/json")),
	});
	return Type.Union(
		[
			ephemlet,
			Type.Object({ ephemlets: Type.Array(ephemlet), guidance: Type.String() }),
			Type.Object({
				visible: Type.Boolean(),
				selected: Type.Union([Type.String(), Type.Null()]),
			}),
			Type.Object({ deleted: Type.String() }),
		],
		{ $defs: { json: jsonSchema() } },
	);
}