/**
 * SSH ControlMaster Remote Execution
 *
 * Like the ssh.ts example but uses a persistent ControlMaster socket so all
 * tool calls (read, write, edit, bash) multiplex over a single SSH connection
 * instead of opening a new one each time. Much faster on slow links or when
 * pi makes many small file reads in sequence.
 *
 * Usage:
 *   pi -e ~/.pi/agent/extensions/ssh-controlmaster.ts --ssh user@host
 *   pi -e ~/.pi/agent/extensions/ssh-controlmaster.ts --ssh user@host:/remote/path
 *
 *   Or, once installed globally, just:
 *   pi --ssh user@host
 *
 * Requirements:
 *   - SSH key-based auth (no password prompts)
 *   - bash on the remote machine
 */

import { spawn } from "node:child_process";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import {
  type BashOperations,
  createBashTool,
  createEditTool,
  createReadTool,
  createWriteTool,
  type EditOperations,
  type ReadOperations,
  type WriteOperations,
} from "@mariozechner/pi-coding-agent";

// ── ControlMaster helpers ────────────────────────────────────────────────────

function controlSocketPath(): string {
  // Keep path short — macOS has a ~104-char Unix socket path limit
  return join(tmpdir(), `pi-cm-${process.pid}.sock`);
}

function startControlMaster(remote: string, socketPath: string): Promise<void> {
  return new Promise((resolve, reject) => {
    // -fN: fork to background after auth, don't run a remote command
    // ControlPersist=yes: keep the master alive in the background indefinitely
    const child = spawn(
      "ssh",
      [
        "-fN",
        "-o", "ControlMaster=yes",
        "-o", `ControlPath=${socketPath}`,
        "-o", "ControlPersist=yes",
        remote,
      ],
      { stdio: "ignore" },
    );
    child.on("error", reject);
    child.on("close", (code) => {
      // -f forks to background; parent exits 0 once master is ready
      if (code === 0) resolve();
      else reject(new Error(`ControlMaster exited with code ${code}`));
    });
  });
}

function stopControlMaster(remote: string, socketPath: string): Promise<void> {
  return new Promise((resolve) => {
    // -O exit sends the exit signal to the master process via the socket
    const child = spawn(
      "ssh",
      ["-O", "exit", "-o", `ControlPath=${socketPath}`, remote],
      { stdio: "ignore" },
    );
    child.on("close", () => resolve()); // best-effort; ignore errors on cleanup
  });
}

// ── SSH exec (multiplexed over ControlMaster socket) ────────────────────────

function sshExec(remote: string, socketPath: string, command: string): Promise<Buffer> {
  return new Promise((resolve, reject) => {
    const child = spawn(
      "ssh",
      [
        "-o", "ControlMaster=no",   // use existing master, don't create a new one
        "-o", `ControlPath=${socketPath}`,
        remote,
        command,
      ],
      { stdio: ["ignore", "pipe", "pipe"] },
    );
    const chunks: Buffer[] = [];
    const errChunks: Buffer[] = [];
    child.stdout.on("data", (d: Buffer) => chunks.push(d));
    child.stderr.on("data", (d: Buffer) => errChunks.push(d));
    child.on("error", reject);
    child.on("close", (code) => {
      if (code !== 0) {
        reject(new Error(`SSH failed (${code}): ${Buffer.concat(errChunks).toString()}`));
      } else {
        resolve(Buffer.concat(chunks));
      }
    });
  });
}

// ── Remote tool operations ───────────────────────────────────────────────────

function createRemoteReadOps(
  remote: string,
  socketPath: string,
  remoteCwd: string,
  localCwd: string,
): ReadOperations {
  const toRemote = (p: string) => p.replace(localCwd, remoteCwd);
  return {
    readFile: (p) => sshExec(remote, socketPath, `cat ${JSON.stringify(toRemote(p))}`),
    access: (p) =>
      sshExec(remote, socketPath, `test -r ${JSON.stringify(toRemote(p))}`).then(() => {}),
    detectImageMimeType: async (p) => {
      try {
        const r = await sshExec(
          remote,
          socketPath,
          `file --mime-type -b ${JSON.stringify(toRemote(p))}`,
        );
        const m = r.toString().trim();
        return ["image/jpeg", "image/png", "image/gif", "image/webp"].includes(m)
          ? (m as "image/jpeg" | "image/png" | "image/gif" | "image/webp")
          : null;
      } catch {
        return null;
      }
    },
  };
}

