#!/usr/bin/env -S deno run --allow-read --allow-write --allow-run --allow-env
/**
 * configure_ssh_key.ts
 * Add or update an SSH config entry for a specific host/user using a specified key.
 * Can optionally generate the key if it does not exist.
 *
 * Usage:
 *   ./configure_ssh_key.ts alice@example.com --key ~/.ssh/id_alice_example --alias work-example
 *   ./configure_ssh_key.ts alice@example.com --generate --alias work-example
 *   ./configure_ssh_key.ts alice@example.com --alias work-example
 *   ./configure_ssh_key.ts alice@example.com:2222
 *
 * Required positional target:
 *   user@host[:port]
 *
 * Options:
 *   --key /path/to/key       Path to private key
 *   --generate               Generate a new ed25519 key if missing
 *   --alias HOST_ALIAS       Host alias (defaults to USER-HOST simplified)
 *   --port PORT              Override / set SSH port
 *   --force                  Overwrite existing config block for this alias
 *   --dry-run                Show intended changes without writing
 *   --help                   Show help
 *
 * If neither --key nor --generate is specified, you'll be prompted to select
 * from available keys in ~/.ssh or ssh-agent.
 */

// Export completion function for loader
interface CompletionContext {
  args: string[];
  current: string;
  position: number;
}

export async function getCompletions(ctx: CompletionContext): Promise<string[]> {
  const { current, args } = ctx;
  
  // Helper to find SSH keys in ~/.ssh/
  const findSSHKeys = async (): Promise<string[]> => {
    try {
      const home = Deno.env.get("HOME") || Deno.env.get("USERPROFILE") || "";
      if (!home) return [];
      const sshDir = `${home}/.ssh`;
      const keys: string[] = [];
      
      try {
        for await (const entry of Deno.readDir(sshDir)) {
          if (entry.isFile && !entry.name.endsWith(".pub") && 
              (entry.name.startsWith("id_") || entry.name.includes("key"))) {
            keys.push(`--key=~/.ssh/${entry.name}`);
          }
        }
      } catch {
        // Fallback to common names if can't read directory
        return ["--key=~/.ssh/id_rsa", "--key=~/.ssh/id_ed25519", "--key=~/.ssh/id_ecdsa"];
      }
      
      return keys.length > 0 ? keys : ["--key=~/.ssh/id_rsa", "--key=~/.ssh/id_ed25519"];
    } catch {
      return ["--key=~/.ssh/id_rsa", "--key=~/.ssh/id_ed25519"];
    }
  };
  
  // If completing a value after --key=
  if (current.startsWith("--key=")) {
    const prefix = "--key=";
    const value = current.slice(prefix.length);
    
    // If value is empty or starts with ~, suggest keys
    if (!value || value.startsWith("~")) {
      return await findSSHKeys();
    }
    // Otherwise suggest path completion hint
    return ["--key=<path-to-private-key>"];
  }
  
  // If completing --alias= value
  if (current.startsWith("--alias=")) {
    // Suggest format based on first positional arg if present
    const target = args.find(a => !a.startsWith("-"));
    if (target && target.includes("@")) {
      const [user, host] = target.split("@");
      const hostPart = host.split(":")[0];
      return [`--alias=${user}-${hostPart.replace(/\./g, "-")}`];
    }
    return ["--alias=<host-alias>"];
  }
  
  // If completing --port= value
  if (current.startsWith("--port=")) {
    return ["--port=22", "--port=2222", "--port=2200"];
  }
  
  // If no flags yet and no target, suggest target format
  const hasTarget = args.some(a => !a.startsWith("-") && a.includes("@"));
  if (!hasTarget && !current.startsWith("-")) {
    return ["user@hostname", "user@hostname:port"];
  }
  
  // Check if certain flags already provided to avoid duplicates
  const hasGenerate = args.includes("--generate");
  const hasKey = args.some(a => a.startsWith("--key"));
  const hasAlias = args.some(a => a.startsWith("--alias"));
  const hasPort = args.some(a => a.startsWith("--port"));
  const hasForce = args.includes("--force");
  const hasDryRun = args.includes("--dry-run");
  
  const options: string[] = [];
  
  // Only suggest --key if --generate not present
  if (!hasGenerate && !hasKey) {
    options.push("--key=");
  }
  
  if (!hasGenerate && !hasKey) {
    options.push("--generate");
  }
  
  if (!hasAlias) {
    options.push("--alias=");
  }
  
  if (!hasPort) {
    options.push("--port=");
  }
  
  if (!hasForce) {
    options.push("--force");
  }
  
  if (!hasDryRun) {
    options.push("--dry-run");
  }
  
  options.push("--help");
  
  return options;
}

