import process from "process";
import { runWithPolicy } from "./interruptLoop.js";
import { errorMessage } from "../util.js";
function formatToolDescription(description, interruptKinds) {
    if (interruptKinds.length === 0)
        return description;
    const kinds = interruptKinds.map((ik) => ik.kind).join(", ");
    return `${description}\n\nInterrupt kinds: ${kinds}`;
}
const MCP_PROTOCOL_VERSION = "2025-06-18";
function success(id, result) {
    return { jsonrpc: "2.0", id, result };
}
function rpcError(id, code, message) {
    return { jsonrpc: "2.0", id, error: { code, message } };
}
function schemaToJsonSchema(schema) {
    const s = schema;
    return s && typeof s.toJSONSchema === "function"
        ? s.toJSONSchema()
        : { type: "object", properties: {} };
}
const POLICY_TOOL_NAMES = {
    GET: "agencyGetPolicy",
    ADD_RULE: "agencyAddRule",
    REMOVE_RULE: "agencyRemoveRule",
    CLEAR: "agencyClearPolicy",
};
const POLICY_TOOL_DEFINITIONS = [
    {
        name: POLICY_TOOL_NAMES.GET,
        description: "Get the current interrupt policy for this agent. Returns a JSON object keyed by interrupt kind, where each kind maps to an ordered array of rules.",
        inputSchema: { type: "object", properties: {} },
    },
    {
        name: POLICY_TOOL_NAMES.ADD_RULE,
        description: "Add a rule to the interrupt policy. Rules control which actions the agent can take autonomously. Each tool lists its interrupt kinds — use those as the 'kind' parameter. Rules are evaluated in order; the first match wins. A rule with no 'match' field is a catch-all.",
        inputSchema: {
            type: "object",
            properties: {
                kind: { type: "string", description: "The interrupt kind to add a rule for (e.g. 'email::send')" },
                action: { type: "string", enum: ["approve", "reject"], description: "What to do when this rule matches" },
                match: { type: "object", additionalProperties: { type: "string" }, description: "Optional. Keys are interrupt data field names, values are glob patterns (e.g. '*@company.com'). If omitted, the rule is a catch-all." },
            },
            required: ["kind", "action"],
        },
    },
    {
        name: POLICY_TOOL_NAMES.REMOVE_RULE,
        description: "Remove a rule from the interrupt policy by index. Use agencyGetPolicy to see current rules and their indices.",
        inputSchema: {
            type: "object",
            properties: {
                kind: { type: "string", description: "The interrupt kind to remove a rule from" },
                ruleIndex: { type: "integer", minimum: 0, description: "Zero-based index of the rule to remove" },
            },
            required: ["kind", "ruleIndex"],
        },
    },
    {
        name: POLICY_TOOL_NAMES.CLEAR,
        description: "Clear the entire interrupt policy, resetting to reject-all. After clearing, all interrupt-producing actions will be rejected until new rules are added.",
        inputSchema: { type: "object", properties: {} },
    },
];
function handlePolicyTool(name, args, policyStore) {
    switch (name) {
        case POLICY_TOOL_NAMES.GET:
            return { content: [{ type: "text", text: JSON.stringify(policyStore.get(), null, 2) }], isError: false };
        case POLICY_TOOL_NAMES.ADD_RULE:
            try {
                policyStore.addRule(args.kind, { action: args.action, ...(args.match && { match: args.match }) });
                return { content: [{ type: "text", text: `Rule added for '${args.kind}'.` }], isError: false };
            }
            catch (err) {
                return { content: [{ type: "text", text: errorMessage(err) }], isError: true };
            }
        case POLICY_TOOL_NAMES.REMOVE_RULE:
            try {
                policyStore.removeRule(args.kind, args.ruleIndex);
                return { content: [{ type: "text", text: `Rule removed from '${args.kind}'.` }], isError: false };
            }
            catch (err) {
                return { content: [{ type: "text", text: errorMessage(err) }], isError: true };
            }
        case POLICY_TOOL_NAMES.CLEAR:
            try {
                policyStore.clear();
                return { content: [{ type: "text", text: "Policy cleared." }], isError: false };
            }
            catch (err) {
                return { content: [{ type: "text", text: errorMessage(err) }], isError: true };
            }
        default:
            return null;
    }
}
/**
 * Run a tool invocation through the policy store (if configured) and wrap
 * the result/error into the MCP `tools/call` response shape. `extractData`
 * lets callers unwrap node results that are `{ data, ... }` envelopes
 * before serializing.
 */
