import std/[os, osproc, sequtils, strformat, strutils, sugar]
import hwylterm


proc checkHost(host: seq[string]): string =
  case host.len:
  of 1: return host[0]
  of 0: quit $bb"[red]expected hostname"
  else: quit $bb"[red]expected one positional argument"

proc check(name: string): bool = (execCmd &"ssh -O check {name}") == 0
proc startSsh(name: string) = discard execCmd &"ssh -M -f -N {name}"
proc exitSsh(name: string):int = execCmd &"ssh -O exit {name}"

proc activateTunnel(name: string, port: int) =
  echo "activating connection to: ", name
  if not check(name):
    startSsh(name)
  quit execCmd &"""ssh -fNL "{port}:localhost:{port}" {name}"""

proc deactivateTunnels(name: string) =
  echo "deactivating connection to: ", name
  quit exitSsh(name)

proc collectControllers(): seq[string] =
  let sshDir = (getEnv "HOME") / ".ssh"
  collect(for _,p in walkDir(sshDir): p)
    .map(extractFilename)
    .filterIt(it.startsWith("control"))

proc killTunnels() =
  let controllers = collectControllers()
  if controllers.len == 0: quit 0
  for c in controllers:
    let host = c.lastPathPart().split("-")[2]
    let code = exitSsh(host)
    if code != 0:
      quit $bb"[red][[Error]failed to exit connection for host: " & host

proc showTunnels() =
  let controllers = collectControllers()
  echo bbfmt"[yellow]{controllers.len}[/] active connections"
  if controllers.len == 0: quit 0
  echo "hosts:"
  echo controllers.mapIt("  " & it.split('-')[1]).join("\n")

when isMainModule:
  import hwylterm/[hwylcli]
  hwylCli:
    name "tunnel"
    settings NoArgsShowHelp
    subcommands:
      --- up
      ? "[b]tunnel up[/] [green]host[/] [[[faint]-h|-p[/]]"
      ... "activate a tunnel"
      required port
      flags:
        port:
          - p
          T int
          ? "port number"
          * 8555
      run:
        let host = checkHost(args)
        activateTunnel(host, port)
      --- down
      ? "[b]tunnel down[/] [green]host[/] [[[faint]-h[/]]"
      ... "disable all tunnels for a host"
      run:
        let host = checkHost(args)
        deactivateTunnels(host)
      --- show
      ... "show active connections"
      run: showTunnels()
      --- kill
      ... "kill all open tunnels"
      run: killTunnels()