function createRemoteWriteOps(
  remote: string,
  socketPath: string,
  remoteCwd: string,
  localCwd: string,
): WriteOperations {
  const toRemote = (p: string) => p.replace(localCwd, remoteCwd);
  return {
    writeFile: async (p, content) => {
      const b64 = Buffer.from(content).toString("base64");
      await sshExec(
        remote,
        socketPath,
        `echo ${JSON.stringify(b64)} | base64 -d > ${JSON.stringify(toRemote(p))}`,
      );
    },
    mkdir: (dir) =>
      sshExec(remote, socketPath, `mkdir -p ${JSON.stringify(toRemote(dir))}`).then(() => {}),
  };
}

function createRemoteEditOps(
  remote: string,
  socketPath: string,
  remoteCwd: string,
  localCwd: string,
): EditOperations {
  const r = createRemoteReadOps(remote, socketPath, remoteCwd, localCwd);
  const w = createRemoteWriteOps(remote, socketPath, remoteCwd, localCwd);
  return { readFile: r.readFile, access: r.access, writeFile: w.writeFile };
}

function createRemoteBashOps(
  remote: string,
  socketPath: string,
  remoteCwd: string,
  localCwd: string,
): BashOperations {
  const toRemote = (p: string) => p.replace(localCwd, remoteCwd);
  return {
    exec: (command, cwd, { onData, signal, timeout }) =>
      new Promise((resolve, reject) => {
        const cmd = `cd ${JSON.stringify(toRemote(cwd))} && ${command}`;
        const child = spawn(
          "ssh",
          [
            "-o", "ControlMaster=no",
            "-o", `ControlPath=${socketPath}`,
            remote,
            cmd,
          ],
          { stdio: ["ignore", "pipe", "pipe"] },
        );

        let timedOut = false;
        const timer = timeout
          ? setTimeout(() => {
              timedOut = true;
              child.kill();
            }, timeout * 1000)
          : undefined;

        child.stdout.on("data", onData);
        child.stderr.on("data", onData);
        child.on("error", (e: Error) => {
          if (timer) clearTimeout(timer);
          reject(e);
        });

        const onAbort = () => child.kill();
        signal?.addEventListener("abort", onAbort, { once: true });

        child.on("close", (code: number) => {
          if (timer) clearTimeout(timer);
          signal?.removeEventListener("abort", onAbort);
          if (signal?.aborted) reject(new Error("aborted"));
          else if (timedOut) reject(new Error(`timeout:${timeout}`));
          else resolve({ exitCode: code });
        });
      }),
  };
}

