/**
 * OMP session control socket.
 *
 * Gives an outside process a way to steer and interrupt a live OMP session it
 * does not own. The bridge can already observe an external session by tailing
 * its JSONL, but observation is one-way: a thread another client is driving is
 * read-only from ChatGPT.
 *
 * This is an ordinary OMP extension rather than a patch, so it survives omp
 * upgrades and loads into every session - including hand-started terminal TUI
 * ones - as soon as it is listed under `extensions` in the OMP config.
 *
 * Each session binds a Unix socket named after its session id, which is the
 * same identifier the bridge already resolves for a thread, so no extra
 * discovery handshake is needed.
 *
 * Protocol: newline-delimited JSON, one request per line, one reply per line.
 *   -> { "op": "status" }
 *   <- { "ok": true, "idle": true, "sessionId": "...", "sessionFile": "...", "pid": 123 }
 *   -> { "op": "subscribe", "streamId": "...", "afterSequence": 12 }
 *   <- { "ok": true, "op": "subscribe", "streamId": "...", "sequence": 14, "earliestSequence": 1, "reset": false, "status": { ... } }
 *   <- { "type": "event", "streamId": "...", "sequence": 13, "event": { ... }, "status": { ... } }
 *   -> { "op": "steer", "text": "actually, do X instead" }
 *   -> { "op": "followUp", "text": "when you're done, also do Y" }
 *   -> { "op": "interrupt" }
 *
 * Anything able to open this socket can inject instructions into a running
 * agent, so the directory is 0700 and each socket is 0600 - owner only.
 */
import crypto from "node:crypto";
import fs from "node:fs";
import fsp from "node:fs/promises";
import net from "node:net";
import os from "node:os";
import path from "node:path";

/** Resolved per bind so an override applies to sessions started after it. */
function controlDir() {
  return process.env.OMP_CONTROL_SOCKET_DIR
    || path.join(os.homedir(), ".omp", "codex-bridge", "control");
}

/** Cap on a single request line; a control message is tiny by construction. */
const MAX_LINE_BYTES = 64 * 1024;

/** Frame serialization cap; oversized frames degrade to invalidation events. */
const MAX_FRAME_BYTES = 64 * 1024;

/** Bounded replay ring capacity. */
const MAX_REPLAY_CAPACITY = 256;

/** Maximum queued bytes before disconnecting a slow subscriber. */
const MAX_SUBSCRIBER_BUFFER = 512 * 1024;

/** One bridge needs one stream; cap accidental/malicious fan-out. */
const MAX_SUBSCRIBERS = 8;

const EXTERNAL_ASK_TIMEOUT_MS = 10 * 60_000;


export function controlSocketPath(sessionId) {
  return path.join(controlDir(), `${sessionId}.sock`);
}

/**
 * Remove a socket left behind by a crashed session.
 *
 * A path that still has a listener belongs to a live session and is left
 * alone; only a refused connection proves the owner is gone.
 */
async function clearStaleSocket(socketPath) {
  if (!fs.existsSync(socketPath)) return true;
  const reachable = await new Promise((resolve) => {
    const probe = net.connect(socketPath);
    const done = (value) => {
      probe.destroy();
      resolve(value);
    };
    probe.once("connect", () => done(true));
    probe.once("error", () => done(false));
    setTimeout(() => done(false), 250).unref?.();
  });
  if (reachable) return false;
  await fsp.rm(socketPath, { force: true });
  return true;
}

const LIFECYCLE_EVENTS = [
  "turn_start",
  "turn_end",
  "agent_start",
  "agent_end",
  "message_start",
  "message_update",
  "message_end",
  "tool_execution_start",
  "tool_execution_update",
  "tool_execution_end",
  "session_compact",
  "auto_compaction_start",
  "auto_compaction_end",
  "compaction_start",
  "compaction_end",
  "goal_updated",
  "goal_update",
  "model_change",
  "model_changed",
  "thinking_level_change",
  "reasoning_effort_change",
  "retry_start",
  "retry_end",
  "auto_retry_start",
  "auto_retry_end",
  "retry_fallback_applied",
  "retry_fallback_succeeded",
  "session_pause",
  "session_resume",
  "session_stop",
];

