import { homedir } from "node:os";
import { posix as posixPath } from "node:path";
import type { CustomEntry, ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
  createBashTool,
  createEditTool,
  createReadTool,
  createWriteTool,
  type BashOperations,
  type EditOperations,
  type ReadOperations,
  type WriteOperations,
} from "@earendil-works/pi-coding-agent";
import { getPermissionsService, PERMISSIONS_READY_CHANNEL } from "@gotgenes/pi-permission-system";
import { installSshPermissionIntegration } from "./permission-integration.ts";
import { loadVault } from "./src/vault.ts";
import { Ssh2Transport, type RemoteTransport } from "./src/ssh2-transport.ts";
import type { PiSshConfig, SshHostConfig } from "./src/config.ts";
import {
  runRemoteFind,
  runRemoteGrep,
  type RemoteFindInput,
  type RemoteGrepInput,
} from "./src/remote-search.ts";

interface SshStoredConfig {
  hostId: string;
  remoteCwd: string;
  remoteHome: string;
}

interface SshConnection {
  hostId: string;
  remote: string;
  port: number;
  remoteCwd: string;
  remoteHome: string;
  localCwd: string;
  localHome: string;
}

interface HostSelection {
  hostId: string;
  remotePath?: string;
}

function parseHostSelection(raw: string): HostSelection {
  const value = raw.trim();
  if (!value) throw new Error("SSH host id is required");
  const colon = value.indexOf(":");
  const hostId = (colon < 0 ? value : value.slice(0, colon)).trim();
  if (!/^[A-Za-z0-9][A-Za-z0-9._-]*$/.test(hostId)) throw new Error(`invalid pi-ssh host id: ${hostId}`);
  if (colon < 0) return { hostId };
  const remotePath = value.slice(colon + 1).trim();
  if (!remotePath || !(remotePath === "~" || remotePath.startsWith("~/") || remotePath.startsWith("/"))) {
    throw new Error("remote path must be absolute or start with ~/");
  }
  return { hostId, remotePath };
}

function mapLocalPathToRemote(path: string, connection: SshConnection): string {
  if (path === connection.localCwd) return connection.remoteCwd;
  if (path.startsWith(`${connection.localCwd}/`)) {
    return `${connection.remoteCwd}${path.slice(connection.localCwd.length)}`;
  }
  if (path === connection.localHome) return connection.remoteHome;
  if (path.startsWith(`${connection.localHome}/`)) {
    return `${connection.remoteHome}${path.slice(connection.localHome.length)}`;
  }
  return path;
}

function createRemoteReadOps(connection: SshConnection, transport: RemoteTransport): ReadOperations {
  return {
    readFile: (absolutePath) => transport.readFile(mapLocalPathToRemote(absolutePath, connection)),
    access: (absolutePath) => transport.ensureReadable(mapLocalPathToRemote(absolutePath, connection)),
    detectImageMimeType: async (absolutePath) => {
      try {
        return await transport.detectImageMimeType(mapLocalPathToRemote(absolutePath, connection));
      } catch {
        return null;
      }
    },
  };
}

function createRemoteWriteOps(connection: SshConnection, transport: RemoteTransport): WriteOperations {
  return {
    mkdir: (absoluteDir) => transport.mkdir(mapLocalPathToRemote(absoluteDir, connection)),
    writeFile: (absolutePath, content) =>
      transport.writeFile(mapLocalPathToRemote(absolutePath, connection), Buffer.from(content, "utf8")),
  };
}

function createRemoteEditOps(connection: SshConnection, transport: RemoteTransport): EditOperations {
  const read = createRemoteReadOps(connection, transport);
  const write = createRemoteWriteOps(connection, transport);
  return {
    readFile: read.readFile,
    writeFile: write.writeFile,
    access: (absolutePath) => transport.ensureReadableWritable(mapLocalPathToRemote(absolutePath, connection)),
  };
}

function createRemoteBashOps(transport: RemoteTransport): BashOperations {
  return {
    exec: (command, cwd, { onData, signal, timeout }) => transport.exec(command, cwd, { onData, signal, timeout }),
  };
}

function getConfiguredHost(config: PiSshConfig, hostId: string): SshHostConfig {
  const host = config.hosts[hostId];
  if (!host) throw new Error(`unknown pi-ssh host '${hostId}'; run ssh_config.sh import ${hostId}`);
  return host;
}