import { parseArgs } from "jsr:@std/cli@^1.0.0/parse-args";
import { expandGlob } from "jsr:@std/fs@^1.0.0/expand-glob";
import { join } from "jsr:@std/path@^1.0.0";

interface ParsedTarget {
  user: string;
  host: string;
  port?: number;
}

interface Options {
  keyPath?: string;
  alias?: string;
  generate: boolean;
  port?: number;
  force: boolean;
  dryRun: boolean;
  isAgentKey?: boolean; // Track if key only exists in ssh-agent
}

interface AgentKey {
  keyType: string;
  publicKey: string;
  comment: string;
  displayName: string;
}

const HOME = Deno.env.get("HOME") || Deno.env.get("USERPROFILE") || "";
const SSH_DIR = join(HOME, ".ssh");
const CONFIG_FILE = join(SSH_DIR, "config");

function printHelp() {
  console.log(`
configure_ssh_key.ts - SSH Config Manager

Usage:
  ./configure_ssh_key.ts user@host[:port] [options]

Examples:
  ./configure_ssh_key.ts alice@example.com --key ~/.ssh/id_alice_example --alias work-example
  ./configure_ssh_key.ts alice@example.com --generate --alias work-example
  ./configure_ssh_key.ts alice@example.com --alias work-example
  ./configure_ssh_key.ts alice@example.com:2222

Options:
  --key PATH        Path to private key
  --generate        Generate a new ed25519 key if missing
  --alias NAME      Host alias (defaults to user-host)
  --port PORT       Override SSH port
  --force           Overwrite existing config block
  --dry-run         Preview changes without writing
  --help, -h        Show this help

Note: If neither --key nor --generate is specified, you'll be prompted to
select from available keys in ~/.ssh or ssh-agent.
`);
}

function parseTarget(target: string): ParsedTarget {
  // user@host:port
  let match = target.match(/^([^@]+)@([^:]+):(\d+)$/);
  if (match) {
    return { user: match[1], host: match[2], port: parseInt(match[3]) };
  }

  // user@host
  match = target.match(/^([^@]+)@([^:]+)$/);
  if (match) {
    return { user: match[1], host: match[2] };
  }

  // host:port
  match = target.match(/^([^:]+):(\d+)$/);
  if (match) {
    throw new Error("Target must include user (format: user@host[:port])");
  }

  throw new Error("Invalid target format. Use: user@host[:port]");
}

async function findPrivateKeys(): Promise<string[]> {
  const keys: string[] = [];
  
  try {
    for await (const entry of expandGlob(`${SSH_DIR}/id_*`, { includeDirs: false })) {
      if (!entry.name.endsWith(".pub")) {
        keys.push(entry.path);
      }
    }
    for await (const entry of expandGlob(`${SSH_DIR}/*.pem`, { includeDirs: false })) {
      keys.push(entry.path);
    }
  } catch {
    // Directory might not exist yet
  }

  return keys.sort();
}

async function listAgentKeys(): Promise<AgentKey[]> {
  const cmd = new Deno.Command("ssh-add", {
    args: ["-L"],
    stdout: "piped",
    stderr: "piped",
  });

  const { success, stdout } = await cmd.output();
  if (!success) {
    return [];
  }

  const output = new TextDecoder().decode(stdout);
  const lines = output.trim().split("\n").filter(line => line.length > 0);
  
  return lines.map(line => {
    const parts = line.split(" ");
    const keyType = parts[0] || "";
    const publicKey = parts[1] || "";
    const comment = parts.slice(2).join(" ") || "(no comment)";
    const displayName = comment.replace(/[/: ]/g, "_").toLowerCase();
    
    return { keyType, publicKey, comment, displayName };
  });
}

