import { Type } from "@mariozechner/pi-ai";
import { defineTool, type ExtensionAPI } from "@mariozechner/pi-coding-agent";

const SIMPLE_EXPR = /^[0-9()+\-*/% \t]+$/;

export default function helloEvalExtension(pi: ExtensionAPI) {
	const helloEvalTool = defineTool({
		name: "hello_eval",
		label: "Hello Eval",
		description: "Evaluate a simple arithmetic expression using bash and return the result.",
		promptSnippet: "Evaluate small arithmetic expressions like 1+1 or (2+3)*4.",
		promptGuidelines: [
			"Use hello_eval for small arithmetic checks instead of doing the arithmetic silently.",
		],
		parameters: Type.Object({
			expression: Type.String({
				description: "Arithmetic expression to evaluate, for example 1+1 or (2+3)*4",
			}),
		}),

		async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
			const expression = params.expression.trim();
			if (!expression) {
				throw new Error("expression must not be empty");
			}
			if (!SIMPLE_EXPR.test(expression)) {
				throw new Error("expression contains unsupported characters");
			}

			const escaped = expression.replace(/\\/g, "\\\\").replace(/"/g, '\\"').replace(/\$/g, "\\$");
			const command = `printf '%s\\n' "$(( ${escaped} ))"`;
			const result = await pi.exec("bash", ["-lc", command], { signal, timeout: 5_000 });

			if (result.code !== 0) {
				throw new Error(result.stderr || `bash exited with code ${result.code}`);
			}

			const value = result.stdout.trim();
			return {
				content: [{ type: "text", text: value }],
				details: {
					expression,
					value,
					command,
				},
			};
		},
	});

	pi.registerTool(helloEvalTool);

	pi.registerCommand("hello-eval-demo", {
		description: "Queue a demo message that asks the model to call hello_eval on 1+1",
		handler: async (_args, ctx) => {
			ctx.ui.notify("Queued hello_eval demo prompt", "info");
			pi.sendUserMessage("Use the hello_eval tool to evaluate 1+1 and tell me the result.", {
				deliverAs: "followUp",
			});
		},
	});
}