function resolveRequestedPath(selection: HostSelection, host: SshHostConfig, remoteHome: string, remotePwd: string): string {
  const requested = selection.remotePath ?? host.defaultCwd ?? remotePwd;
  if (requested === "~") return remoteHome;
  if (requested.startsWith("~/")) return posixPath.join(remoteHome, requested.slice(2));
  return requested;
}

async function captureChecked(transport: RemoteTransport, command: string, cwd = "."): Promise<string> {
  const result = await transport.capture(command, cwd, 20);
  if (result.exitCode !== 0) {
    const message = result.output.toString("utf8").trim();
    throw new Error(message || `remote command failed with exit code ${result.exitCode}`);
  }
  return result.output.toString("utf8").trim();
}

async function connectSelection(
  selection: HostSelection,
  localCwd: string,
  localHome: string,
): Promise<{ connection: SshConnection; transport: Ssh2Transport }> {
  const config = loadVault();
  const host = getConfiguredHost(config, selection.hostId);
  const transport = new Ssh2Transport(host);
  try {
    await transport.connect();
    const remoteHome = await captureChecked(transport, 'printf "%s" "$HOME"');
    const remotePwd = await captureChecked(transport, "pwd");
    if (!remoteHome || !remotePwd) throw new Error("remote HOME/cwd probe returned empty output");
    const requestedPath = resolveRequestedPath(selection, host, remoteHome, remotePwd);
    const remoteCwd = await captureChecked(transport, "pwd", requestedPath);
    return {
      connection: {
        hostId: selection.hostId,
        remote: `${selection.hostId} [${host.user}@${host.hostName}]`,
        port: host.port,
        remoteCwd,
        remoteHome,
        localCwd,
        localHome,
      },
      transport,
    };
  } catch (error) {
    await transport.dispose();
    throw error;
  }
}

function pickerEntries(config: PiSshConfig): Array<{ value: string; hostId: string }> {
  const groupByHost = new Map<string, string>();
  for (const group of Object.values(config.groups ?? {})) {
    for (const hostId of group.hosts) if (!groupByHost.has(hostId)) groupByHost.set(hostId, group.label);
  }
  return Object.entries(config.hosts).map(([hostId, host]) => ({
    hostId,
    value: `${groupByHost.get(hostId) ? `${groupByHost.get(hostId)} / ` : ""}${host.label ?? hostId} [${hostId}]`,
  }));
}

