import { Type } from "typebox";
import { DEFAULT_CODE_MODE_OUTPUT_TOKENS, MAX_CODE_MODE_OUTPUT_TOKENS, } from "./host-protocol.js";
import { EXEC_DESCRIPTION, WAIT_DESCRIPTION, } from "./custom-tool-prompt.js";
import { createCodeModeRenderTracker } from "./render-tracker.js";
import { renderExecCall, renderTrackedCodeModeResult, renderWaitCall, } from "./rendering.js";
import { toCodeModeToolResult } from "./tool-result.js";
const DEFAULT_WAIT_MS = 10_000;
const EXEC_PARAMETERS = Type.Object({
    code: Type.String(),
});
const WAIT_PARAMETERS = Type.Object({
    cell_id: Type.String(),
    yield_time_ms: Type.Optional(Type.Integer({
        minimum: 0,
        default: DEFAULT_WAIT_MS,
    })),
    max_tokens: Type.Optional(Type.Integer({
        minimum: 1,
        maximum: MAX_CODE_MODE_OUTPUT_TOKENS,
        default: DEFAULT_CODE_MODE_OUTPUT_TOKENS,
    })),
    terminate: Type.Optional(Type.Boolean()),
});
export function registerPublicCodeModeTools(pi, runtime) {
    const tracker = createCodeModeRenderTracker();
    const renderResult = createResultRenderer(runtime, tracker);
    pi.registerTool(createExecTool(runtime, tracker, renderResult));
    pi.registerTool(createWaitTool(runtime, tracker, renderResult));
}
function createExecTool(runtime, tracker, renderResult) {
    return {
        name: "exec",
        label: "Exec",
        description: EXEC_DESCRIPTION,
        promptSnippet: "Compose tools with JavaScript.",
        parameters: EXEC_PARAMETERS,
        async execute(id, params, signal, onUpdate, ctx) {
            tracker.start(id);
            try {
                const response = await (await runtime.getClient()).execute(params.code, { cwd: ctx.cwd, extensionContext: ctx, onUpdate }, signal, runtime.collectTools(ctx));
                tracker.finish(id, response.kind === "yielded" ? "yielded" : "done");
                return toCodeModeToolResult(response);
            }
            catch (error) {
                tracker.finish(id);
                throw error;
            }
        },
        renderCall: ((args, theme, context) => renderExecCall(args, theme, context, tracker, runtime.useRichRendering())),
        renderResult: renderResult,
    };
}
function createWaitTool(runtime, tracker, renderResult) {
    return {
        name: "wait",
        label: "Wait",
        description: WAIT_DESCRIPTION,
        promptSnippet: "Resume or terminate an exec cell.",
        parameters: WAIT_PARAMETERS,
        async execute(id, params, signal, onUpdate, ctx) {
            tracker.start(id);
            try {
                const client = await runtime.getClient();
                const context = { cwd: ctx.cwd, extensionContext: ctx, onUpdate };
                const response = params.terminate
                    ? await client.terminate(params.cell_id, context, signal)
                    : await client.wait(params.cell_id, params.yield_time_ms ?? DEFAULT_WAIT_MS, context, signal);
                tracker.finish(id, response.kind === "yielded" ? "yielded" : "done");
                return toCodeModeToolResult(response, params.max_tokens);
            }
            catch (error) {
                tracker.finish(id);
                throw error;
            }
        },
        renderCall: ((args, theme, context) => renderWaitCall(args, theme, context, tracker, runtime.useRichRendering())),
        renderResult: renderResult,
    };
}
function createResultRenderer(runtime, tracker) {
    return (result, options, theme, context) => renderTrackedCodeModeResult(result, options, theme, context, tracker, runtime.collectRenderTools(), runtime.useRichRendering());
}
