#!/bin/bash
set -e

#==============================================================================
# Configuration & Constants
#==============================================================================
readonly AGENT_DIR="./agent"
readonly DEFAULT_PORT=7567
readonly DEFAULT_TAG="latest"
readonly DEFAULT_CORE_URL="https://api.serversinc.io"
readonly DEFAULT_EMAIL="hello@serversinc.io"
readonly DOMAIN_SUFFIX="serversinc.io"
readonly DNS_WAIT_TIMEOUT=30
readonly DNS_WAIT_INTERVAL=2

TOKEN=""
CORE_URL=""
TAG="$DEFAULT_TAG"
PORT="$DEFAULT_PORT"
EMAIL=""

#==============================================================================
# Utilities
#==============================================================================
print_usage() {
  echo "Usage: $0 --token <token> [--core <CORE_URL>] [--tag <image-tag>] [--email <email>]"
  echo "  --token is required."
  echo "  --core defaults to $DEFAULT_CORE_URL if not given."
  echo "  --email defaults to $DEFAULT_EMAIL if not given (used for the Traefik Let's Encrypt account)."
  exit 1
}

confirm_or_exit() {
  local message="$1"
  
  if [[ -t 0 ]]; then
    read -rp "$message" confirm
    if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
      echo "Aborting setup."
      exit 1
    fi
  else
    echo "Non-interactive shell detected, auto-confirming."
  fi
}

#==============================================================================
# Installation Functions
#==============================================================================
ensure_jq_installed() {
  if ! command -v jq &> /dev/null; then
    echo "jq not found. Installing..."
    sudo apt-get update -y
    sudo apt-get install -y jq
  fi
}

ensure_docker_installed() {
  if ! command -v docker &> /dev/null; then
    echo "Docker not found. Installing..."
    curl -sSL https://get.serversinc.io/docker.sh | bash
    echo "Docker installed. You may need to log out and back in for permissions to apply."
  fi
}

#==============================================================================
# System Information
#==============================================================================
gather_system_info() {
  SERVER_NAME=$(hostname)
  CPU=$(nproc)
  MEMORY=$(grep MemTotal /proc/meminfo | awk '{print int($2/1024)}')
  DISK=$(df -BG --output=size / | tail -1 | tr -d '[:space:]' | sed 's/G//')
  IPV4=$(curl -s https://ipinfo.io/ip || hostname -I | awk '{print $1}')
  
  echo "Detected system details:"
  echo "  CPU Cores: $CPU"
  echo "  Memory: ${MEMORY} MB"
  echo "  Disk: ${DISK} GB"
  echo "  IPv4: $IPV4"
  echo
}

build_system_payload() {
  jq -n \
    --arg cpu "$CPU" \
    --arg memory "$MEMORY" \
    --arg disk "$DISK" \
    --arg ip_v4 "$IPV4" \
    --arg name "$SERVER_NAME" \
    '{
      cpu: ($cpu | tonumber),
      memory: ($memory | tonumber),
      disk: ($disk | tonumber),
      ip_v4: $ip_v4,
      name: $name
    }'
}

#==============================================================================
# Core API Integration
#==============================================================================
register_with_core() {
  local payload="$1"
  
  echo "Registering server with Core..."
  RESPONSE=$(curl -s -X POST "${CORE_URL%/}/v1/exchange/$TOKEN" \
    -H "Content-Type: application/json" \
    -d "$payload")
}

extract_credentials_from_response() {
  SSH_KEY=$(echo "$RESPONSE" | jq -r '.data.ssh_key // empty')
  PUBLIC_KEY=$(echo "$RESPONSE" | jq -r '.data.public_key // empty')
  SECRET_KEY=$(echo "$RESPONSE" | jq -r '.data.secret_key // empty')
  SERVER_ID=$(echo "$RESPONSE" | jq -r '.data.server_id // empty')

  local email_from_core=$(echo "$RESPONSE" | jq -r '.data.letsencrypt_email // empty')
  if [ -z "$EMAIL" ] && [ -n "$email_from_core" ] && [ "$email_from_core" != "null" ]; then
    EMAIL="$email_from_core"
  fi

  if [ -z "$EMAIL" ]; then
    EMAIL="$DEFAULT_EMAIL"
  fi

  export EMAIL
}

wait_for_dns_propagation() {
  local domain="${SERVER_ID,,}.${DOMAIN_SUFFIX}"
  local elapsed=0

  echo "Waiting for $domain to resolve before starting Traefik (Traefik requests its ACME cert as soon as the Agent container starts, and won't retry on its own if that request loses the race against DNS propagation)..."

  while [ "$elapsed" -lt "$DNS_WAIT_TIMEOUT" ]; do
    if getent hosts "$domain" >/dev/null 2>&1; then
      echo "$domain resolved after ${elapsed}s."
      return 0
    fi
    sleep "$DNS_WAIT_INTERVAL"
    elapsed=$((elapsed + DNS_WAIT_INTERVAL))
  done

  echo "Warning: $domain did not resolve within ${DNS_WAIT_TIMEOUT}s. Continuing anyway — Traefik's first ACME attempt for this domain may fail and require a manual 'docker restart traefik' once DNS catches up." >&2
}

validate_credentials() {
  if [ -z "$SSH_KEY" ] || [ "$SSH_KEY" == "null" ]; then
    echo "Error: failed to receive SSH key from Core"
    exit 1
  fi
  
  if [ -z "$PUBLIC_KEY" ] || [ "$PUBLIC_KEY" == "null" ] || \
     [ -z "$SECRET_KEY" ] || [ "$SECRET_KEY" == "null" ] || \
     [ -z "$SERVER_ID" ] || [ "$SERVER_ID" == "null" ]; then
    echo "Error: response did not include PUBLIC_KEY, SECRET_KEY, and SERVER_ID." >&2
    exit 1
  fi
}

