#!/bin/bash
set -euo pipefail

source ${SCRIPT_DIR}/logging
source ${SCRIPT_DIR}/kasm_api

kasm_id=$1
container_name=$2
container_ns_id=$3
container_pid=$4
script_data_encoded=$5
notification_pipe=$6
container_setup_progress=0

identity_dir=/var/run/kasm-sidecar/$container_ns_id/ziti
ziti_log_file=/var/run/kasm-sidecar/$container_ns_id/ziti.log

append_ziti_log() {
  printf '%s\n' "$1" >> "$ziti_log_file"
}

# Query the Ziti controller version from the identity's ztAPI endpoint.
# Prints the version string (without leading 'v') on stdout, returns 1 on failure.
get_controller_version() {
  local identity_file=$1
  local zt_api
  zt_api=$(jq -r '.ztAPI // empty' "$identity_file" 2>/dev/null)
  if [ -z "$zt_api" ]; then
    error "Could not extract ztAPI from identity file" >&2
    return 1
  fi
  local version
  version=$(curl -sk --max-time 10 "${zt_api}/version" 2>/dev/null \
    | jq -r '.data.version // empty' \
    | sed 's/^v//')
  if [ -z "$version" ]; then
    error "Could not retrieve controller version from ${zt_api}/version" >&2
    return 1
  fi
  echo "$version"
}

# Return the path to the best-matching installed ziti binary for a given
# controller version.
select_ziti_bin() {
  local controller_version=$1
  local versions
  versions=$(ls /usr/local/bin/ziti-* 2>/dev/null | sed 's|/usr/local/bin/ziti-||' | sort -V)

  if [ -z "$versions" ]; then
    error "No versioned ziti binaries found in /usr/local/bin" >&2
    return 1
  fi

  local controller_minor
  controller_minor=$(echo "$controller_version" | cut -d. -f1,2)

  # Prefer highest installed version with matching major.minor
  local best_minor=""
  while IFS= read -r ver; do
    if [ "$(echo "$ver" | cut -d. -f1,2)" = "$controller_minor" ]; then
      best_minor="$ver"
    fi
  done <<< "$versions"

  if [ -n "$best_minor" ]; then
    echo "/usr/local/bin/ziti-${best_minor}"
    return 0
  fi

  error "No installed ziti binary matches controller major.minor version ${controller_minor}" >&2
  return 1
}

notify_failure() {
  set +e

  # update the workspace status
  set_container_status "starting" "$container_setup_progress" "Ziti tunnel connection failed. Contact an administrator if the issue persists."

  # log the error
  if [ -f "$ziti_log_file" ]; then
    ziti_log=$(cat "$ziti_log_file")
    error "Ziti tunnel log: $ziti_log"
  fi

  # show the error notification
  if [ "$show_vpn_status" == "1" ]; then
    show_container_notification "error" "Ziti Connection Failed" "An error occurred starting the Ziti tunnel. Refer to the logs for more information."
  fi

  # mark the container as failed with a delay to ensure the workspace error is visible to the user since
  # the container relies on REST API calls to update the status
  sleep 10
  nsenter --mount=/proc/$container_pid/ns/mnt --net=/proc/$container_pid/ns/net -- sh -c "echo 'error' > /dockerstartup/.egress_status"
  exit 0
}

if [ -z "$container_name" ] || [ -z "$container_ns_id" ] || [ -z "$container_pid" ] || [ -z "$script_data_encoded" ]; then
  error "Missing parameters. Got $@. Exiting."
  exit 1
fi

trap notify_failure ERR

container_setup_progress=5
set_container_status "starting" "$container_setup_progress" "Connecting to Ziti network..."

# extract script data
script_data=$(echo "$script_data_encoded" | base64 -d)
user_identity_json=$(echo "$script_data" | jq -r '.user_identity_json // empty')
image_identity_json=$(echo "$script_data" | jq -r '.image_identity_json // empty')
show_vpn_status=$(echo $script_data | jq -r '.show_vpn_status')

if [ "$show_vpn_status" == "1" ]; then
  show_container_notification "offline" "Connecting to Ziti" "Please wait while the Ziti tunnel is being established..."
fi

# Extract identity JSONs from script_data (provided by the Kasm backend)
if [ -z "$user_identity_json" ] && [ -z "$image_identity_json" ]; then
  error "No Ziti identities provided in script_data. At least one of user_identity_json or image_identity_json is required."
  notify_failure
  exit 0
fi

# Write identity files to the session directory
mkdir -p "$identity_dir"

if [ -n "$user_identity_json" ]; then
  user_identity_file="${identity_dir}/user_identity.json"
  echo "$user_identity_json" > "$user_identity_file"
  chmod 600 "$user_identity_file"
  info "User identity written to ${user_identity_file}"
fi

if [ -n "$image_identity_json" ]; then
  image_identity_file="${identity_dir}/image_identity.json"
  echo "$image_identity_json" > "$image_identity_file"
  chmod 600 "$image_identity_file"
  info "Image identity written to ${image_identity_file}"
fi

# Select the ziti binary that matches the controller version
first_identity_file=""
if [ -n "$user_identity_json" ]; then
  first_identity_file="${identity_dir}/user_identity.json"
elif [ -n "$image_identity_json" ]; then
  first_identity_file="${identity_dir}/image_identity.json"
fi

if [ -z "$first_identity_file" ]; then
  error "No identity file available to detect controller version"
  notify_failure
  exit 0
fi

