import assert from "node:assert/strict";
import test from "node:test";
import {
	createBashToolDefinition,
	createEditToolDefinition,
	createFindToolDefinition,
	createGrepToolDefinition,
	createLsToolDefinition,
	createPowerShellToolDefinition,
	createReadToolDefinition,
	createWriteToolDefinition,
	type ExtensionAPI,
	type ToolDefinition,
	type ToolRendererResolver,
} from "@selesai/code";
import toolDisplayExtension from "../src/index.ts";
import { BUILT_IN_TOOL_OVERRIDE_NAMES } from "../src/types.ts";
import { resolveToolRenderers } from "./tool-renderer-harness.ts";

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

interface CapturedHandler {
	event: string;
	handler: (...args: unknown[]) => unknown;
}

interface CapturedCommand {
	name: string;
	description?: string;
	handler?: (...args: unknown[]) => unknown;
}

function createApiStub(
	overrides: Partial<{
		registerTool: (tool: unknown) => void;
		registerToolRenderer: (resolver: ToolRendererResolver) => void;
		registerCommand: (name: string, cmd: unknown) => void;
		on: (event: string, handler: (...args: unknown[]) => unknown) => void;
		getAllTools: () => unknown[];
		getCommands: () => Array<{ name: string }>;
	}> = {},
): {
	api: ExtensionAPI;
	capturedTools: Array<{ name: string } & Record<string, unknown>>;
	capturedResolvers: ToolRendererResolver[];
	capturedCommands: CapturedCommand[];
	capturedHandlers: CapturedHandler[];
} {
	const capturedTools: Array<{ name: string } & Record<string, unknown>> = [];
	const capturedResolvers: ToolRendererResolver[] = [];
	const capturedCommands: CapturedCommand[] = [];
	const capturedHandlers: CapturedHandler[] = [];

	const api = {
		registerTool(tool: unknown): void {
			capturedTools.push(tool as { name: string } & Record<string, unknown>);
			overrides.registerTool?.(tool);
		},
		registerToolRenderer(resolver: ToolRendererResolver): void {
			capturedResolvers.push(resolver);
			overrides.registerToolRenderer?.(resolver);
		},
		registerCommand(name: string, cmd: unknown): void {
			capturedCommands.push({ name, ...(cmd as object) } as CapturedCommand);
			overrides.registerCommand?.(name, cmd);
		},
		on(event: string, handler: (...args: unknown[]) => unknown): void {
			capturedHandlers.push({ event, handler });
			overrides.on?.(event, handler);
		},
		getAllTools(): unknown[] {
			return overrides.getAllTools?.() ?? [];
		},
		getCommands(): Array<{ name: string }> {
			return overrides.getCommands?.() ?? [];
		},
	} as unknown as ExtensionAPI;

	return { api, capturedTools, capturedResolvers, capturedCommands, capturedHandlers };
}

