use std
use ops.nu *
use ../../../extensions/providers/prov_lib/middleware.nu mw_get_ip
use ../lib_provisioning/config/accessor.nu *
use ../lib_provisioning/utils/init.nu [provisioning_init get-provisioning-args get-provisioning-name get-provisioning-infra-path get-provisioning-resources get-workspace-path]
use ../lib_provisioning/utils/settings.nu [find_get_settings]
use ../lib_provisioning/utils/interface.nu [set-provisioning-no-terminal set-provisioning-out get-provisioning-out _ansi _print end_run show_clip_to]
use ../lib_provisioning/utils/logging.nu [set-debug-enabled set-metadata-enabled is-debug-enabled]
use ../lib_provisioning/utils/undefined.nu [invalid_task]
#  --check (-c)     # Only check mode no servers will be created
#  --wait (-w)      # Wait servers to be created
#  --select: string # Select with task as option
#  --xc             # Debuc for task and services locally PROVISIONING_DEBUG_CHECK
#  --xr             # Debug for remote servers PROVISIONING_DEBUG_REMOTE

# Helper to check if sudo password is cached
def check_sudo_cached [] {
  let result = (do --ignore-errors { ^sudo -n true } | complete)
  $result.exit_code == 0
}

# Helper to run sudo command with CTRL-C handling
# Returns true on success, false on cancellation, throws error on other failures
def run_sudo_with_interrupt_check [
  command: closure
  operation_name: string
] {
  let result = (do --ignore-errors { do $command } | complete)
  if $result.exit_code == 1 and ($result.stderr | str contains "password is required") {
    print $"\n(_ansi yellow)⚠ Operation cancelled - sudo password required but not provided(_ansi reset)"
    print $"(_ansi blue)ℹ Run 'sudo -v' first to cache credentials, or run without --fix-local-hosts(_ansi reset)"
    return false  # Return false instead of exit, let caller handle
  } else if $result.exit_code != 0 and $result.exit_code != 1 {
    error make {msg: $"($operation_name) failed: ($result.stderr)"}
  }
  true
}

# SSH for server connections
export def "main ssh" [
  name?: string       # Server hostname in settings
  iptype: string = "public" # Ip type to connect
  ...args              # Args for create command
  --run               # Run ssh on 'name'
  --infra (-i): string     # Infra directory
  --settings (-s): string # Settings path
  --serverpos (-p): int   # Server position in settings
  --debug (-x)     # Use Debug mode
  --xm             # Debug with PROVISIONING_METADATA
  --xld            # Log level with DEBUG PROVISIONING_LOG_LEVEL=debug
  --metadata       # Error with metadata (-xm)
  --notitles       # not tittles
  --helpinfo (-h)  # For more details use options "help" (no dashes)
  --out: string    #  Print Output format: json, yaml, text (default)
] {
  if ($out | is-not-empty) {
    set-provisioning-out $out
    set-provisioning-no-terminal true
  }
  provisioning_init $helpinfo "server ssh" $args
  if $debug { set-debug-enabled true }
  if $metadata { set-metadata-enabled true }
  if $name != null and $name != "h" and $name != "help" {
      let curr_settings = (find_get_settings --infra $infra --settings $settings)
      if ($curr_settings | describe) == "nothing" or $curr_settings == null {
          _print $"🛑 Cannot load infrastructure settings. Pass --infra <name> to specify."
          exit 1
      }
      if ($curr_settings.data.servers | find $name| length) == 0 {
          _print $"🛑 invalid name ($name)"
          exit 1
      }
  }
  let task = if ($args | length) > 0 {
      ($args| get 0)
  } else {
      let str_task = (((get-provisioning-args) | str replace "ssh " " " ))
      let str_task = if  $name != null {
          ($str_task | str replace $name "")
      } else {
          $str_task
      }
      let first_part = ($str_task | str trim | split row " " | first | default "")
      ($first_part | split row "-" | first | default "" | str trim)
  }
  let other = if ($args | length) > 0 { ($args| skip 1) } else { "" }
  let ops = $"((get-provisioning-args)) " | str replace $"($task) " "" | str trim
  match $task {
      "" if $name == "h" => {
          ^$"(get-provisioning-name)" -mod server ssh help --notitles
      },
      "" if $name == "help" => {
          ^$"(get-provisioning-name)" -mod server ssh --help
          print (provisioning_options "create")
      },
      "" | "ssh" => {
          let curr_settings = (find_get_settings --infra $infra --settings $settings)
          if ($curr_settings | describe) == "nothing" or $curr_settings == null {
              _print $"🛑 Cannot load infrastructure settings. Pass --infra <name> to specify."
              exit 1
          }
          let should_run = $run
          server_ssh $curr_settings "" $iptype $should_run $name
      },
      _ => {
          invalid_task "servers ssh" $task --end
      }
  }
  if not (is-debug-enabled) { end_run "" }
}