export default function controlSocketExtension(api) {
  let streamId = crypto.randomUUID();
  let sequence = 0;
  const replayBuffer = [];
  const subscribers = new Set();
  const pendingAsks = new Map();
  let askCounter = 0;
  let server = null;
  let boundPath = null;
  let heartbeatTimer = null;
  let currentCtx = null;

  const getStatus = (ctx) => {
    const activeCtx = ctx || currentCtx;
    return {
      idle: Boolean(activeCtx?.isIdle?.()),
      hasPendingMessages: Boolean(activeCtx?.hasPendingMessages?.()),
      sessionId: activeCtx?.sessionManager?.getSessionId?.() ?? null,
      sessionFile: activeCtx?.sessionManager?.getSessionFile?.() ?? null,
      cwd: activeCtx?.cwd ?? null,
      model: activeCtx?.model?.id ?? null,
      pid: process.pid,
    };
  };

  const publishEvent = (event, ctx) => {
    if (!event || typeof event !== "object") {
      event = { type: String(event || "unknown") };
    }
    sequence += 1;
    const status = getStatus(ctx);
    const frame = {
      type: "event",
      streamId,
      sequence,
      event,
      status,
    };

    let jsonLine;
    try {
      const serialized = JSON.stringify(frame);
      if (Buffer.byteLength(serialized, "utf8") <= MAX_FRAME_BYTES) {
        jsonLine = serialized;
      }
    } catch {
      // Serialization error (e.g. circular references)
    }

    let finalFrame = frame;
    if (!jsonLine) {
      finalFrame = {
        type: "event",
        streamId,
        sequence,
        event: {
          type: "invalidation",
          originalType: typeof event.type === "string" ? event.type : "unknown",
          reason: "oversized_or_unserializable",
        },
        status,
      };
      jsonLine = JSON.stringify(finalFrame);
    }

    replayBuffer.push(finalFrame);
    if (replayBuffer.length > MAX_REPLAY_CAPACITY) {
      replayBuffer.shift();
    }

    for (const socket of subscribers) {
      if (socket.destroyed || socket.writableEnded || !socket.writable) {
        subscribers.delete(socket);
        continue;
      }
      if (socket.bufferSize > MAX_SUBSCRIBER_BUFFER) {
        subscribers.delete(socket);
        socket.destroy(new Error("Slow subscriber"));
        continue;
      }
      const ok = socket.write(`${jsonLine}\n`);
      if (!ok && socket.bufferSize > MAX_SUBSCRIBER_BUFFER) {
        subscribers.delete(socket);
        socket.destroy(new Error("Slow subscriber"));
      }
    }
  };

  const askParameters = {
    type: "object",
    additionalProperties: false,
    required: ["questions"],
    properties: {
      questions: {
        type: "array",
        minItems: 1,
        maxItems: 5,
        items: {
          type: "object",
          additionalProperties: false,
          required: ["id", "question", "options"],
          properties: {
            id: { type: "string" },
            header: { type: "string" },
            question: { type: "string" },
            multi: { type: "boolean" },
            recommended: { type: "number" },
            options: {
              type: "array",
              items: {
                type: "object",
                additionalProperties: false,
                required: ["label"],
                properties: {
                  label: { type: "string" },
                  description: { type: "string" },
                  preview: { type: "string" },
                },
              },
            },
          },
        },
      },
    },
  };
  const Type = api.typebox;
  const askToolParameters = Type?.Object
    ? Type.Object({
        questions: Type.Array(Type.Object({
          id: Type.String(),
          header: Type.Optional(Type.String()),
          question: Type.String(),
          multi: Type.Optional(Type.Boolean()),
          recommended: Type.Optional(Type.Number()),
          options: Type.Array(Type.Object({
            label: Type.String(),
            description: Type.Optional(Type.String()),
            preview: Type.Optional(Type.String()),
          })),
        }), { minItems: 1, maxItems: 5 }),
      }, { additionalProperties: false })
    : askParameters;

  const askResult = (questions, response) => {
    if (response?.cancelled) {
      return {
        content: [{ type: "text", text: "Ask input was cancelled" }],
        details: { results: [] },
        isError: true,
      };
    }
    const answers = response?.answers && typeof response.answers === "object" ? response.answers : {};
    const results = questions.map((question) => {
      const selectedOptions = Array.isArray(answers[question.id]?.answers)
        ? answers[question.id].answers.map(String)
        : [];
      const optionLabels = (question.options || []).map((option) => String(option.label || ""));
      const customInput = selectedOptions.find((answer) => !optionLabels.includes(answer));
      return {
        id: question.id,
        question: question.question,
        options: optionLabels,
        multi: Boolean(question.multi),
        selectedOptions: customInput ? selectedOptions.filter((answer) => answer !== customInput) : selectedOptions,
        ...(customInput ? { customInput } : {}),
      };
    });
    let text;
    if (results.length === 1) {
      const result = results[0];
      text = result.customInput
        ? `User provided custom input: ${result.customInput}`
        : `User selected: ${result.selectedOptions.join(", ")}`;
    } else {
      text = `User answers:\n${results.map((result) => {
        const answer = result.customInput || result.selectedOptions.join(", ");
        return `${result.id}: ${JSON.stringify(answer)}`;
      }).join("\n")}`;
    }
    const first = results[0];
    return {
      content: [{ type: "text", text }],
      details: results.length === 1
        ? {
            question: first.question,
            options: first.options,
            multi: first.multi,
            selectedOptions: first.selectedOptions,
            ...(first.customInput ? { customInput: first.customInput } : {}),
          }
        : { results },
      isError: false,
    };
  };

  const executeRelayedAsk = async (toolCallId, params, signal, ctx) => {
    const requestId = `ask:${streamId}:${++askCounter}`;
    let timer = null;
    let abortHandler = null;
    const response = await new Promise((resolve, reject) => {
      const settle = (value, error = null) => {
        clearTimeout(timer);
        signal?.removeEventListener("abort", abortHandler);
        pendingAsks.delete(requestId);
        if (error) reject(error);
        else resolve(value);
      };
      abortHandler = () => settle(null, signal.reason instanceof Error ? signal.reason : new Error("Ask interrupted"));
      signal?.addEventListener("abort", abortHandler, { once: true });
      timer = setTimeout(() => settle({ cancelled: true }), EXTERNAL_ASK_TIMEOUT_MS);
      timer.unref?.();
      pendingAsks.set(requestId, {
        requestId,
        toolCallId,
        questions: params.questions,
        createdAt: Date.now(),
        settle,
      });
      publishEvent({ type: "ask_request", requestId, toolCallId, questions: params.questions }, ctx);
    });
    return askResult(params.questions, response);
  };

  if (typeof api.registerTool === "function") {
    api.registerTool({
      name: "ask",
      label: "Ask",
      description: "Ask the user one or more clarifying questions.",
      parameters: askToolParameters,
      approval: "read",
      strict: true,
      async execute(toolCallId, params, signal, onUpdate, ctx) {
        if (subscribers.size === 0) {
          if (typeof ctx?.invokeTool === "function") {
            return ctx.invokeTool(params, { signal, onUpdate });
          }
          return askResult(params.questions, { cancelled: true });
        }
        return executeRelayedAsk(toolCallId, params, signal, ctx);
      },
    });
  }

  for (const eventName of LIFECYCLE_EVENTS) {
    try {
      api.on?.(eventName, (event, eventCtx) => {
        publishEvent(
          event && typeof event === "object"
            ? (event.type ? event : { type: eventName, ...event })
            : { type: eventName },
          eventCtx || currentCtx
        );
      });
    } catch {
      // Ignore if api.on fails for any event name
    }
  }

  const respond = (socket, payload) => {
    if (socket.destroyed || socket.writableEnded || !socket.writable) return false;
    try {
      const line = JSON.stringify(payload);
      return socket.write(`${line}\n`);
    } catch {
      return false;
    }
  };

  const handleRequest = async (request, ctx) => {
    const op = String(request?.op ?? "");
    switch (op) {
      case "status":
        return {
          ok: true,
          op,
          ...getStatus(ctx),
        };
      case "subscribe": {
        const hasStreamId = typeof request.streamId === "string" && request.streamId.length > 0;
        const hasAfterSeq = typeof request.afterSequence === "number" && Number.isInteger(request.afterSequence) && request.afterSequence >= 0;
        const isReconnecting = hasStreamId && hasAfterSeq;

        let reset = false;
        let replayFrames = [];
        const earliestSequence = replayBuffer.length > 0 ? replayBuffer[0].sequence : (sequence > 0 ? sequence + 1 : 1);

        if (isReconnecting) {
          if (request.streamId !== streamId) {
            reset = true;
          } else if (request.afterSequence > sequence) {
            reset = true;
          } else if (request.afterSequence === sequence) {
            reset = false;
            replayFrames = [];
          } else {
            // request.afterSequence < sequence
            if (request.afterSequence >= earliestSequence - 1) {
              reset = false;
              replayFrames = replayBuffer.filter((f) => f.sequence > request.afterSequence);
            } else {
              reset = true;
              replayFrames = [];
            }
          }
        } else if (hasStreamId && request.streamId !== streamId) {
          reset = true;
        } else {
          // A first-time subscriber has no replay cursor. It must hydrate a
          // durable snapshot after the ACK to close the seed-to-subscribe race.
          reset = true;
        }

        const ack = {
          ok: true,
          op: "subscribe",
          streamId,
          sequence,
          earliestSequence,
          reset,
          status: getStatus(ctx),
          pendingAsks: Array.from(pendingAsks.values(), (pending) => ({
            requestId: pending.requestId,
            toolCallId: pending.toolCallId,
            questions: pending.questions,
            createdAt: pending.createdAt,
          })),
        };

        return {
          ...ack,
          __replayFrames: replayFrames,
        };
      }
      case "configure": {
        if (request.model) {
          const selector = String(request.model);
          const available = ctx.modelRegistry?.getAvailable?.() || [];
          let model = available.find((candidate) =>
            candidate?.id === selector || candidate?.selector === selector
          );
          if (!model) {
            const separator = selector.indexOf("/");
            const provider = separator > 0 ? selector.slice(0, separator) : null;
            const modelId = separator > 0 ? selector.slice(separator + 1) : selector;
            model = provider ? ctx.modelRegistry?.find?.(provider, modelId) : null;
          }
          if (!model) return { ok: false, op, error: `Unknown model: ${selector}` };
          const changed = await api.setModel(model);
          if (changed === false) return { ok: false, op, error: `Model unavailable: ${selector}` };
        }
        if (request.reasoningEffort) {
          await api.setThinkingLevel(String(request.reasoningEffort));
        }
        return {
          ok: true,
          op,
          model: ctx.model?.id ?? null,
          reasoningEffort: request.reasoningEffort || null,
        };
      }
      case "steer":
      case "followUp": {
        const text = typeof request.text === "string" ? request.text : "";
        if (!text.trim()) return { ok: false, op, error: "text is required" };
        // Idle sessions start a turn; streaming ones queue at the next
        // tool/step boundary, which is what makes this safe to call blind.
        api.sendUserMessage(text, { deliverAs: op === "steer" ? "steer" : "followUp" });
        return { ok: true, op, delivered: op, idle: ctx.isIdle() };
      }
      case "ask_response": {
        const pending = pendingAsks.get(String(request.requestId || ""));
        if (!pending) return { ok: false, op, error: "Ask request is no longer pending" };
        pending.settle({ answers: request.answers || {} });
        publishEvent({ type: "ask_resolved", requestId: pending.requestId, toolCallId: pending.toolCallId }, ctx);
        return { ok: true, op, requestId: pending.requestId };
      }
      case "ask_cancel": {
        const pending = pendingAsks.get(String(request.requestId || ""));
        if (!pending) return { ok: false, op, error: "Ask request is no longer pending" };
        pending.settle({ cancelled: true });
        publishEvent({ type: "ask_cancelled", requestId: pending.requestId, toolCallId: pending.toolCallId }, ctx);
        return { ok: true, op, requestId: pending.requestId };
      }
      case "compact": {
        if (typeof ctx.compact !== "function") return { ok: false, op, error: "compaction is unavailable" };
        const instructions = typeof request.instructions === "string" && request.instructions.trim()
          ? request.instructions
          : undefined;
        // Compaction outlives a control round trip; acknowledge the start and
        // let the session_compact event report the result.
        void Promise.resolve(ctx.compact(instructions)).catch(() => {});
        return { ok: true, op, started: true };
      }
      case "interrupt":
        ctx.abort();
        return { ok: true, op };
      default:
        return { ok: false, op, error: `Unsupported control op: ${op || "(missing)"}` };
    }
  };

  const handleConnection = (socket, ctx) => {
    let subscribed = false;
    socket.setEncoding("utf8");
    let buffer = "";
    let requestChain = Promise.resolve();

    const cleanup = () => {
      subscribers.delete(socket);
    };

    socket.on("close", cleanup);
    socket.on("end", cleanup);
    socket.on("error", () => {
      cleanup();
      socket.destroy();
    });

    socket.on("data", (chunk) => {
      buffer += chunk;
      if (buffer.length > MAX_LINE_BYTES) {
        respond(socket, { ok: false, error: "Request too large" });
        cleanup();
        socket.destroy();
        return;
      }
      let newline = buffer.indexOf("\n");
      while (newline >= 0) {
        const line = buffer.slice(0, newline).trim();
        buffer = buffer.slice(newline + 1);
        if (line) {
          requestChain = requestChain
            .then(async () => {
              let request;
              try {
                request = JSON.parse(line);
              } catch {
                respond(socket, { ok: false, error: "Invalid JSON" });
                return;
              }
              if (request?.op === "subscribe" && subscribed) {
                respond(socket, { ok: false, op: "subscribe", error: "Already subscribed" });
                return;
              }
              const reply = await handleRequest(request, ctx);
              if (!reply) return;
              const replay = reply.__replayFrames;
              if (replay) delete reply.__replayFrames;
              if (!respond(socket, reply)) {
                socket.destroy(new Error("Subscription ACK backpressure"));
                return;
              }
              if (replay && Array.isArray(replay)) {
                let replayOk = true;
                for (const frame of replay) {
                  if (
                    socket.destroyed
                    || socket.writableEnded
                    || !socket.writable
                    || socket.bufferSize > MAX_SUBSCRIBER_BUFFER
                    || !respond(socket, frame)
                  ) {
                    replayOk = false;
                    break;
                  }
                }
                if (!replayOk || subscribers.size >= MAX_SUBSCRIBERS) {
                  socket.destroy(new Error(replayOk ? "Too many subscribers" : "Replay backpressure"));
                  return;
                }
                subscribers.add(socket);
                subscribed = true;
              }
            })
            .catch((error) => {
              respond(socket, {
                ok: false,
                error: error instanceof Error ? error.message : String(error),
              });
            });
        }
        newline = buffer.indexOf("\n");
      }
    });
  };

  const stop = async () => {
    for (const socket of subscribers) {
      try {
        socket.destroy();
      } catch {}
    }
    subscribers.clear();
    replayBuffer.length = 0;
    currentCtx = null;
    clearInterval(heartbeatTimer);
    for (const pending of [...pendingAsks.values()]) pending.settle({ cancelled: true });
    pendingAsks.clear();
    heartbeatTimer = null;

    const listener = server;
    const socketPath = boundPath;
    server = null;
    boundPath = null;
    if (listener) await new Promise((resolve) => listener.close(() => resolve()));
    if (socketPath) await fsp.rm(socketPath, { force: true }).catch(() => {});
  };

  api.on?.("session_start", async (_event, ctx) => {
    currentCtx = ctx;
    streamId = crypto.randomUUID();
    sequence = 0;
    for (const pending of [...pendingAsks.values()]) pending.settle({ cancelled: true });
    pendingAsks.clear();
    replayBuffer.length = 0;
    for (const socket of subscribers) {
      try { socket.destroy(); } catch {}
    }
    clearInterval(heartbeatTimer);
    heartbeatTimer = setInterval(() => {
      if (subscribers.size > 0) publishEvent({ type: "heartbeat" }, ctx);
    }, 5_000);
    heartbeatTimer.unref?.();
    subscribers.clear();

    const sessionId = ctx.sessionManager?.getSessionId?.();
    if (!sessionId) return;
    const socketPath = controlSocketPath(sessionId);
    const dir = controlDir();
    await fsp.mkdir(dir, { recursive: true, mode: 0o700 }).catch(() => {});
    await fsp.chmod(dir, 0o700).catch(() => {});
    if (!await clearStaleSocket(socketPath)) return;

    const listener = net.createServer((socket) => handleConnection(socket, ctx));
    listener.unref();
    await new Promise((resolve) => {
      listener.once("error", () => resolve());
      listener.listen(socketPath, () => resolve());
    });
    if (!listener.listening) return;
    await fsp.chmod(socketPath, 0o600).catch(() => {});
    server = listener;
    boundPath = socketPath;
  });

  api.on?.("session_shutdown", stop);
}