async function runToolInvocation(id, invoke, policyConfig, extractData = (r) => r) {
    try {
        const result = policyConfig
            ? await runWithPolicy(invoke, policyConfig.policyStore, policyConfig.interruptHandlers)
            : await invoke();
        return success(id, {
            content: [{ type: "text", text: JSON.stringify(extractData(result), null, 2) }],
            isError: false,
        });
    }
    catch (err) {
        return success(id, {
            content: [{ type: "text", text: errorMessage(err) }],
            isError: true,
        });
    }
}
async function handleToolCall(message, functionsByName, nodesByName, policyConfig) {
    const name = message.params?.name;
    const args = message.params?.arguments ?? {};
    const id = message.id ?? null;
    if (policyConfig) {
        const policyResult = handlePolicyTool(name, args, policyConfig.policyStore);
        if (policyResult)
            return success(id, policyResult);
    }
    const fn = functionsByName[name];
    if (fn) {
        return runToolInvocation(id, () => fn.agencyFunction.invoke({ type: "named", positionalArgs: [], namedArgs: args }), policyConfig);
    }
    const node = nodesByName[name];
    if (node) {
        const positional = node.parameters.map((p) => args[p.name]);
        return runToolInvocation(id, () => node.invoke(...positional), policyConfig, (r) => (r && typeof r === "object" && "data" in r ? r.data : r));
    }
    return rpcError(id, -32602, `Unknown tool '${name}'`);
}
export function createMcpHandler(config) {
    const { serverName, serverVersion, exports } = config;
    const functions = exports.filter((e) => e.kind === "function");
    const nodes = exports.filter((e) => e.kind === "node");
    const functionsByName = Object.fromEntries(functions.map((f) => [f.name, f]));
    const nodesByName = Object.fromEntries(nodes.map((n) => [n.name, n]));
    const functionToolEntries = functions.map((f) => ({
        name: f.name,
        description: formatToolDescription(f.description, f.interruptKinds),
        inputSchema: schemaToJsonSchema(f.agencyFunction.toolDefinition?.schema),
        ...(f.agencyFunction.safe ? { annotations: { readOnlyHint: true } } : {}),
    }));
    const nodeToolEntries = nodes.map((n) => ({
        name: n.name,
        description: formatToolDescription(`Run the '${n.name}' node`, n.interruptKinds),
        inputSchema: {
            type: "object",
            properties: Object.fromEntries(n.parameters.map((p) => [p.name, { type: "string" }])),
            required: n.parameters.map((p) => p.name),
        },
    }));
    const toolsListPayload = [...functionToolEntries, ...nodeToolEntries];
    const { policyConfig } = config;
    if (policyConfig) {
        toolsListPayload.push(...POLICY_TOOL_DEFINITIONS);
    }
    return async (message) => {
        if (message.jsonrpc !== "2.0") {
            return rpcError(message.id ?? null, -32600, "Expected JSON-RPC 2.0 message");
        }
        switch (message.method) {
            case "initialize":
                return success(message.id ?? null, {
                    protocolVersion: MCP_PROTOCOL_VERSION,
                    capabilities: { tools: { listChanged: false } },
                    serverInfo: { name: serverName, version: serverVersion },
                });
            case "notifications/initialized":
                return null;
            case "ping":
                return success(message.id ?? null, {});
            case "tools/list":
                return success(message.id ?? null, { tools: toolsListPayload });
            case "tools/call":
                return handleToolCall(message, functionsByName, nodesByName, policyConfig);
            case "shutdown":
                return success(message.id ?? null, {});
            case "exit":
                process.exit(0);
            default:
                if (message.id !== undefined) {
                    return rpcError(message.id, -32601, `Method not found: ${message.method}`);
                }
                return null;
        }
    };
}
function sendResponse(response) {
    if (response) {
        process.stdout.write(`${JSON.stringify(response)}\n`);
    }
}
function processLine(line, handler) {
    try {
        handler(JSON.parse(line))
            .then(sendResponse)
            .catch((err) => {
            sendResponse(rpcError(null, -32603, `Handler error: ${errorMessage(err)}`));
        });
    }
    catch (err) {
        sendResponse(rpcError(null, -32700, `Invalid JSON: ${errorMessage(err)}`));
    }
}
/**
 * Maximum size of an unterminated stdio line. A peer that streams data
 * without ever sending a newline would otherwise grow the buffer
 * unboundedly and OOM the process.
 */
const MAX_STDIO_LINE_BYTES = 10 * 1024 * 1024; // 10 MB
export function startStdioServer(handler) {
    process.stdin.setEncoding("utf-8");
    let buffer = "";
    process.stdin.on("data", (chunk) => {
        buffer += chunk;
        let newlineIndex = buffer.indexOf("\n");
        while (newlineIndex !== -1) {
            const line = buffer.slice(0, newlineIndex).trim();
            buffer = buffer.slice(newlineIndex + 1);
            if (line.length > 0)
                processLine(line, handler);
            newlineIndex = buffer.indexOf("\n");
        }
        const bufferBytes = Buffer.byteLength(buffer, "utf-8");
        if (bufferBytes > MAX_STDIO_LINE_BYTES) {
            // Drop the oversized partial and continue reading. We can't reply
            // because we have no message id; surface the problem to stderr.
            process.stderr.write(`MCP stdio: discarding partial input of ${bufferBytes} bytes (no newline within ${MAX_STDIO_LINE_BYTES} bytes)\n`);
            buffer = "";
        }
    });
}
