#!/bin/bash

set -e

# ─── Colors ───────────────────────────────────────────────────────────────────
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_ok()   { echo -e "${GREEN}[+]${NC} $*"; }
log_err()  { echo -e "${RED}[!]${NC} $*"; }
log_warn() { echo -e "${YELLOW}[~]${NC} $*"; }
log_info() { echo -e "${BLUE}[*]${NC} $*"; }

# ─── Logging to file (without passwords) ─────────────────────────────────────
LOG_FILE="/var/log/setup.log"
# Redirect to both terminal and log file, but we'll handle passwords separately
exec > >(tee -a "$LOG_FILE") 2>&1
log_info "Log saved to $LOG_FILE"

# ─── Help ─────────────────────────────────────────────────────────────────────
show_help() {
  echo ""
  echo "Usage: $0 [options]"
  echo ""
  echo "Options:"
  echo "  --port     <port>      Cockpit port (default: 12345)"
  echo "  --user     <username>  User to create (default: user)"
  echo "  --password <password>  Password for the user (default: auto-generated)"
  echo "  --timezone <tz>        System timezone (default: UTC, e.g. Europe/Moscow)"
  echo "  --help                 Show this help message"
  echo ""
  echo "Examples:"
  echo "  $0"
  echo "  $0 --port 8443"
  echo "  $0 --port 8443 --user admin"
  echo "  $0 --port 8443 --user admin --password MySecret123!"
  echo "  $0 --port 8443 --timezone Europe/Moscow"
  echo ""
}

# ─── Parse arguments ──────────────────────────────────────────────────────────
COCKPIT_PORT=12345
NEW_USER="user"
USER_PASSWORD=""
TIMEZONE="UTC"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --port)
      COCKPIT_PORT="$2"
      shift 2
      ;;
    --user)
      NEW_USER="$2"
      shift 2
      ;;
    --password)
      USER_PASSWORD="$2"
      shift 2
      ;;
    --timezone)
      TIMEZONE="$2"
      shift 2
      ;;
    --help)
      show_help
      exit 0
      ;;
    *)
      log_err "Unknown argument: $1"
      show_help
      exit 1
      ;;
  esac
done

# ─── Must be root ─────────────────────────────────────────────────────────────
if [[ $EUID -ne 0 ]]; then
  log_err "This script must be run as root"
  exit 1
fi

# ─── Validate OS ──────────────────────────────────────────────────────────────
if ! command -v apt &>/dev/null; then
  log_err "This script requires a Debian/Ubuntu-based system (apt not found)"
  exit 1
fi
log_ok "OS check passed"

# ─── Validate port ────────────────────────────────────────────────────────────
validate_port() {
  local port="$1"
  if ! [[ "$port" =~ ^[0-9]+$ ]] || \
       [ "$port" -lt 1 ] || \
       [ "$port" -gt 65535 ]; then
    return 1
  fi
  return 0
}

if ! validate_port "$COCKPIT_PORT"; then
  log_err "Invalid port: '$COCKPIT_PORT'. Must be a number between 1 and 65535"
  exit 1
fi
log_ok "Port validation passed: $COCKPIT_PORT"

# ─── Check port is not in use ─────────────────────────────────────────────────
if ss -tlnp | grep -q ":${COCKPIT_PORT} "; then
  log_warn "Port $COCKPIT_PORT is already in use"
  while true; do
    # Print prompt directly to terminal, not to log
    printf "${YELLOW}[~]${NC} Enter a new port: " > /dev/tty
    read -r NEW_PORT < /dev/tty
    if ! validate_port "$NEW_PORT"; then
      log_err "Invalid port: '$NEW_PORT'. Must be a number between 1 and 65535"
      continue
    fi
    if ss -tlnp | grep -q ":${NEW_PORT} "; then
      log_err "Port $NEW_PORT is also in use. Try another"
      continue
    fi
    COCKPIT_PORT="$NEW_PORT"
    log_ok "Using port $COCKPIT_PORT"
    break
  done
fi

# ─── Validate timezone ────────────────────────────────────────────────────────
if ! timedatectl list-timezones | grep -qx "$TIMEZONE"; then
  log_err "Invalid timezone: '$TIMEZONE'"
  log_info "Example valid timezones: UTC, Europe/Moscow, America/New_York"
  exit 1
fi
log_ok "Timezone validation passed: $TIMEZONE"

# ─── Generate password if not provided ───────────────────────────────────────
generate_password() {
  local upper lower special digits all
  upper="ABCDEFGHIJKLMNOPQRSTUVWXYZ"
  lower="abcdefghijklmnopqrstuvwxyz"
  special="!@#\$%^&*()-_=+[]{}|;:,.<>?"
  digits="0123456789"
  all="${upper}${lower}${special}${digits}"

  local pass=""
  pass+=$(tr -dc "$upper"   < /dev/urandom | head -c 3)
  pass+=$(tr -dc "$lower"   < /dev/urandom | head -c 3)
  pass+=$(tr -dc "$special" < /dev/urandom | head -c 3)
  pass+=$(tr -dc "$digits"  < /dev/urandom | head -c 2)
  pass+=$(tr -dc "$all"     < /dev/urandom | head -c 4)

  # Shuffle so groups don't appear in order
  echo "$pass" | fold -w1 | shuf | tr -d '\n'
}