function createHostDefinitions(cwd: string): Map<string, ToolDefinition> {
	const definitions = [
		createReadToolDefinition(cwd),
		createGrepToolDefinition(cwd),
		createFindToolDefinition(cwd),
		createLsToolDefinition(cwd),
		createBashToolDefinition(cwd),
		createPowerShellToolDefinition(cwd),
		createEditToolDefinition(cwd),
		createWriteToolDefinition(cwd),
	];
	return new Map(definitions.map((definition) => [definition.name, definition]));
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

test("entry point registers expected lifecycle handlers", () => {
	const { api, capturedHandlers } = createApiStub();
	toolDisplayExtension(api);

	const eventNames = capturedHandlers.map((h) => h.event);
	// Thinking-label handlers
	assert.ok(eventNames.includes("message_update"), "message_update handler registered");
	assert.ok(eventNames.includes("message_end"), "message_end handler registered");
	assert.ok(eventNames.includes("context"), "context handler registered");
	// Lifecycle handlers from index.ts directly
	assert.ok(eventNames.includes("session_start"), "session_start handler registered");
	assert.ok(eventNames.includes("before_agent_start"), "before_agent_start handler registered");
	// Write diff capture handlers.
	assert.ok(eventNames.includes("tool_call"), "tool_call handler registered");
	assert.ok(eventNames.includes("tool_result"), "tool_result handler registered");
	const sessionStartCount = eventNames.filter((e) => e === "session_start").length;
	assert.ok(sessionStartCount >= 1, "at least one session_start handler registered");
	const beforeAgentStartCount = eventNames.filter((e) => e === "before_agent_start").length;
	assert.ok(beforeAgentStartCount >= 1, "at least one before_agent_start handler registered");
});

test("entry point registers tool-display command", () => {
	const { api, capturedCommands } = createApiStub();
	toolDisplayExtension(api);

	const cmdNames = capturedCommands.map((c) => c.name);
	assert.ok(cmdNames.includes("tool-display"), "tool-display command registered");
});

test("entry point registers exactly one renderer resolver and no tools", () => {
	const { api, capturedTools, capturedResolvers } = createApiStub();
	toolDisplayExtension(api);

	assert.equal(capturedResolvers.length, 1, "one resolver covers all built-ins");
	assert.equal(capturedTools.length, 0, "built-in tools are never re-registered");

	for (const toolName of BUILT_IN_TOOL_OVERRIDE_NAMES) {
		const renderers = resolveToolRenderers(capturedResolvers, toolName, () => undefined);
		assert.equal(typeof renderers?.renderCall, "function", `${toolName} has renderCall`);
		assert.equal(typeof renderers?.renderResult, "function", `${toolName} has renderResult`);
	}
});

test("session_start handler refreshes capabilities and notifies pending errors", async () => {
	const { api, capturedHandlers } = createApiStub();
	toolDisplayExtension(api);

	const sessionHandler = capturedHandlers.find((h) => h.event === "session_start")?.handler;
	assert.ok(sessionHandler, "session_start handler captured");

	const ctx = {
		ui: {
			theme: { fg: (_c: string, t: string) => t },
			notify: (_msg: string, _level: string) => {
				/* no-op */
			},
		},
	};

	// Should not throw
	await assert.doesNotReject(async () => sessionHandler({}, ctx));
});

test("before_agent_start handler refreshes capabilities without crashing", async () => {
	const { api, capturedHandlers } = createApiStub();
	toolDisplayExtension(api);

	const beforeHandler = capturedHandlers.find((h) => h.event === "before_agent_start")?.handler;
	assert.ok(beforeHandler, "before_agent_start handler captured");

	// Should not throw
	await assert.doesNotReject(async () => beforeHandler());
});

test("multiple calls to toolDisplayExtension are idempotent per load", () => {
	const { api, capturedTools, capturedResolvers, capturedCommands } = createApiStub();

	// Call twice
	toolDisplayExtension(api);
	toolDisplayExtension(api);

	assert.equal(capturedResolvers.length, 2, "each load registers its own resolver");
	assert.equal(capturedTools.length, 0, "no tool is registered by either load");

	const cmdNames = capturedCommands.map((c) => c.name);
	assert.ok(cmdNames.filter((n) => n === "tool-display").length >= 1, "command registered at least once");
});

test("entry point tolerates empty getAllTools and getCommands results", () => {
	const { api } = createApiStub({
		getAllTools: () => [],
		getCommands: () => [],
	});

	assert.doesNotThrow(() => toolDisplayExtension(api));
});

test("entry point tolerates tools with existing owners in getAllTools", () => {
	const { api, capturedTools, capturedResolvers } = createApiStub({
		getAllTools: () => [
			{ name: "read", sourceInfo: { source: "local", path: "/ext/read.ts" } },
			{ name: "edit", sourceInfo: { source: "local", path: "/ext/edit.ts" } },
			{ name: "grep", sourceInfo: { source: "local", path: "/ext/grep.ts" } },
		],
		getCommands: () => [{ name: "custom" }],
	});

	assert.doesNotThrow(() => toolDisplayExtension(api));
	assert.equal(capturedTools.length, 0, "externally owned built-ins are not touched");
	assert.ok(resolveToolRenderers(capturedResolvers, "read", () => undefined), "renderers still apply");
});

test("graceful degradation: extension throws when registerCommand is missing", () => {
	const minimalApi = {
		registerTool(): void {
			/* no-op */
		},
		registerToolRenderer(): void {
			/* no-op */
		},
		on(): void {
			/* no-op */
		},
		getAllTools(): unknown[] {
			return [];
		},
		getCommands(): Array<{ name: string }> {
			return [];
		},
	} as unknown as ExtensionAPI;

	// registerToolDisplayCommand calls pi.registerCommand directly, so this
	// is expected to throw in a peer-dep mismatch scenario.
	assert.throws(
		() => toolDisplayExtension(minimalApi),
		/registerCommand/i,
		"missing registerCommand should propagate",
	);
});

test("graceful degradation: extension throws when on is missing", () => {
	const minimalApi = {
		registerTool(): void {
			/* no-op */
		},
		registerToolRenderer(): void {
			/* no-op */
		},
		registerCommand(): void {
			/* no-op */
		},
		getAllTools(): unknown[] {
			return [];
		},
		getCommands(): Array<{ name: string }> {
			return [];
		},
	} as unknown as ExtensionAPI;

	// registerNativeUserMessageBox calls pi.on, so this should throw when on is missing
	assert.throws(
		() => toolDisplayExtension(minimalApi),
		/pi\.on is not a function|on is not a function/i,
		"missing on should propagate",
	);
});

test("lifecycle events fire in expected order during a session lifecycle", async () => {
	// Simulate the sequence: setup → before_agent_start → session_start
	const { api, capturedHandlers } = createApiStub();

	toolDisplayExtension(api);

	// Manually invoke handlers in expected lifecycle order
	const beforeHandler = capturedHandlers.find((h) => h.event === "before_agent_start")?.handler;
	const sessionHandler = capturedHandlers.find((h) => h.event === "session_start")?.handler;
	const messageUpdateHandler = capturedHandlers.find((h) => h.event === "message_update")?.handler;
	const messageEndHandler = capturedHandlers.find((h) => h.event === "message_end")?.handler;
	const contextHandler = capturedHandlers.find((h) => h.event === "context")?.handler;

	assert.ok(beforeHandler, "before_agent_start handler found");
	assert.ok(sessionHandler, "session_start handler found");

	// Simulate a session lifecycle
	await beforeHandler();
	await sessionHandler(
		{},
		{ ui: { theme: { fg: (_c: string, t: string) => t }, notify: () => {} } },
	);

	// Simulate message lifecycle for thinking labels
	if (messageUpdateHandler) {
		await messageUpdateHandler(
			{
				message: {
					role: "assistant",
					api: "anthropic-messages",
					content: [{ type: "thinking", thinking: "test" }],
				},
			},
			{ ui: { theme: { fg: (_c: string, t: string) => `[${_c}]${t}` } } },
		);
	}

	if (messageEndHandler) {
		await messageEndHandler(
			{
				message: {
					role: "assistant",
					api: "openai-chat",
					content: [{ type: "thinking", thinking: "done" }],
				},
			},
			{},
		);
	}

	if (contextHandler) {
		await contextHandler(
			{
				messages: [
					{
						role: "assistant",
						content: [{ type: "thinking", thinking: "\x1b[31mThinking: \x1b[0mcontext" }],
					},
				],
			},
			{},
		);
	}

	// All handlers executed without throwing - this is the main assertion
	assert.ok(true, "lifecycle handlers completed without error");
});

test("session_start handler tolerates missing ctx.ui", async () => {
	const { api, capturedHandlers } = createApiStub();
	toolDisplayExtension(api);

	const sessionHandler = capturedHandlers.find((h) => h.event === "session_start")?.handler;
	assert.ok(sessionHandler);

	// ctx with no ui (edge case from older pi versions)
	await assert.doesNotReject(async () => sessionHandler({}, {}));
});

test("before_agent_start handler tolerates being called multiple times", async () => {
	const { api, capturedHandlers } = createApiStub();
	toolDisplayExtension(api);

	const beforeHandler = capturedHandlers.find((h) => h.event === "before_agent_start")?.handler;
	assert.ok(beforeHandler);

	await assert.doesNotReject(async () => beforeHandler());
	await assert.doesNotReject(async () => beforeHandler());
	await assert.doesNotReject(async () => beforeHandler());
});

test("session_start handler tolerates being called multiple times", async () => {
	const { api, capturedHandlers } = createApiStub();
	toolDisplayExtension(api);

	const sessionHandler = capturedHandlers.find((h) => h.event === "session_start")?.handler;
	assert.ok(sessionHandler);

	const ctx = { ui: { theme: {}, notify: () => {} } };
	await assert.doesNotReject(async () => sessionHandler({}, ctx));
	await assert.doesNotReject(async () => sessionHandler({}, ctx));
	await assert.doesNotReject(async () => sessionHandler({}, ctx));
});

test("the host keeps its own definitions for every built-in after load and lifecycle", async () => {
	const cwd = process.cwd();
	const definitions = createHostDefinitions(cwd);
	const { api, capturedTools, capturedResolvers, capturedHandlers } = createApiStub();

	toolDisplayExtension(api);
	for (const handler of capturedHandlers) {
		if (handler.event === "session_start" || handler.event === "before_agent_start") {
			await handler.handler({}, { ui: { theme: {}, notify: () => {} } });
		}
	}

	assert.equal(capturedTools.length, 0);
	for (const toolName of BUILT_IN_TOOL_OVERRIDE_NAMES) {
		const definition = definitions.get(toolName);
		assert.ok(definition, `missing host definition for ${toolName}`);
		assert.equal(definition.name, toolName);
		assert.ok(definition.description.length > 0);
		assert.ok(definition.parameters);
	}

	// Renderers carry no prompt/parameter metadata: A3 holds because the host
	// definition is untouched, not because the extension re-created it.
	for (const toolName of BUILT_IN_TOOL_OVERRIDE_NAMES) {
		const renderers = resolveToolRenderers(capturedResolvers, toolName, () => undefined);
		assert.ok(renderers);
		assert.deepEqual(
			Object.keys(renderers).sort(),
			["renderCall", "renderResult", "renderShell"],
			`${toolName} renderers expose only rendering hooks`,
		);
	}
});