async function generateKey(keyPath: string, user: string, host: string) {
  try {
    await Deno.stat(keyPath);
    console.log(`Key already exists: ${keyPath} (skipping generation)`);
    return;
  } catch {
    // Key doesn't exist, generate it
  }

  console.log(`Generating ed25519 key: ${keyPath}`);
  
  const cmd = new Deno.Command("ssh-keygen", {
    args: [
      "-t", "ed25519",
      "-a", "100",
      "-f", keyPath,
      "-C", `${user}@${host}`,
      "-N", "",
    ],
  });

  const { success } = await cmd.output();
  if (!success) {
    throw new Error("Failed to generate SSH key");
  }
}

async function ensurePermissions(keyPath: string) {
  // Verify this is actually a private key file (not .pub)
  if (keyPath.endsWith(".pub")) {
    throw new Error("Key path should point to private key, not .pub file");
  }
  
  // Set private key permissions
  try {
    await Deno.chmod(keyPath, 0o600);
  } catch (e) {
    // If file doesn't exist yet (e.g., agent key), that's ok
    if (!(e instanceof Deno.errors.NotFound)) {
      throw e;
    }
  }
  
  // Set public key permissions if it exists
  const pubKeyPath = `${keyPath}.pub`;
  try {
    await Deno.stat(pubKeyPath);
    await Deno.chmod(pubKeyPath, 0o644);
  } catch {
    // .pub file doesn't exist, that's ok
  }
}

function generateConfigBlock(
  alias: string,
  host: string,
  user: string,
  keyPath: string,
  port?: number,
): string {
  const lines = [
    `Host ${alias}`,
    `  HostName ${host}`,
    `  User ${user}`,
  ];

  if (port) {
    lines.push(`  Port ${port}`);
  }

  lines.push(
    `  IdentityFile ${keyPath}`,
    `  IdentitiesOnly yes`,
    `  PreferredAuthentications publickey`,
  );

  return lines.join("\n");
}

async function readConfig(): Promise<string> {
  try {
    return await Deno.readTextFile(CONFIG_FILE);
  } catch {
    return "";
  }
}

function removeHostBlock(config: string, alias: string): string {
  const lines = config.split("\n");
  const result: string[] = [];
  let skipping = false;

  for (const line of lines) {
    if (line.match(/^Host\s+/)) {
      const hostMatch = line.match(/^Host\s+(\S+)/);
      if (hostMatch && hostMatch[1] === alias) {
        skipping = true;
        continue;
      } else {
        skipping = false;
      }
    }

    if (!skipping) {
      result.push(line);
    }
  }

  return result.join("\n");
}

function hasHostBlock(config: string, alias: string): boolean {
  const pattern = new RegExp(`^Host\\s+${alias.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}(\\s|$)`, "m");
  return pattern.test(config);
}

async function backupConfig() {
  try {
    const timestamp = new Date().toISOString().replace(/[:.]/g, "-").split("T").join("_").slice(0, 19);
    const backupPath = `${CONFIG_FILE}.bak.${timestamp}`;
    await Deno.copyFile(CONFIG_FILE, backupPath);
  } catch {
    // Backup failed, continue anyway
  }
}

function expandTilde(path: string): string {
  if (path.startsWith("~/")) {
    return join(HOME, path.slice(2));
  }
  return path;
}

function normalizeKeyPath(path: string): string {
  // Remove .pub extension if present - we always work with private key paths
  return path.endsWith(".pub") ? path.slice(0, -4) : path;
}

function simplifyForAlias(str: string): string {
  return str.toLowerCase().replace(/[^a-z0-9]/g, "-");
}

// Utility: Check for mutually exclusive options
function checkExclusiveOptions(opts: boolean[]): boolean {
  return opts.filter(Boolean).length <= 1;
}