if [[ -z "$USER_PASSWORD" ]]; then
  USER_PASSWORD=$(generate_password)
  log_ok "Password auto-generated"
else
  log_info "Using provided password"
fi

log_info "Cockpit port : $COCKPIT_PORT"
log_info "Username     : $NEW_USER"
log_info "Timezone     : $TIMEZONE"

# ─── 1. System update ─────────────────────────────────────────────────────────
log_ok "Updating system..."
apt update && apt upgrade -y

# ─── 2. Set timezone ──────────────────────────────────────────────────────────
log_ok "Setting timezone to $TIMEZONE..."
timedatectl set-timezone "$TIMEZONE"
log_ok "Timezone set to $TIMEZONE"

# ─── 3. Install ufw ───────────────────────────────────────────────────────────
log_ok "Installing ufw..."
apt install -y ufw

ufw default deny incoming
ufw default allow outgoing

# Always ensure SSH is open before enabling UFW
if ! ufw status | grep -q '22/tcp'; then
  log_warn "SSH port (22) not found in UFW rules, adding it now to prevent lockout..."
  ufw allow 22/tcp
fi

ufw allow "${COCKPIT_PORT}/tcp"
ufw --force enable
log_ok "ufw enabled. Open ports: 22/tcp, ${COCKPIT_PORT}/tcp"

# ─── 4. Install fail2ban ──────────────────────────────────────────────────────
log_ok "Installing fail2ban..."
apt install -y fail2ban

cat > /etc/fail2ban/jail.local <<EOF
[DEFAULT]
bantime  = 1h
findtime = 10m
maxretry = 5

[sshd]
enabled = true
port    = 22

[cockpit]
enabled  = true
port     = ${COCKPIT_PORT}
filter   = cockpit
logpath  = /var/log/auth.log
maxretry = 5
EOF

cat > /etc/fail2ban/filter.d/cockpit.conf <<EOF
[Definition]
failregex = ^.*cockpit.*authentication failure.*rhost=<HOST>.*$
            ^.*pam_unix.*cockpit.*authentication failure.*$
ignoreregex =
EOF

systemctl enable fail2ban
systemctl restart fail2ban
log_ok "fail2ban installed and configured"

# ─── 5. Create user with sudo rights ─────────────────────────────────────────
log_ok "Creating user '$NEW_USER'..."
if id "$NEW_USER" &>/dev/null; then
  log_warn "User '$NEW_USER' already exists — updating password and sudo rights"
  usermod -aG sudo "$NEW_USER"
else
  adduser --disabled-password --gecos "" "$NEW_USER"
  usermod -aG sudo "$NEW_USER"
fi

# Use printf to safely handle special characters in password
printf '%s:%s\n' "$NEW_USER" "$USER_PASSWORD" | chpasswd

# Save password to file, excluded from log output
{
  echo "User:     $NEW_USER"
  echo "Password: $USER_PASSWORD"
} > /root/${NEW_USER}_password.txt
chmod 600 /root/${NEW_USER}_password.txt

log_ok "User '$NEW_USER' ready. Credentials saved to /root/${NEW_USER}_password.txt"

# ─── 6. Install Cockpit ───────────────────────────────────────────────────────
log_ok "Installing Cockpit..."
apt install -y cockpit

systemctl enable cockpit.socket
systemctl start cockpit.socket

# ─── 7. Change Cockpit port ───────────────────────────────────────────────────
log_ok "Configuring Cockpit on port $COCKPIT_PORT..."

mkdir -p /etc/cockpit

cat > /etc/cockpit/cockpit.conf <<EOF
[WebService]
ListenStream=${COCKPIT_PORT}
Origins = https://localhost:${COCKPIT_PORT}
ProtocolHeader = X-Forwarded-Proto
EOF

mkdir -p /etc/systemd/system/cockpit.socket.d

cat > /etc/systemd/system/cockpit.socket.d/listen.conf <<EOF
[Socket]
ListenStream=
ListenStream=${COCKPIT_PORT}
EOF

systemctl daemon-reload
systemctl restart cockpit.socket
log_ok "Cockpit restarted on port $COCKPIT_PORT"

# ─── 8. Create certificate script ─────────────────────────────────────────────
log_ok "Creating certificate script..."

CERT_SCRIPT="/usr/local/bin/selfcert-renew.sh"

cat > "$CERT_SCRIPT" <<'CERTSCRIPT'
#!/bin/bash

set -e

DOMAIN_OR_CN="$(hostname)"
CERT_DIR="/etc/ssl/selfcert"
CERT="$CERT_DIR/cert.pem"
KEY="$CERT_DIR/key.pem"