// ── Extension ────────────────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {
  pi.registerFlag("ssh", {
    description: "SSH remote: user@host or user@host:/remote/path",
    type: "string",
  });

  const localCwd = process.cwd();
  const localRead = createReadTool(localCwd);
  const localWrite = createWriteTool(localCwd);
  const localEdit = createEditTool(localCwd);
  const localBash = createBashTool(localCwd);

  // Resolved lazily in session_start once CLI flags are available
  let resolvedSsh: {
    remote: string;
    remoteCwd: string;
    socketPath: string;
  } | null = null;

  const getSsh = () => resolvedSsh;

  // ── Tool overrides ───────────────────────────────────────────────────────

  pi.registerTool({
    ...localRead,
    async execute(id, params, signal, onUpdate, _ctx) {
      const ssh = getSsh();
      if (ssh) {
        const tool = createReadTool(localCwd, {
          operations: createRemoteReadOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd),
        });
        return tool.execute(id, params, signal, onUpdate);
      }
      return localRead.execute(id, params, signal, onUpdate);
    },
  });

  pi.registerTool({
    ...localWrite,
    async execute(id, params, signal, onUpdate, _ctx) {
      const ssh = getSsh();
      if (ssh) {
        const tool = createWriteTool(localCwd, {
          operations: createRemoteWriteOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd),
        });
        return tool.execute(id, params, signal, onUpdate);
      }
      return localWrite.execute(id, params, signal, onUpdate);
    },
  });

  pi.registerTool({
    ...localEdit,
    async execute(id, params, signal, onUpdate, _ctx) {
      const ssh = getSsh();
      if (ssh) {
        const tool = createEditTool(localCwd, {
          operations: createRemoteEditOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd),
        });
        return tool.execute(id, params, signal, onUpdate);
      }
      return localEdit.execute(id, params, signal, onUpdate);
    },
  });

  pi.registerTool({
    ...localBash,
    async execute(id, params, signal, onUpdate, _ctx) {
      const ssh = getSsh();
      if (ssh) {
        const tool = createBashTool(localCwd, {
          operations: createRemoteBashOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd),
        });
        return tool.execute(id, params, signal, onUpdate);
      }
      return localBash.execute(id, params, signal, onUpdate);
    },
  });

  // ── Session lifecycle ────────────────────────────────────────────────────

  pi.on("session_start", async (_event, ctx) => {
    const arg = pi.getFlag("ssh") as string | undefined;
    if (!arg) return;

    let remote: string;
    let remoteCwd: string;

    if (arg.includes(":")) {
      [remote, remoteCwd] = arg.split(":");
    } else {
      remote = arg;
      // Resolve remote home dir via a plain one-shot SSH call before the
      // ControlMaster is up — this is the only un-multiplexed connection.
      const result = await new Promise<string>((resolve, reject) => {
        const child = spawn("ssh", [remote, "pwd"], {
          stdio: ["ignore", "pipe", "pipe"],
        });
        const chunks: Buffer[] = [];
        child.stdout.on("data", (d: Buffer) => chunks.push(d));
        child.on("error", reject);
        child.on("close", (code) => {
          if (code === 0) resolve(Buffer.concat(chunks).toString().trim());
          else reject(new Error(`Could not resolve remote pwd (exit ${code})`));
        });
      });
      remoteCwd = result;
    }

    const socketPath = controlSocketPath();
    ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH: ${remote} ⟳ connecting…`));

    try {
      await startControlMaster(remote, socketPath);
    } catch (err) {
      ctx.ui.setStatus("ssh", ctx.ui.theme.fg("error", `SSH: ${remote} ✗ failed`));
      ctx.ui.notify(`SSH ControlMaster failed: ${err}`, "error");
      return;
    }

    resolvedSsh = { remote, remoteCwd, socketPath };
    ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH ⚡ ${remote}:${remoteCwd}`));
    ctx.ui.notify(`SSH ControlMaster ready — ${remote}:${remoteCwd}`, "success");
  });

  pi.on("session_shutdown", async (_event, _ctx) => {
    const ssh = getSsh();
    if (ssh) {
      await stopControlMaster(ssh.remote, ssh.socketPath);
      resolvedSsh = null;
    }
  });

  // ── User ! commands via SSH ──────────────────────────────────────────────

  pi.on("user_bash", (_event) => {
    const ssh = getSsh();
    if (!ssh) return;
    return {
      operations: createRemoteBashOps(ssh.remote, ssh.socketPath, ssh.remoteCwd, localCwd),
    };
  });

  // ── Tell the LLM where it's operating ───────────────────────────────────

  pi.on("before_agent_start", async (event) => {
    const ssh = getSsh();
    if (ssh) {
      const modified = event.systemPrompt.replace(
        `Current working directory: ${localCwd}`,
        `Current working directory: ${ssh.remoteCwd} (via SSH ControlMaster: ${ssh.remote})`,
      );
      return { systemPrompt: modified };
    }
  });
}