#!/usr/bin/env node
// pi-ask-mcp/server.js
//
// Minimal MCP stdio server that exposes ONE tool: `ask`.
// Bridges Claude → pi via a Unix-domain socket: when Claude calls the tool,
// this server forwards the question(s) to pi over $PI_ASK_SOCKET, awaits
// the user's answer, and returns it as the tool result.
//
// Wire format with Claude (stdin/stdout):  JSON-RPC 2.0 over NDJSON.
// Wire format with pi   (PI_ASK_SOCKET):   NDJSON request/response, see
//                                          ../../shared/pi-ask-bridge.ts.
//
// This file is INTENTIONALLY plain JavaScript (no transpile step, no
// node_modules) — Claude CLI spawns it via `node <path>`. Keep it small,
// dependency-free, and self-contained.

import { connect } from "node:net";
import { randomUUID } from "node:crypto";
import { createInterface } from "node:readline";

// ── Configuration ──────────────────────────────────────────────────────────
const SOCKET = process.env.PI_ASK_SOCKET;
if (!SOCKET) {
	process.stderr.write("[pi-ask-mcp] PI_ASK_SOCKET env var is required\n");
	process.exit(2);
}

const SERVER_INFO      = { name: "pi", version: "0.1.0" };
const PROTOCOL_VERSION = "2024-11-05";
const SOCKET_TIMEOUT_MS = 15 * 60 * 1000; // matches runClaude's default

// ── Tool schema (kept in sync with pi-ask-tool/index.ts AskParamsSchema) ──
const ASK_INPUT_SCHEMA = {
	type: "object",
	required: ["questions"],
	properties: {
		questions: {
			type: "array",
			minItems: 1,
			description: "One or more questions to ask the user.",
			items: {
				type: "object",
				required: ["id", "question", "options"],
				properties: {
					id:       { type: "string", description: "Stable id (e.g. 'auth', 'cache')." },
					question: { type: "string", description: "Question text shown to the user." },
					options: {
						type: "array",
						minItems: 1,
						description: "2-5 concise options. Do NOT include 'Other' (UI adds it).",
						items: {
							type: "object",
							required: ["label"],
							properties: {
								label: { type: "string", description: "Option display label." },
							},
						},
					},
					multi:       { type: "boolean", description: "Allow multi-select. Defaults to false." },
					recommended: { type: "number",  description: "0-indexed recommended option (default highlight)." },
				},
			},
		},
	},
};

const ASK_DESCRIPTION = [
	"Ask the user one or more structured questions through pi's native TUI.",
	"Use this whenever a choice materially affects the outcome — instead of",
	"guessing or the built-in AskUserQuestion. Provide 2-5 concise options.",
	"Set multi=true when multiple answers are valid. Do NOT include an 'Other'",
	"option (UI adds it automatically). The result is a JSON array of",
	"{id, selectedOptions[], customInput?} per question — empty selectedOptions",
	"means the user cancelled.",
].join(" ");

// ── stdio framing: NDJSON ──────────────────────────────────────────────────
const rl = createInterface({ input: process.stdin });
const send = (msg) => process.stdout.write(JSON.stringify(msg) + "\n");
const log  = (msg) => process.stderr.write(`[pi-ask-mcp] ${msg}\n`);

// ── socket round-trip to pi-ask-bridge ─────────────────────────────────────
function askPi(args) {
	return new Promise((resolve, reject) => {
		const sock = connect(SOCKET);
		const id   = randomUUID();
		let buf    = "";
		let settled = false;
		const finish = (fn, val) => { if (settled) return; settled = true; clearTimeout(t); fn(val); try { sock.end(); } catch {} };
		const t = setTimeout(
			() => finish(reject, new Error(`pi-ask bridge timeout after ${SOCKET_TIMEOUT_MS / 1000}s`)),
			SOCKET_TIMEOUT_MS,
		);

		sock.on("connect", () => sock.write(JSON.stringify({ id, type: "ask", ...args }) + "\n"));
		sock.on("data", (d) => {
			buf += d.toString();
			const nl = buf.indexOf("\n");
			if (nl < 0) return;
			try { finish(resolve, JSON.parse(buf.slice(0, nl))); }
			catch (err) { finish(reject, err); }
		});
		sock.on("error", (err) => finish(reject, err));
		sock.on("close", () => {
			if (!settled) finish(reject, new Error("pi-ask bridge closed connection without reply"));
		});
	});
}

// ── JSON-RPC method handlers ───────────────────────────────────────────────
async function handleRequest(req) {
	const { id, method, params } = req;
	try {
		switch (method) {
			case "initialize":
				return ok(id, {
					protocolVersion: PROTOCOL_VERSION,
					capabilities:    { tools: {} },
					serverInfo:      SERVER_INFO,
				});
			case "tools/list":
				return ok(id, {
					tools: [{ name: "ask", description: ASK_DESCRIPTION, inputSchema: ASK_INPUT_SCHEMA }],
				});
			case "tools/call": {
				const name = params?.name;
				const args = params?.arguments ?? {};
				if (name !== "ask") return err(id, -32602, `unknown tool: ${name}`);
				const reply = await askPi(args);
				if (reply.type === "error") {
					return ok(id, {
						isError: true,
						content: [{ type: "text", text: `(user did not answer: ${reply.message})` }],
					});
				}
				return ok(id, {
					content: [{ type: "text", text: JSON.stringify(reply.results, null, 2) }],
				});
			}
			case "ping":            return ok(id, {});
			case "resources/list":  return ok(id, { resources: [] });
			case "prompts/list":    return ok(id, { prompts: [] });
			default:                return err(id, -32601, `method not found: ${method}`);
		}
	} catch (e) {
		return err(id, -32603, e instanceof Error ? e.message : String(e));
	}
}

const ok  = (id, result) => ({ jsonrpc: "2.0", id, result });
const err = (id, code, message) => ({ jsonrpc: "2.0", id, error: { code, message } });

// ── main loop ──────────────────────────────────────────────────────────────
//
// Track in-flight handlers so we don't exit before they finish. Without this,
// `node server.js <<<input` (or any case where stdin closes mid-request) would
// race the async tools/call handler against rl 'close' → process.exit, and
// the reply would silently disappear.
let inflight = 0;
let stdinClosed = false;
function drainAndExit(code = 0) {
	if (inflight === 0) process.exit(code);
}

async function handleOne(msg) {
	// Notifications carry no id and expect no response.
	if (msg.id === undefined || msg.id === null) {
		if (msg.method === "exit") drainAndExit(0);
		return; // notifications/initialized, notifications/cancelled, etc.
	}
	inflight += 1;
	try {
		const reply = await handleRequest(msg);
		if (reply) send(reply);
	} finally {
		inflight -= 1;
		if (stdinClosed) drainAndExit(0);
	}
}

rl.on("line", (line) => {
	if (!line.trim()) return;
	let msg;
	try { msg = JSON.parse(line); } catch { return; }
	if (Array.isArray(msg)) {
		for (const m of msg) void handleOne(m);
	} else {
		void handleOne(msg);
	}
});

rl.on("close",     () => { stdinClosed = true; drainAndExit(0); });
process.on("SIGTERM", () => process.exit(0));
process.on("SIGINT",  () => process.exit(0));
log("ready, socket=" + SOCKET);