COCKPIT_CERT_DIR="/etc/cockpit/ws-certs.d"
COCKPIT_CERT="$COCKPIT_CERT_DIR/01-self.cert"
COCKPIT_KEY="$COCKPIT_CERT_DIR/01-self.key"

DAYS_VALID=365
RENEW_BEFORE_DAYS=7

mkdir -p "$CERT_DIR"
mkdir -p "$COCKPIT_CERT_DIR"

echo "[+] Using CN: $DOMAIN_OR_CN"

generate_cert() {
  echo "[+] Generating certificate..."

  openssl req -x509 -newkey rsa:4096 -nodes \
    -keyout "$KEY" \
    -out "$CERT" \
    -days "$DAYS_VALID" \
    -subj "/CN=$DOMAIN_OR_CN"

  chmod 600 "$KEY"
  chmod 644 "$CERT"

  echo "[+] Certificate generated"
}

deploy_cert() {
  echo "[+] Deploying certificate to Cockpit..."

  cp "$CERT" "$COCKPIT_CERT"
  cp "$KEY"  "$COCKPIT_KEY"

  chmod 644 "$COCKPIT_CERT"
  chmod 640 "$COCKPIT_KEY"

  # cockpit-ws group may not exist on all systems
  if getent group cockpit-ws &>/dev/null; then
    chown root:cockpit-ws "$COCKPIT_KEY"
  else
    echo "[~] Group cockpit-ws not found, skipping chown"
  fi

  echo "[+] Certificate deployed to $COCKPIT_CERT_DIR"
}

need_renew() {
  if [ ! -f "$CERT" ]; then
    return 0
  fi

  local expiry
  expiry=$(openssl x509 -enddate -noout -in "$CERT" | cut -d= -f2)
  local expiry_sec
  expiry_sec=$(date -d "$expiry" +%s)
  local now_sec
  now_sec=$(date +%s)

  local left_days=$(( (expiry_sec - now_sec) / 86400 ))

  echo "[+] Days left: $left_days"

  if [ "$left_days" -le "$RENEW_BEFORE_DAYS" ]; then
    return 0
  fi

  return 1
}

reload_services() {
  echo "[+] Reloading services..."

  # Only reload nginx if it's actually running
  if systemctl is-active --quiet nginx; then
    systemctl reload nginx
    echo "[+] nginx reloaded"
  fi

  systemctl restart cockpit 2>/dev/null || true

  echo "[+] Services reloaded"
}

# ── main ──
if [ ! -f "$CERT" ] || need_renew; then
  generate_cert
  deploy_cert
  reload_services
else
  echo "[+] Certificate still valid, nothing to do"
fi
CERTSCRIPT

chmod +x "$CERT_SCRIPT"
log_ok "Certificate script created at $CERT_SCRIPT"

bash "$CERT_SCRIPT"

# ─── 9. Create cron job ───────────────────────────────────────────────────────
log_ok "Setting up daily cron job for certificate renewal..."

cat > /etc/cron.d/selfcert-renew <<EOF
0 3 * * * root /usr/local/bin/selfcert-renew.sh >> /var/log/selfcert-renew.log 2>&1
EOF

chmod 644 /etc/cron.d/selfcert-renew
log_ok "Cron job created at /etc/cron.d/selfcert-renew"

# ─── 10. Print summary ────────────────────────────────────────────────────────
SERVER_IP=$(hostname -I | awk '{print $1}')

# Stop logging to file before printing sensitive info
exec 1>/dev/tty 2>/dev/tty

echo ""
echo -e "${GREEN}════════════════════════════════════════════════${NC}"
echo -e "${GREEN}  ✅  Setup complete!${NC}"
echo -e "${GREEN}════════════════════════════════════════════════${NC}"
echo ""
echo -e "  🌐  ${BLUE}Cockpit URL:${NC}"
echo -e "      https://${SERVER_IP}:${COCKPIT_PORT}"
echo ""
echo -e "  👤  ${BLUE}User credentials:${NC}"
echo -e "      Login:    ${NEW_USER}"
echo -e "      Password: ${USER_PASSWORD}"
echo ""
echo -e "  🛡️   ${BLUE}fail2ban:${NC}   active (SSH + Cockpit)"
echo -e "  🕐  ${BLUE}Timezone:${NC}   ${TIMEZONE}"
echo -e "  📄  ${BLUE}Certificate:${NC} /etc/ssl/selfcert/cert.pem"
echo -e "  🔑  ${BLUE}Key:${NC}         /etc/ssl/selfcert/key.pem"
echo -e "  📋  ${BLUE}Setup log:${NC}   $LOG_FILE"
echo -e "  📋  ${BLUE}Cert log:${NC}    /var/log/selfcert-renew.log"
echo ""
echo -e "  ${YELLOW}⚠️  After saving your password, delete the file:${NC}"
echo -e "      ${YELLOW}rm -f /root/${NEW_USER}_password.txt${NC}"
echo ""
echo -e "${GREEN}════════════════════════════════════════════════${NC}"