controller_version=$(get_controller_version "$first_identity_file")
if [ -z "$controller_version" ]; then
  error "Failed to detect Ziti controller version from ${first_identity_file}"
  notify_failure
  exit 0
fi

ziti_bin=$(select_ziti_bin "$controller_version")
info "Detected Ziti controller version ${controller_version}; using binary: ${ziti_bin}"

if [ ! -x "$ziti_bin" ]; then
  error "Ziti binary not found or not executable: ${ziti_bin}"
  notify_failure
  exit 0
fi

# Docker's embedded DNS at 127.0.0.11 is broken after nssetup replaces the
# container's original network interfaces. The container can reach the internet
# through the bridge NAT, so we use upstream resolvers from host resolver files.
detect_upstream_dns() {
  # Prefer systemd-resolved's upstream view first, then generic resolv.conf.
  for resolv_file in /run/systemd/resolve/resolv.conf /etc/resolv.conf; do
    if [ -r "$resolv_file" ]; then
      while IFS= read -r dns; do
        dns=$(echo "$dns" | sed 's/%.*$//')
        if [ -z "$dns" ] || echo "$dns" | grep -Eq '^(127\.|::1$)'; then
          continue
        fi

        # Pick the first resolver that is actually routable from container netns.
        if echo "$dns" | grep -q ':'; then
          if nsenter --net=/proc/${container_pid}/ns/net -- ip -6 route get "$dns" >/dev/null 2>&1; then
            echo "$dns"
            return 0
          fi
        else
          if nsenter --net=/proc/${container_pid}/ns/net -- ip route get "$dns" >/dev/null 2>&1; then
            echo "$dns"
            return 0
          fi
        fi

        info "Skipping unreachable host upstream DNS candidate: ${dns}"
      done < <(awk '/^nameserver[[:space:]]+/ {print $2}' "$resolv_file")
    fi
  done

  return 1
}

# Start logging to the ziti log
touch "$ziti_log_file"
info "Starting ziti tunnel (log: ${ziti_log_file})"
append_ziti_log "Starting ziti tunnel (log: ${ziti_log_file})"

# Figure out what DNS to use in addition to the tunnel DNS
upstream_dns=$(detect_upstream_dns || true)
if [ -z "$upstream_dns" ]; then
  dns_log_message="No reachable host upstream DNS detected; failing Ziti setup"
  error "$dns_log_message"
  append_ziti_log "$dns_log_message"
  notify_failure
  exit 0
fi
dns_log_message="Using host upstream DNS: ${upstream_dns}"
info "$dns_log_message"
append_ziti_log "$dns_log_message"

# Start ziti tunnel (Go SDK) in the container's network namespace.
# We update container /etc/resolv.conf for container-wide DNS behavior, then
# run ziti in netns + private mount ns with a bound /etc/resolv.conf so the
# Go resolver self-test sees 127.0.0.1 first.
#
# resolv.conf layout:
#   nameserver 127.0.0.1  - tunnel's DNS (self-test requires this first)
#   nameserver <upstream> - host upstream DNS for forwarding non-Ziti queries
nsenter --net=/proc/${container_pid}/ns/net --mount=/proc/${container_pid}/ns/mnt sh -c "
  printf 'nameserver 127.0.0.1\nnameserver %s\n' '${upstream_dns}' > /etc/resolv.conf
"
nsenter --net=/proc/${container_pid}/ns/net -- unshare -m sh -c "
  printf 'nameserver 127.0.0.1\nnameserver %s\n' '${upstream_dns}' > /tmp/ziti-resolv.conf
  mount --bind /tmp/ziti-resolv.conf /etc/resolv.conf

  '${ziti_bin}' tunnel run --identity-dir '${identity_dir}' --verbose >> '${ziti_log_file}' 2>&1 &
  echo \$! >> /var/run/kasm-sidecar/${container_ns_id}/pids
" >> "$ziti_log_file" 2>&1

# Wait for the tunnel to initialize
info "Waiting for Ziti tunnel to initialize"
SECONDS=0
MAX_WAIT=30
FATAL_LOG_REGEX="failed to authenticate|system resolver test failed|lookup .*127\\.0\\.0\\.53:53|can't execute|No such file or directory"
while true; do
  if grep -Eiq "$FATAL_LOG_REGEX" "$ziti_log_file" 2>/dev/null; then
    error "Ziti tunnel initialization failed (fatal error detected in ziti log)"
    notify_failure
    exit 0
  fi

  if grep -q "dns intercept IP range" "$ziti_log_file" 2>/dev/null; then
    # Give the tunnel a short window to surface immediate auth/resolver failures.
    sleep 2
    if grep -Eiq "$FATAL_LOG_REGEX" "$ziti_log_file" 2>/dev/null; then
      error "Ziti tunnel initialization failed after DNS intercept setup"
      notify_failure
      exit 0
    fi
    break
  fi

  if [ $SECONDS -ge $MAX_WAIT ]; then
    error "Ziti tunnel initialization timeout after ${MAX_WAIT} seconds"
    notify_failure
    exit 0
  fi

  sleep 1
done

info "Ziti tunnel initialized"

# Mark the container as ready
nsenter --mount=/proc/$container_pid/ns/mnt --net=/proc/$container_pid/ns/net -- sh -c "echo 'ready' > /dockerstartup/.egress_status"

container_setup_progress=10
set_container_status "running" "$container_setup_progress" "Connected via Ziti network."
show_container_notification "info" "Ziti Connected" "Ziti tunnel established successfully."