// Utility: Central error handler
function handleError(e: unknown) {
  console.error(`Error: ${e instanceof Error ? e.message : String(e)}`);
  Deno.exit(1);
}

// Utility: Unified key selection
async function selectKeyUnified(options: Options): Promise<{ keyPath: string, isAgentKey: boolean }> {
  if (options.keyPath) return { keyPath: options.keyPath, isAgentKey: false };
  
  // Interactive combined selection (file keys + agent keys)
  const fileKeys = await findPrivateKeys();
  const agentKeys = await listAgentKeys();
  // ...existing code for combinedKeys...
  const fileKeyInfos: Array<{ path: string, pub: string | null }> = [];
  for (const path of fileKeys) {
    let pub: string | null = null;
    try {
      pub = await Deno.readTextFile(path + ".pub");
    } catch {
      pub = null;
    }
    fileKeyInfos.push({ path, pub });
  }
  async function agentKeyFilename(agentKey: AgentKey) {
    const hash = Array.from(new Uint8Array(await crypto.subtle.digest("SHA-256", new TextEncoder().encode(agentKey.publicKey))))
      .map(b => b.toString(16).padStart(2, "0")).join("");
    return `id_agent_${hash}`;
  }
  const combinedKeys: Array<{ type: "file" | "agent", path?: string, agentKey?: AgentKey, display: string, linkedFile?: string }> = [];
  for (const fileInfo of fileKeyInfos) {
    combinedKeys.push({ type: "file", path: fileInfo.path, display: `[file] ${fileInfo.path}` });
  }
  for (const agentKey of agentKeys) {
    let linkedFile: string | undefined = undefined;
    for (const fileInfo of fileKeyInfos) {
      if (fileInfo.pub && fileInfo.pub.includes(agentKey.publicKey)) {
        linkedFile = fileInfo.path;
        break;
      }
    }
    combinedKeys.push({
      type: "agent",
      agentKey,
      display: `[agent] ${agentKey.comment} (${agentKey.keyType})${linkedFile ? ` [linked: ${linkedFile}]` : ''}`,
      linkedFile
    });
  }
  if (combinedKeys.length === 0) {
    throw new Error("No candidate private keys found in ~/.ssh or ssh-agent");
  }
  console.log("\nAvailable keys:");
  combinedKeys.forEach((key, i) => {
    console.log(`  [${i}] ${key.display}`);
  });
  let selectedIdx: number | undefined = undefined;
  // Detect non-interactive environment
  const isInteractive = typeof Deno.stdin.setRaw === "function" && Deno.stdin.isTerminal();
  if (!isInteractive) {
    // If keyPath is not provided, fail in non-interactive mode
    if (!options.keyPath) {
      console.error("Error: Interactive key selection required, but running in non-interactive mode. Use --key to specify a key.");
      Deno.exit(1);
    }
    // If keyPath is provided, select it automatically
    selectedIdx = combinedKeys.findIndex(k => k.type === "file" && k.path === options.keyPath);
    if (selectedIdx === -1) {
      console.error(`Error: Specified key not found: ${options.keyPath}`);
      Deno.exit(1);
    }
  } else {
    while (selectedIdx === undefined) {
      const input = prompt("\nSelect key number (or q to quit):");
      if (input === null || input.toLowerCase() === "q") {
        console.error("Error: No key selected. Exiting.");
        Deno.exit(1);
      }
      const num = parseInt(input ?? "");
      if (!isNaN(num) && num >= 0 && num < combinedKeys.length) {
        selectedIdx = num;
      } else {
        console.error("Invalid selection. Try again.");
      }
    }
  }
  const selected = combinedKeys[selectedIdx];
  if (selected.type === "file" && selected.path) {
    return { keyPath: selected.path, isAgentKey: false };
  } else if (selected.type === "agent" && selected.agentKey) {
    if (selected.linkedFile) {
      return { keyPath: selected.linkedFile, isAgentKey: false };
    } else {
      const filename = await agentKeyFilename(selected.agentKey);
      const pubKeyPath = join(SSH_DIR, `${filename}.pub`);
      let pubKeyExists = false;
      try {
        await Deno.stat(pubKeyPath);
        pubKeyExists = true;
      } catch {
        // .pub file does not exist
      }
      if (!pubKeyExists) {
        const pubKeyContent = `${selected.agentKey.keyType} ${selected.agentKey.publicKey} ${selected.agentKey.comment}\n`;
        await Deno.writeTextFile(pubKeyPath, pubKeyContent);
        await Deno.chmod(pubKeyPath, 0o644);
        console.log(`\n✅ Saved public key to: ${pubKeyPath}`);
      }
      // Do not create a dummy private key file for agent keys
      return { keyPath: pubKeyPath, isAgentKey: true };
    }
  }
  throw new Error("No key path selected or generated.");
}