export default function piSshExtension(pi: ExtensionAPI): void {
  pi.registerFlag("ssh", {
    description: "Configured pi-ssh host id, optionally followed by :/absolute/remote/path",
    type: "string",
  });

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

  let connection: SshConnection | null = null;
  let transport: Ssh2Transport | null = null;
  const getConnection = () => connection;
  installSshPermissionIntegration(pi, getConnection, {
    getPermissionsService,
    permissionsReadyChannel: PERMISSIONS_READY_CHANNEL,
  });

  const requireSsh = (toolName: string): { connection: SshConnection; transport: Ssh2Transport } => {
    if (!connection || !transport) {
      throw new Error(`${toolName} requires an active SSH2 connection. Run /ssh and select an imported host.`);
    }
    return { connection, transport };
  };

  const activateConnection = async (
    nextConnection: SshConnection,
    nextTransport: Ssh2Transport,
    ctx: ExtensionContext,
    options: { persist: boolean; verb: string },
  ): Promise<void> => {
    if (transport) await transport.dispose();
    connection = nextConnection;
    transport = nextTransport;
    if (options.persist) {
      pi.appendEntry("pi-ssh-config", {
        hostId: nextConnection.hostId,
        remoteCwd: nextConnection.remoteCwd,
        remoteHome: nextConnection.remoteHome,
      } satisfies SshStoredConfig);
    }
    const message = `pi-ssh ${options.verb}: ${nextConnection.remote}:${nextConnection.remoteCwd} (port ${nextConnection.port})`;
    console.log(message);
    if (ctx.hasUI) {
      ctx.ui.setStatus("pi-ssh", ctx.ui.theme.fg("accent", `SSH ${nextConnection.hostId}:${nextConnection.remoteCwd}`));
      ctx.ui.notify(message, "info");
    }
  };

  const deactivateConnection = async (ctx: ExtensionContext): Promise<void> => {
    if (transport) await transport.dispose();
    transport = null;
    connection = null;
    if (ctx.hasUI) ctx.ui.setStatus("pi-ssh", undefined);
  };

  pi.registerTool({
    ...localRead,
    name: "ssh_read",
    label: "ssh_read",
    description: `Read a file through the active SSH2 connection. ${localRead.description}`,
    async execute(id, params, signal, onUpdate) {
      const active = requireSsh("ssh_read");
      return createReadTool(localCwd, { operations: createRemoteReadOps(active.connection, active.transport) })
        .execute(id, params, signal, onUpdate);
    },
  });

  pi.registerTool({
    ...localWrite,
    name: "ssh_write",
    label: "ssh_write",
    description: `Write a file through the active SSH2 connection. ${localWrite.description}`,
    async execute(id, params, signal, onUpdate) {
      const active = requireSsh("ssh_write");
      return createWriteTool(localCwd, { operations: createRemoteWriteOps(active.connection, active.transport) })
        .execute(id, params, signal, onUpdate);
    },
  });

  pi.registerTool({
    ...localEdit,
    name: "ssh_edit",
    label: "ssh_edit",
    description: `Edit a file through the active SSH2 connection. ${localEdit.description}`,
    async execute(id, params, signal, onUpdate) {
      const active = requireSsh("ssh_edit");
      return createEditTool(localCwd, { operations: createRemoteEditOps(active.connection, active.transport) })
        .execute(id, params, signal, onUpdate);
    },
  });

  pi.registerTool({
    name: "ssh_find",
    label: "ssh_find",
    description: "Find remote files by a fixed filename/path substring using fd, git ls-files, or find. Results are bounded and require an active SSH2 connection.",
    parameters: {
      type: "object",
      properties: {
        pattern: { type: "string", description: "Fixed substring to match in remote file paths" },
        path: { type: "string", description: "Remote root path; defaults to the active remote cwd" },
        caseSensitive: { type: "boolean", description: "Use case-sensitive matching (default: false)" },
        includeHidden: { type: "boolean", description: "Include hidden paths (default: false)" },
        limit: { type: "integer", minimum: 1, maximum: 200, description: "Maximum results (default: 50)" },
      },
      required: ["pattern"],
      additionalProperties: false,
    },
    async execute(_id, params, signal) {
      const active = requireSsh("ssh_find");
      const result = await runRemoteFind(
        active.transport,
        params as RemoteFindInput,
        active.connection.remoteCwd,
        active.connection.remoteHome,
        signal,
      );
      return {
        content: [{ type: "text", text: result.text }],
        details: { backend: result.backend, matchCount: result.matchCount, truncated: result.truncated },
      };
    },
  });

  pi.registerTool({
    name: "ssh_grep",
    label: "ssh_grep",
    description: "Search remote file contents using ripgrep, git grep, or grep. Literal case-insensitive matching is the default; results are bounded and require an active SSH2 connection.",
    parameters: {
      type: "object",
      properties: {
        pattern: { type: "string", description: "Text or regular expression to search for" },
        path: { type: "string", description: "Remote root path; defaults to the active remote cwd" },
        literal: { type: "boolean", description: "Treat pattern as fixed text (default: true)" },
        caseSensitive: { type: "boolean", description: "Use case-sensitive matching (default: false)" },
        include: { type: "string", description: "Optional file glob such as *.ts" },
        limit: { type: "integer", minimum: 1, maximum: 200, description: "Maximum result lines (default: 50)" },
      },
      required: ["pattern"],
      additionalProperties: false,
    },
    async execute(_id, params, signal) {
      const active = requireSsh("ssh_grep");
      const result = await runRemoteGrep(
        active.transport,
        params as RemoteGrepInput,
        active.connection.remoteCwd,
        active.connection.remoteHome,
        signal,
      );
      return {
        content: [{ type: "text", text: result.text }],
        details: { backend: result.backend, matchCount: result.matchCount, truncated: result.truncated },
      };
    },
  });

  pi.registerTool({
    ...localBash,
    name: "ssh_bash",
    label: "ssh_bash",
    description: `Run a shell command through the active SSH2 connection. ${localBash.description}`,
    async execute(id, params, signal, onUpdate) {
      const active = requireSsh("ssh_bash");
      return createBashTool(localCwd, { operations: createRemoteBashOps(active.transport) })
        .execute(id, params, signal, onUpdate);
    },
  });

  pi.on("session_start", async (event, ctx) => {
    const flag = pi.getFlag("ssh") as string | undefined;
    if (flag) {
      try {
        const connected = await connectSelection(parseHostSelection(flag), localCwd, localHome);
        await activateConnection(connected.connection, connected.transport, ctx, { persist: true, verb: "enabled" });
        return;
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        await deactivateConnection(ctx);
        console.error(`pi-ssh failed to connect: ${message}`);
        if (ctx.hasUI) ctx.ui.notify(`pi-ssh failed to connect: ${message}`, "error");
        throw error;
      }
    }

    if (event.reason !== "startup" && event.reason !== "resume") return;
    const entries = ctx.sessionManager.getEntries();
    let stored: SshStoredConfig | undefined;
    for (let index = entries.length - 1; index >= 0; index -= 1) {
      const entry = entries[index];
      if (entry.type === "custom" && (entry as CustomEntry<unknown>).customType === "pi-ssh-config") {
        stored = (entry as CustomEntry<SshStoredConfig>).data;
        if (stored) break;
      }
    }
    if (!stored) return;
    try {
      const connected = await connectSelection(
        { hostId: stored.hostId, remotePath: stored.remoteCwd },
        localCwd,
        localHome,
      );
      await activateConnection(connected.connection, connected.transport, ctx, { persist: false, verb: "resumed" });
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      await deactivateConnection(ctx);
      console.error(`pi-ssh resume failed: ${message}`);
      if (ctx.hasUI) ctx.ui.notify(`pi-ssh resume failed: ${message}`, "warning");
    }
  });

  pi.registerCommand("ssh", {
    description: "Connect configured SSH2 hosts: /ssh [host-id[:/path]], /ssh status, /ssh off",
    getArgumentCompletions: (prefix) => {
      try {
        const options = ["off", "status", ...Object.keys(loadVault().hosts)];
        const filtered = options.filter((option) => option.startsWith(prefix));
        return filtered.length > 0 ? filtered.map((option) => ({ value: option, label: option })) : null;
      } catch {
        return ["off", "status"].filter((option) => option.startsWith(prefix)).map((option) => ({ value: option, label: option }));
      }
    },
    handler: async (args, ctx) => {
      const input = args.trim();
      if (input === "status") {
        ctx.ui.notify(connection
          ? `pi-ssh: ${connection.remote}:${connection.remoteCwd} (port ${connection.port})`
          : "pi-ssh: not connected (local tools active)", "info");
        return;
      }
      if (input === "off") {
        await deactivateConnection(ctx);
        ctx.ui.notify("pi-ssh: disconnected", "info");
        return;
      }

      let target = input;
      if (!target) {
        let config: PiSshConfig;
        try {
          config = loadVault();
        } catch (error) {
          ctx.ui.notify(error instanceof Error ? error.message : String(error), "warning");
          return;
        }
        const entries = pickerEntries(config);
        if (entries.length === 0) {
          ctx.ui.notify("No pi-ssh hosts configured. Run ssh_config.sh import.", "warning");
          return;
        }
        const values = [...(connection ? ["Disconnect [off]"] : []), ...entries.map((entry) => entry.value)];
        const selected = await ctx.ui.select("SSH2 host", values);
        if (!selected) return;
        if (selected === "Disconnect [off]") {
          await deactivateConnection(ctx);
          ctx.ui.notify("pi-ssh: disconnected", "info");
          return;
        }
        target = entries.find((entry) => entry.value === selected)?.hostId ?? "";
      }

      try {
        const connected = await connectSelection(parseHostSelection(target), localCwd, localHome);
        await activateConnection(connected.connection, connected.transport, ctx, { persist: true, verb: "connected" });
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        ctx.ui.notify(`pi-ssh: failed to connect: ${message}`, "error");
      }
    },
  });

  pi.on("session_shutdown", async () => {
    if (transport) await transport.dispose();
    transport = null;
  });

  pi.on("user_bash", () => transport ? { operations: createRemoteBashOps(transport) } : undefined);

  pi.on("before_agent_start", async (event) => {
    if (!connection) return;
    const guidance =
      `\n\n# Remote SSH connection\n\n` +
      `An SSH2 connection to configured host ${connection.remote} on port ${connection.port} is active. ` +
      `The default read/write/edit/bash/find/grep tools act on the LOCAL machine. ` +
      `Use ssh_read, ssh_write, ssh_edit, ssh_find, ssh_grep, and ssh_bash for explicit remote operations. ` +
      `Use ssh_find before ssh_grep to narrow remote searches; both tools return bounded results and choose the fastest available remote backend. ` +
      `(User \`!\` commands run remotely.) ` +
      `Remote operations are rooted at ${connection.remoteCwd}; relative paths resolve against that directory.`;
    return { systemPrompt: `${event.systemPrompt}${guidance}` };
  });
}