import { spawn } from "child_process";
import { createInterface } from "readline";
import { CappedOutputBuffer } from "../../util/output-buffer.js";
import { getOmkResourceSettings } from "../../util/resource-profile.js";
import { getOmkVersionSync } from "../../util/version.js";
import { checkCommand, resolveKimiBin } from "../../util/shell.js";
import { buildSafeKimiChildEnv } from "./runner.js";
import { terminateProcessTree } from "../../util/process-tree.js";
let requestId = 0;
function nextId() {
    return `omk-${++requestId}`;
}
function wireProcTarget(proc) {
    return {
        pid: proc.pid,
        get exitCode() {
            return proc.exitCode ?? null;
        },
        get signalCode() {
            return proc.signalCode ?? null;
        },
        kill(signal) {
            return proc.kill(signal);
        },
        once(event, listener) {
            proc.once(event, listener);
            return undefined;
        },
    };
}
export class KimiWireClient {
    options;
    proc;
    rl;
    pending = new Map();
    // Strongly-typed event handlers
    eventHandlers = [];
    legacyEventHandlers = [];
    // Typed request handlers
    approvalHandlers = [];
    toolCallHandlers = [];
    questionHandlers = [];
    hookHandlers = [];
    genericRequestHandlers = [];
    static MAX_PENDING_RPCS = 64;
    // Initialize response cache
    _initializeResult;
    constructor(options = {}) {
        this.options = options;
    }
    // ─── Server info getters ─────────────────────────────────────
    get serverInfo() {
        return this._initializeResult?.server;
    }
    get slashCommands() {
        return this._initializeResult?.slash_commands;
    }
    get serverCapabilities() {
        return this._initializeResult?.capabilities;
    }
    get registeredTools() {
        return this._initializeResult?.external_tools;
    }
    get hooksInfo() {
        return this._initializeResult?.hooks;
    }
    // ─── Event handlers ──────────────────────────────────────────
    /** Register a strongly-typed event handler */
    onEvent(handler) {
        this.eventHandlers.push(handler);
        return () => {
            this.eventHandlers = this.eventHandlers.filter((h) => h !== handler);
        };
    }
    /** @deprecated Use onEvent with WireEvent instead */
    onLegacyEvent(handler) {
        this.legacyEventHandlers.push(handler);
        return () => {
            this.legacyEventHandlers = this.legacyEventHandlers.filter((h) => h !== handler);
        };
    }
    emitEvent(event) {
        for (const h of this.eventHandlers) {
            try {
                h(event);
            }
            catch {
                // ignore handler errors
            }
        }
    }
    emitLegacyEvent(event) {
        for (const h of this.legacyEventHandlers) {
            try {
                h(event);
            }
            catch {
                // ignore handler errors
            }
        }
    }
    // ─── Request handlers ────────────────────────────────────────
    onApprovalRequest(handler) {
        this.approvalHandlers.push(handler);
        return () => {
            this.approvalHandlers = this.approvalHandlers.filter((h) => h !== handler);
        };
    }
    onToolCallRequest(handler) {
        this.toolCallHandlers.push(handler);
        return () => {
            this.toolCallHandlers = this.toolCallHandlers.filter((h) => h !== handler);
        };
    }
    onQuestionRequest(handler) {
        this.questionHandlers.push(handler);
        return () => {
            this.questionHandlers = this.questionHandlers.filter((h) => h !== handler);
        };
    }
    onHookRequest(handler) {
        this.hookHandlers.push(handler);
        return () => {
            this.hookHandlers = this.hookHandlers.filter((h) => h !== handler);
        };
    }
    /** Generic fallback request handler for unhandled request types */
    onRequest(handler) {
        this.genericRequestHandlers.push(handler);
        return () => {
            this.genericRequestHandlers = this.genericRequestHandlers.filter((h) => h !== handler);
        };
    }
    // ─── Lifecycle ───────────────────────────────────────────────
    async start() {
        const kimiBin = resolveKimiBin();
        const kimiAvailable = await checkCommand(kimiBin);
        if (!kimiAvailable) {
            throw new Error("[omk] `kimi` command not found in PATH. " +
                "Install Kimi CLI first: npm i -g @anthropic-ai/kimi-code\n" +
                "If already installed, check your PATH or set KIMI_BIN env var.");
        }
        const args = ["--wire"];
        if (this.options.agentFile)
            args.push("--agent-file", this.options.agentFile);
        if (this.options.configFile)
            args.push("--config-file", this.options.configFile);
        if (this.options.mcpConfigFile)
            args.push("--mcp-config-file", this.options.mcpConfigFile);
        if (this.options.continue)
            args.push("--continue");
        if (this.options.session)
            args.push("--session", this.options.session);
        if (this.options.model)
            args.push("--model", this.options.model);
        if (this.options.yolo)
            args.push("--yolo");
        const childEnv = buildSafeKimiChildEnv(process.env, this.options.env ?? {}, {}, {
            warnExplicitSecrets: true,
            explicitEnvContext: "Kimi wire client env",
        });
        this.proc = spawn(kimiBin, args, {
            cwd: this.options.cwd,
            env: childEnv,
            stdio: ["pipe", "pipe", "pipe"],
        });
        this.rl = createInterface({ input: this.proc.stdout });
        this.rl.on("line", (line) => {
            const MAX_LINE_LENGTH = 10 * 1024 * 1024; // 10MB
            if (line.length > MAX_LINE_LENGTH)
                return;
            const trimmed = line.trim();
            if (!trimmed)
                return;
            try {
                const msg = JSON.parse(trimmed);
                if ("id" in msg && ("result" in msg || "error" in msg)) {
                    const pending = this.pending.get(msg.id);
                    if (pending) {
                        clearTimeout(pending.timer);
                        this.pending.delete(msg.id);
                        if ("error" in msg && msg.error) {
                            pending.reject(new Error(msg.error.message));
                        }
                        else {
                            pending.resolve(msg);
                        }
                    }
                }
                else if ("method" in msg) {
                    this.handleServerMessage(msg);
                }
            }
            catch {
                // Non-JSON line: emit as raw assistant message for backward compat
                this.emitLegacyEvent({ type: "message", role: "assistant", content: trimmed });
                this.emitEvent({ type: "ContentPart", payload: { type: "text", text: trimmed } });
            }
        });
        this.proc.stderr?.on("data", (data) => {
            const text = data.toString("utf-8");
            this.emitLegacyEvent({ type: "error", message: text });
        });
        this.proc.on("exit", (code) => {
            this.emitLegacyEvent({ type: "error", message: `Kimi process exited with code ${code}` });
            this.rl?.close();
            this.rl = undefined;
            for (const [id, p] of this.pending) {
                clearTimeout(p.timer);
                p.reject(new Error(`Kimi process exited (code ${code}) before response to request ${id}`));
            }
            this.pending.clear();
        });
        let initResult;
        try {
            initResult = await this.call("initialize", {
                protocol_version: "1.7",
                client: { name: "open-multi-agent-kit", version: getOmkVersionSync() },
                capabilities: { supports_question: true, supports_plan_mode: true },
                external_tools: [
                    {
                        name: "omk_claim_task",
                        description: "Receive a DAG node task assignment",
                        parameters: { type: "object", properties: {} },
                    },
                    {
                        name: "omk_update_task",
                        description: "Update task status",
                        parameters: {
                            type: "object",
                            properties: {
                                task_id: { type: "string" },
                                status: { type: "string", enum: ["running", "done", "failed", "blocked"] },
                            },
                            required: ["task_id", "status"],
                        },
                    },
                    {
                        name: "omk_read_memory",
                        description: "Read project memory",
                        parameters: {
                            type: "object",
                            properties: { path: { type: "string" } },
                            required: ["path"],
                        },
                    },
                    {
                        name: "omk_write_memory",
                        description: "Write project memory",
                        parameters: {
                            type: "object",
                            properties: { path: { type: "string" }, content: { type: "string" } },
                            required: ["path", "content"],
                        },
                    },
                    {
                        name: "omk_emit_metric",
                        description: "Record metrics",
                        parameters: {
                            type: "object",
                            properties: { key: { type: "string" }, value: { type: "number" } },
                            required: ["key", "value"],
                        },
                    },
                    {
                        name: "omk_report_blocker",
                        description: "Report blockers",
                        parameters: {
                            type: "object",
                            properties: { reason: { type: "string" } },
                            required: ["reason"],
                        },
                    },
                ],
                hooks: this.options.hooks,
            }, undefined, 30000);
        }
        catch (err) {
            const msg = String(err);
            if (msg.includes("-32601") || msg.includes("method not found")) {
                // No-handshake fallback: proceed without initialize result
                initResult = undefined;
            }
            else {
                throw err;
            }
        }
        this._initializeResult = initResult;
    }
    // ─── JSON-RPC call machinery ─────────────────────────────────
    async call(method, params, signal, timeoutMs) {
        if (!this.proc?.stdin)
            throw new Error("Wire client not started");
        if (this.pending.size >= KimiWireClient.MAX_PENDING_RPCS) {
            throw new Error(`Too many pending RPCs (${this.pending.size}), max is ${KimiWireClient.MAX_PENDING_RPCS}`);
        }
        if (signal?.aborted) {
            throw new Error(`Wire RPC aborted: ${method}`);
        }
        const id = nextId();
        const req = { jsonrpc: "2.0", id, method, params: params };
        const effectiveTimeoutMs = timeoutMs ?? Number(process.env.OMK_WIRE_TIMEOUT_MS || "120000");
        let timer;
        const promise = new Promise((resolve, reject) => {
            timer = setTimeout(() => {
                this.pending.delete(id);
                reject(new Error(`Wire RPC timeout: ${method} (id: ${id}, timeout: ${effectiveTimeoutMs}ms)`));
            }, effectiveTimeoutMs);
            timer.unref();
            this.pending.set(id, { resolve, reject, timer });
        });
        if (signal?.aborted) {
            this.pending.delete(id);
            clearTimeout(timer);
            throw new Error(`Wire RPC aborted: ${method} (id: ${id})`);
        }
        try {
            if (this.proc.stdin.destroyed || this.proc.stdin.writableEnded) {
                throw new Error("Kimi process stdin is closed");
            }
            this.proc.stdin.write(JSON.stringify(req) + "\n");
        }
        catch (err) {
            this.pending.delete(id);
            clearTimeout(timer);
            throw err;
        }
        const res = await (signal
            ? Promise.race([
                promise,
                new Promise((_, reject) => {
                    const onAbort = () => {
                        this.pending.delete(id);
                        clearTimeout(timer);
                        reject(new Error(`Wire RPC aborted: ${method} (id: ${id})`));
                    };
                    if (signal.aborted) {
                        onAbort();
                    }
                    else {
                        signal.addEventListener("abort", onAbort, { once: true });
                    }
                }),
            ])
            : promise);
        if ("error" in res && res.error)
            throw new Error(res.error.message);
        return res.result;
    }
    // ─── Server message dispatch ─────────────────────────────────
    handleServerMessage(msg) {
        // event notification (no id, or we treat it as notification)
        if (msg.method === "event") {
            const params = msg.params;
            if (!params)
                return;
            this.dispatchEvent(params.type, params.payload);
            return;
        }
        // request from agent (requires response)
        if (msg.method === "request") {
            const params = msg.params;
            if (!params)
                return;
            this.dispatchRequest(msg.id, params.type, params.payload);
            return;
        }
        // Legacy: direct "status" method (pre-v1.7 servers may send this)
        if (msg.method === "status") {
            const p = msg.params;
            const statusUpdate = {
                context_usage: typeof p.context_usage === "number" ? p.context_usage : null,
                context_tokens: typeof p.context_tokens === "number" ? p.context_tokens : null,
                max_context_tokens: typeof p.max_context_tokens === "number" ? p.max_context_tokens : null,
                token_usage: typeof p.token_usage === "object" && p.token_usage !== null
                    ? p.token_usage
                    : null,
                message_id: typeof p.message_id === "string" ? p.message_id : null,
                plan_mode: typeof p.plan_mode === "boolean" ? p.plan_mode : null,
            };
            this.emitEvent({ type: "StatusUpdate", payload: statusUpdate });
            // Also emit legacy status event
            this.emitLegacyEvent({
                type: "status",
                contextUsage: Number(p.context_usage ?? p.context_usage ?? 0),
                maxContextTokens: Number(p.max_context_tokens ?? 256000),
                tokenUsage: Number(p.token_usage ? p.token_usage.output : 0),
                planMode: Boolean(p.plan_mode ?? false),
            });
            return;
        }
    }
    dispatchEvent(eventType, payload) {
        switch (eventType) {
            case "TurnBegin": {
                const p = payload;
                this.emitEvent({ type: "TurnBegin", payload: p });
                break;
            }
            case "TurnEnd": {
                this.emitEvent({ type: "TurnEnd" });
                break;
            }
            case "StepBegin": {
                const p = payload;
                this.emitEvent({ type: "StepBegin", payload: p });
                break;
            }
            case "StepInterrupted": {
                this.emitEvent({ type: "StepInterrupted" });
                break;
            }
            case "CompactionBegin": {
                this.emitEvent({ type: "CompactionBegin" });
                break;
            }
            case "CompactionEnd": {
                this.emitEvent({ type: "CompactionEnd" });
                break;
            }
            case "StatusUpdate": {
                const p = payload;
                this.emitEvent({ type: "StatusUpdate", payload: p });
                // Legacy emit
                this.emitLegacyEvent({
                    type: "status",
                    contextUsage: p.context_usage ?? 0,
                    maxContextTokens: p.max_context_tokens ?? 256000,
                    tokenUsage: p.token_usage?.output ?? 0,
                    planMode: p.plan_mode ?? false,
                });
                break;
            }
            case "ContentPart": {
                const part = payload;
                this.emitEvent({ type: "ContentPart", payload: part });
                if (part.type === "text") {
                    this.emitLegacyEvent({ type: "message", role: "assistant", content: part.text });
                }
                break;
            }
            case "ToolCall": {
                const p = payload;
                this.emitEvent({ type: "ToolCall", payload: p });
                this.emitLegacyEvent({ type: "tool_use", name: p.function.name, input: p.function.arguments });
                break;
            }
            case "ToolCallPart": {
                this.emitEvent({ type: "ToolCallPart", payload: payload });
                break;
            }
            case "ToolResult": {
                const p = payload;
                this.emitEvent({ type: "ToolResult", payload: p });
                this.emitLegacyEvent({
                    type: "tool_result",
                    name: p.tool_call_id,
                    output: p.return_value.output,
                });
                break;
            }
            case "ApprovalResponse": {
                this.emitEvent({ type: "ApprovalResponse", payload: payload });
                break;
            }
            case "SubagentEvent": {
                this.emitEvent({ type: "SubagentEvent", payload: payload });
                break;
            }
            case "SteerInput": {
                this.emitEvent({ type: "SteerInput", payload: payload });
                break;
            }
            case "PlanDisplay": {
                this.emitEvent({ type: "PlanDisplay", payload: payload });
                break;
            }
            case "HookTriggered": {
                this.emitEvent({ type: "HookTriggered", payload: payload });
                break;
            }
            case "HookResolved": {
                this.emitEvent({ type: "HookResolved", payload: payload });
                break;
            }
            default: {
                // Unknown event type — ignore or log
                break;
            }
        }
    }
    dispatchRequest(rpcId, requestType, payload) {
        let responded = false;
        const respond = (result) => {
            if (responded)
                return;
            responded = true;
            const res = { jsonrpc: "2.0", id: rpcId, result };
            this.proc?.stdin?.write(JSON.stringify(res) + "\n");
        };
        const reject = (error) => {
            if (responded)
                return;
            responded = true;
            const res = { jsonrpc: "2.0", id: rpcId, error };
            this.proc?.stdin?.write(JSON.stringify(res) + "\n");
        };
        switch (requestType) {
            case "ApprovalRequest": {
                const req = payload;
                const wireReq = { type: "ApprovalRequest", payload: req };
                if (this.approvalHandlers.length > 0) {
                    for (const h of this.approvalHandlers) {
                        try {
                            h(req, respond, reject);
                        }
                        catch {
                            // handler error
                        }
                    }
                }
                else {
                    for (const h of this.genericRequestHandlers) {
                        try {
                            h(wireReq, respond, reject);
                        }
                        catch {
                            // handler error
                        }
                    }
                }
                // Also emit legacy request event
                this.emitLegacyEvent({ type: "request", method: requestType, params: payload, respond, reject });
                break;
            }
            case "ToolCallRequest": {
                const req = payload;
                const wireReq = { type: "ToolCallRequest", payload: req };
                if (this.toolCallHandlers.length > 0) {
                    for (const h of this.toolCallHandlers) {
                        try {
                            h(req, respond, reject);
                        }
                        catch {
                            // handler error
                        }
                    }
                }
                else {
                    for (const h of this.genericRequestHandlers) {
                        try {
                            h(wireReq, respond, reject);
                        }
                        catch {
                            // handler error
                        }
                    }
                }
                this.emitLegacyEvent({ type: "request", method: requestType, params: payload, respond, reject });
                break;
            }
            case "QuestionRequest": {
                const req = payload;
                const wireReq = { type: "QuestionRequest", payload: req };
                if (this.questionHandlers.length > 0) {
                    for (const h of this.questionHandlers) {
                        try {
                            h(req, respond, reject);
                        }
                        catch {
                            // handler error
                        }
                    }
                }
                else {
                    for (const h of this.genericRequestHandlers) {
                        try {
                            h(wireReq, respond, reject);
                        }
                        catch {
                            // handler error
                        }
                    }
                }
                this.emitLegacyEvent({ type: "request", method: requestType, params: payload, respond, reject });
                break;
            }
            case "HookRequest": {
                const req = payload;
                const wireReq = { type: "HookRequest", payload: req };
                if (this.hookHandlers.length > 0) {
                    for (const h of this.hookHandlers) {
                        try {
                            h(req, respond, reject);
                        }
                        catch {
                            // handler error
                        }
                    }
                }
                else {
                    for (const h of this.genericRequestHandlers) {
                        try {
                            h(wireReq, respond, reject);
                        }
                        catch {
                            // handler error
                        }
                    }
                }
                this.emitLegacyEvent({ type: "request", method: requestType, params: payload, respond, reject });
                break;
            }
            default: {
                const wireReq = { type: requestType, payload: payload };
                for (const h of this.genericRequestHandlers) {
                    try {
                        h(wireReq, respond, reject);
                    }
                    catch {
                        // handler error
                    }
                }
                this.emitLegacyEvent({ type: "request", method: requestType, params: payload, respond, reject });
            }
        }
    }
    // ─── Public RPC methods ──────────────────────────────────────
    async prompt(userInput) {
        return this.call("prompt", { user_input: userInput });
    }
    async steer(userInput) {
        return this.call("steer", { user_input: userInput });
    }
    async cancel() {
        return this.call("cancel", {});
    }
    async replay() {
        return this.call("replay", {});
    }
    async setPlanMode(enabled) {
        return this.call("set_plan_mode", { enabled });
    }
    // ─── Shutdown ────────────────────────────────────────────────
    async stop() {
        for (const [id, p] of this.pending) {
            clearTimeout(p.timer);
            p.reject(new Error(`Wire client stopped before response to request ${id}`));
        }
        this.pending.clear();
        this.proc?.removeAllListeners();
        this.proc?.stderr?.removeAllListeners("data");
        this.rl?.removeAllListeners();
        this.eventHandlers = [];
        this.legacyEventHandlers = [];
        this.approvalHandlers = [];
        this.toolCallHandlers = [];
        this.questionHandlers = [];
        this.hookHandlers = [];
        this.genericRequestHandlers = [];
        if (this.rl) {
            this.rl.close();
            this.rl = undefined;
        }
        if (this.proc) {
            const proc = this.proc;
            this.proc = undefined;
            await terminateProcessTree(wireProcTarget(proc), { graceMs: 1000, waitMs: 5000 });
        }
    }
}
// Re-export all wire protocol types for convenience
export * from "./wire-protocol-types.js";
// ─── Task runner factory ───────────────────────────────────────
export function createKimiTaskRunner(client) {
    return {
        async run(node, _env) {
            const resources = await getOmkResourceSettings();
            const prompt = `[${node.role}] ${node.name}`;
            const stdout = new CappedOutputBuffer(resources.wireOutputBytes, "wire stdout");
            const stderr = new CappedOutputBuffer(resources.wireOutputBytes, "wire stderr");
            const offEvent = client.onEvent((event) => {
                switch (event.type) {
                    case "ContentPart": {
                        if (event.payload.type === "text") {
                            stdout.append(`${event.payload.text}\n`);
                        }
                        break;
                    }
                    case "ToolResult": {
                        const output = event.payload.return_value.output;
                        const text = typeof output === "string" ? output : JSON.stringify(output);
                        stdout.append(`${text}\n`);
                        break;
                    }
                }
            });
            const offLegacy = client.onLegacyEvent((event) => {
                if (event.type === "error")
                    stderr.append(`${event.message}\n`);
            });
            try {
                const result = await client.prompt(prompt);
                const success = result.status === "finished";
                return {
                    success,
                    stdout: stdout.toString(),
                    stderr: stderr.toString(),
                };
            }
            catch (err) {
                stderr.append(`\n${String(err)}`);
                return {
                    success: false,
                    stdout: stdout.toString(),
                    stderr: stderr.toString(),
                };
            }
            finally {
                offEvent();
                offLegacy();
            }
        },
    };
}