setup_ssh_access() {
  mkdir -p ~/.ssh
  chmod 700 ~/.ssh
  
  if ! grep -qxF "$SSH_KEY" ~/.ssh/authorized_keys 2>/dev/null; then
    echo "$SSH_KEY" >> ~/.ssh/authorized_keys
  fi
  
  chmod 600 ~/.ssh/authorized_keys
  echo "Core SSH key added to ~/.ssh/authorized_keys"
}

#==============================================================================
# Traefik Management
#==============================================================================
check_traefik_status() {
  local container_line=$(docker ps --format '{{.ID}} {{.Image}} {{.Names}}' | grep -i traefik | head -n1 || true)
  
  if [ -z "$container_line" ]; then
    container_line=$(docker ps -a --format '{{.ID}} {{.Image}} {{.Names}}' | grep -i traefik | head -n1 || true)
  fi
  
  if [ -z "$container_line" ]; then
    echo "no-container"
    return 0
  fi
  
  local container_id=$(echo "$container_line" | awk '{print $1}')
  local container_name=$(echo "$container_line" | awk '{print $3}')
  
  if is_traefik_configured "$container_id"; then
    echo "ok:$container_name"
  else
    echo "misconfigured:$container_name"
  fi
}

is_traefik_configured() {
  local container_id="$1"
  local socket_ok=false
  local provider_ok=false
  
  # Check Docker socket mount
  if docker inspect -f '{{json .Mounts}}' "$container_id" 2>/dev/null | grep -q '/var/run/docker.sock' || \
     docker inspect -f '{{json .HostConfig.Binds}}' "$container_id" 2>/dev/null | grep -q '/var/run/docker.sock'; then
    socket_ok=true
  fi
  
  # Check Docker provider configuration
  if docker inspect -f '{{range .Config.Env}}{{println .}}{{end}}' "$container_id" 2>/dev/null | grep -iq 'TRAEFIK_PROVIDERS_DOCKER\|PROVIDERS_DOCKER' || \
     docker inspect -f '{{json .Config.Cmd}}' "$container_id" 2>/dev/null | grep -qi 'providers.docker' || \
     docker inspect -f '{{json .Args}}' "$container_id" 2>/dev/null | grep -qi 'providers.docker'; then
    provider_ok=true
  fi
  
  [ "$socket_ok" = true ] && [ "$provider_ok" = true ]
}

install_traefik() {
  echo "Traefik container not found. Installing Traefik..."
  curl -sSL https://get.serversinc.io/traefik.sh | bash -s
  sleep 3
}

fix_traefik_configuration() {
  local container_name="$1"
  echo "Traefik container '$container_name' is misconfigured." >&2
  echo "Running Traefik installer to fix configuration..."
  curl -sSL https://get.serversinc.io/traefik.sh | bash
  sleep 3
}

ensure_traefik_running() {
  local status=$(check_traefik_status)
  
  case "$status" in
    no-container)
      install_traefik
      status=$(check_traefik_status)
      ;;
    misconfigured:*)
      fix_traefik_configuration "${status#misconfigured:}"
      status=$(check_traefik_status)
      ;;
  esac
  
  if [[ "$status" == ok:* ]]; then
    echo "Traefik configured correctly: ${status#ok:}"
  else
    echo "Warning: Traefik not properly configured. Status: $status" >&2
  fi
}

#==============================================================================
# Agent Setup
#==============================================================================
create_agent_config() {
  mkdir -p "$AGENT_DIR"

  # The Agent reads its trusted public key from a file (PUBLIC_KEY_PATH), not
  # from a raw env var — write the PEM Core returned to disk and point the
  # container-internal path at it (agent.sh mounts $AGENT_DIR at /agent).
  echo "$PUBLIC_KEY" > "$AGENT_DIR/agent_public_key.pem"

  cat > "$AGENT_DIR/.env" <<EOF
PORT=$PORT
CORE_URL=$CORE_URL
PUBLIC_KEY_PATH=/agent/agent_public_key.pem
SECRET_KEY=$SECRET_KEY
SERVER_ID=$SERVER_ID
EOF
}

install_agent() {
  echo "Installing Agent (image tag: $TAG) with CORE_URL=$CORE_URL..."
  curl -sSL https://get.serversinc.io/agent.sh | bash -s -- "$CORE_URL" "$TAG"
}

#==============================================================================
# Main Execution
#==============================================================================
parse_arguments() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      -t|--token)
        TOKEN="$2"; shift 2;;
      -c|--core|--core-url)
        CORE_URL="$2"; shift 2;;
      -g|--tag)
        TAG="$2"; shift 2;;
      -e|--letsencrypt-email)
        EMAIL="$2"; shift 2;;
      -h|--help)
        print_usage;;
      *)
        echo "Unknown arg: $1"; print_usage;;
    esac
  done
}

validate_and_default_arguments() {
  if [ -z "$TOKEN" ]; then
    echo "Error: --token is required." >&2
    print_usage
  fi

  if [ -z "$CORE_URL" ]; then
    CORE_URL="$DEFAULT_CORE_URL"
  fi
}

main() {
  parse_arguments "$@"
  validate_and_default_arguments
  
  ensure_jq_installed
  
  gather_system_info
  confirm_or_exit "Are these details correct? (y/n): "
  
  local payload=$(build_system_payload)
  register_with_core "$payload"
  extract_credentials_from_response
  validate_credentials
  
  setup_ssh_access
  ensure_docker_installed
  ensure_traefik_running

  wait_for_dns_propagation

  create_agent_config
  install_agent
  
  echo "Setup complete! 🎉"
}

main "$@"