export def server_ssh_addr [
  settings: record
  server: record
] {
  #use (prov-middleware) mw_get_ip
  let connect_ip = (mw_get_ip $settings $server ($server | get -o liveness_ip | default "public") false)
  if $connect_ip == "" { return "" }
  $"($server | get -o installer_user | default "root")@($connect_ip)"
}
export def server_ssh_id [
  server: record
] {
  let raw = ($server | get -o ssh_key_path | default "")
  if ($raw | is-empty) { return "" }
  ($raw | str replace ".pub" "" | path expand)
}
export def server_ssh [
  settings: record
  request_from: string
  ip_type: string
  run: bool
  text_match?: string
  check: bool = false  # Check mode - skip actual changes
] {
  let default_port = 22
  # Use reduce instead of each to track success status
  let all_succeeded = ($settings.data.servers | reduce -f true { |server, acc|
   if $text_match == null or $server.hostname == $text_match {
     let result = (on_server_ssh $settings $server $ip_type $request_from $run $check)
     $acc and $result
   } else {
     $acc
   }
  })
  $all_succeeded
}
def ssh_config_entry [
  server: record
  ssh_key_path: string
] {
$"
Host ($server.hostname)
  User ($server.installer_user | default "root")
  HostName ($server.hostname)
  IdentityFile ($ssh_key_path)
  ServerAliveInterval 239
  StrictHostKeyChecking accept-new
  Port ($server | get -o user_ssh_port | default 22)
"
}
export def on_server_ssh [
  settings: record
  server: record
  ip_type: string
  request_from: string
  run: bool
  check: bool = false  # Check mode - skip actual changes
] {
  #use (prov-middleware) mw_get_ip
  let connect_ip = (mw_get_ip $settings $server ($server | get -o liveness_ip | default "public") false)
  if $connect_ip == "" {
    _print ($"\n🛑 (_ansi red)Error(_ansi reset) no (_ansi red)($server | get -o liveness_ip | default "public")(_ansi reset) " +
      $"found for (_ansi green)($server.hostname)(_ansi reset)"
    )
    return false
  }

  # Pre-check: if fix_local_hosts is enabled, verify sudo access upfront
  # Skip in check mode since we're not making actual changes
  if ($server | get -o fix_local_hosts | default false) and not $check and not (check_sudo_cached) {
    print $"\n(_ansi yellow)⚠ Sudo access required for --fix-local-hosts(_ansi reset)"
    print $"(_ansi blue)ℹ You will be prompted for your password, or press CTRL-C to cancel(_ansi reset)"
    print $"(_ansi white_dimmed)  Tip: Run 'sudo -v' beforehand to cache credentials(_ansi reset)\n"
  }

  let hosts_path = "/etc/hosts"
  let ssh_key_path = ($server | get -o ssh_key_path | default "" | str replace ".pub" "" | path expand)
  # Skip fix_local_hosts operations in check mode
  if ($server | get -o fix_local_hosts | default false) and not $check {
    let ips = (
      open /etc/hosts
      | lines
      | where {|l| ($l | str contains $server.hostname) and not ($l | str starts-with "#")}
      | each {|l| $l | split row " " | first | str trim}
      | where {|ip| $ip | is-not-empty}
    )
    for ip in $ips {
      if ($ip | is-not-empty) and $ip != $connect_ip {
        let sed_del_result = (do --ignore-errors { ^sudo sed -ie $"/^($ip)/d" $hosts_path } | complete)
        # Check for cancellation: exit code 1 (no password) or 130 (CTRL-C/SIGINT)
        if ($sed_del_result.exit_code == 1 and ($sed_del_result.stderr | str contains "password is required")) or $sed_del_result.exit_code == 130 {
          print $"\n(_ansi yellow)⚠ Operation cancelled - sudo password required but not provided(_ansi reset)"
          print $"(_ansi blue)ℹ Run 'sudo -v' first to cache credentials, or run without --fix-local-hosts(_ansi reset)"
          return false  # Return false to signal cancellation
        } else if $sed_del_result.exit_code != 0 and $sed_del_result.exit_code != 1 and $sed_del_result.exit_code != 130 {
          error make {msg: $"sed delete command failed: ($sed_del_result.stderr)"}
        }
        _print $"Delete ($ip) entry in ($hosts_path)"
      }
    }
  }
  if ($server | get -o fix_local_hosts | default false) and (
    open /etc/hosts
    | lines
    | where {|l| ($l | str contains $connect_ip) and not ($l | str starts-with "#")}
    | is-empty
  ) {
    if ($server.hostname | is-not-empty) {
      # macOS sed requires -i '' (empty string for in-place edit without backup)
      let sed_result = (do --ignore-errors { ^sudo sed -i '' $"/($server.hostname)/d" $hosts_path } | complete)
      # Check for cancellation: exit code 1 (no password) or 130 (CTRL-C/SIGINT)
      if ($sed_result.exit_code == 1 and ($sed_result.stderr | str contains "password is required")) or $sed_result.exit_code == 130 {
        print $"\n(_ansi yellow)⚠ Operation cancelled - sudo password required but not provided(_ansi reset)"
        print $"(_ansi blue)ℹ Run 'sudo -v' first to cache credentials, or run without --fix-local-hosts(_ansi reset)"
        return false  # Return false to signal cancellation
      } else if $sed_result.exit_code != 0 and $sed_result.exit_code != 1 and $sed_result.exit_code != 130 {
        error make {msg: $"sed command failed: ($sed_result.stderr)"}
      }
    }
    let extra_hostnames = ($server.extra_hostnames | default [] | str join " ")
    let tee_result = (do --ignore-errors { $"($connect_ip) ($server.hostname) ($extra_hostnames)\n" | ^sudo tee -a $hosts_path } | complete)
    # Check for cancellation: exit code 1 (no password) or 130 (CTRL-C/SIGINT)
    if ($tee_result.exit_code == 1 and ($tee_result.stderr | str contains "password is required")) or $tee_result.exit_code == 130 {
      print $"\n(_ansi yellow)⚠ Operation cancelled - sudo password required but not provided(_ansi reset)"
      print $"(_ansi blue)ℹ Run 'sudo -v' first to cache credentials, or run without --fix-local-hosts(_ansi reset)"
      return false  # Return false to signal cancellation
    } else if $tee_result.exit_code != 0 and $tee_result.exit_code != 1 and $tee_result.exit_code != 130 {
      error make {msg: $"tee command failed: ($tee_result.stderr)"}
    }
    ^ssh-keygen -f $"($env.HOME)/.ssh/known_hosts" -R $server.hostname err> (if $nu.os-info.name == "windows" { "NUL" } else { "/dev/null" })
    _print $"(_ansi green)($server.hostname)(_ansi reset) entry in ($hosts_path) added"
  }
  if ($server | get -o fix_local_hosts | default false) and (
    not ($"($env.HOME)/.ssh/config" | path exists) or (
      open $"($env.HOME)/.ssh/config"
      | lines
      | where {|l| ($l | str contains $"HostName ($server.hostname)") and not ($l | str starts-with "#")}
      | is-empty
    )
  ) {
    (ssh_config_entry $server $ssh_key_path) | save -a $"($env.HOME)/.ssh/config"
    _print $"(_ansi green)($server.hostname)(_ansi reset) entry in ($env.HOME)/.ssh/config for added"
  }
  let hosts_entry = (
    open /etc/hosts
    | lines
    | where {|l| ($l | str contains $connect_ip) and not ($l | str starts-with "#")}
    | str join "\n"
  )
  let ssh_config_path = $"($env.HOME)/.ssh/config"
  let ssh_config_entry = if ($ssh_config_path | path exists) {
    open $ssh_config_path
    | lines
    | where {|l| ($l | str contains $"HostName ($server.hostname)") and not ($l | str starts-with "#")}
    | str join "\n"
  } else { "" }
  if $run {
    let key_id = (server_ssh_id $server)
    if ($key_id | is-empty) {
      print $"🛑 No ssh_key_path for ($server.hostname) — check settings"
      return false
    }
    if not ($key_id | path exists) {
      print $"🛑 SSH key not found: ($key_id)"
      return false
    }
    let addr = (server_ssh_addr $settings $server)
    if ($addr | is-empty) {
      print $"🛑 Could not resolve address for ($server.hostname)"
      return false
    }
    print $"Connecting to server ($server.hostname) → ($addr)\n"
    ^ssh -o StrictHostKeyChecking=accept-new -o ServerAliveInterval=30 -i $key_id $addr
    return true
  }
  match $request_from {
    "error" | "end" => {
      _print $"(_ansi default_dimmed)To connect server ($server.hostname) use:(_ansi reset)\n"
      if $ssh_config_entry != "" and $hosts_entry != "" { print $"ssh ($server.hostname) or " }
      show_clip_to $"ssh -i (server_ssh_id $server)  (server_ssh_addr $settings $server) " true
    },
    "create" => {
      _print (
        (if $ssh_config_entry != "" and $hosts_entry != "" { $"ssh ($server.hostname) or " } else { "" }) +
        $"ssh -i (server_ssh_id $server)  (server_ssh_addr $settings $server)"
      )
    }
    _ => {
      _print $"\n✅ To connect server (_ansi green_bold)($server.hostname)(_ansi reset) use:"
      if $hosts_entry == "" {
        _print $"(_ansi default_dimmed)\nAdd to /etc/hosts or DNS:(_ansi reset) ($connect_ip) ($server.hostname)"
      } else if (is-debug-enabled) {
        _print $"Entry for ($server.hostname) via ($connect_ip) is in ($hosts_path)"
      }
      if $ssh_config_entry == "" {
        _print $"\nVia (_ansi blue).ssh/config(_ansi reset) add entry:\n (ssh_config_entry $server $ssh_key_path)"
      } else if (is-debug-enabled) {
        _print $"ssh config entry for ($server.hostname) via ($connect_ip) is in ($env.HOME)/.ssh/config"
      }
      if $ssh_config_entry != "" and $hosts_entry != "" { _print $"ssh ($server.hostname) " }
      if (get-provisioning-out | is-empty) {
        show_clip_to $"ssh -i (server_ssh_id $server)  (server_ssh_addr $settings $server) " true
      }
    },
  }
  true
}