async function main() {
  try {
    const args = parseArgs(Deno.args, {
      string: ["key", "alias", "port"],
      boolean: ["generate", "force", "dry-run", "help", "h"],
      alias: { h: "help" },
      "--": true,
    });

    if (args.help || args.h) {
      printHelp();
      Deno.exit(0);
    }

    const targetSpec = args._[0]?.toString();
    if (!targetSpec) {
      throw new Error("Missing required target (user@host[:port])\nRun with --help for usage information");
    }

    let target: ParsedTarget;
    try {
      target = parseTarget(targetSpec);
    } catch (e) {
      throw e;
    }

    if (args.port) {
      target.port = parseInt(args.port);
    }

    const keyProvided = !!args.key;
    const generate = args.generate || false;
    if (!checkExclusiveOptions([keyProvided, generate])) {
      throw new Error("Use only one of --key or --generate");
    }

    const options: Options = {
      keyPath: args.key ? normalizeKeyPath(expandTilde(args.key)) : undefined,
      alias: args.alias,
      generate,
      port: target.port,
      force: args.force || false,
      dryRun: args["dry-run"] || false,
    };

    if (!options.alias) {
      options.alias = `${target.user}-${simplifyForAlias(target.host)}`;
    }

    // Unified key selection
    const keySelection = await selectKeyUnified(options);
    options.keyPath = keySelection.keyPath;
    options.isAgentKey = keySelection.isAgentKey;

    if (!options.keyPath) {
      throw new Error("No key path selected or generated.");
    }

    await Deno.mkdir(SSH_DIR, { recursive: true });

    if (options.generate) {
      await generateKey(options.keyPath, target.user, target.host);
    }

    // Only verify file exists and set permissions for non-agent keys
    if (!options.isAgentKey) {
      try {
        await Deno.stat(options.keyPath);
      } catch {
        throw new Error(`Private key not found at ${options.keyPath}`);
      }

      await ensurePermissions(options.keyPath);
    }

    const configBlock = generateConfigBlock(
      options.alias!,
      target.host,
      target.user,
      options.keyPath,
      options.port,
    );

    let existingConfig = await readConfig();
    if (hasHostBlock(existingConfig, options.alias!)) {
      if (!options.force) {
        throw new Error(`Host block for '${options.alias}' already exists. Use --force to replace.`);
      }
      console.log(`Replacing existing Host block for '${options.alias}'.`);
      existingConfig = removeHostBlock(existingConfig, options.alias!);
    }

    if (options.dryRun) {
      console.log("\n--- SSH Config Block ---\n");
      console.log(configBlock);
      console.log("\n(Dry run: no changes written)");
      return;
    }

    await backupConfig();
    const newConfig = existingConfig.trim().length > 0
      ? `${existingConfig.trim()}\n\n${configBlock}\n`
      : `${configBlock}\n`;
    await Deno.writeTextFile(CONFIG_FILE, newConfig);
    console.log(`\n✅ SSH config updated: ${CONFIG_FILE}`);
    console.log(`\nTo connect, use: ssh ${options.alias}`);
  } catch (e) {
    handleError(e);
  }
}

// Run main function
if (import.meta.main) {
  main();
}