#!/bin/bash

####################################################################
#### Code-Server Install Script                                 ####
####                                                            ####
#### Installs code-server with Nginx, Apache, or Caddy reverse  ####
#### proxy, Let's Encrypt TLS certificate, and systemd service. ####
####                                                            ####
#### Supported: RHEL/Rocky/Alma 8+, Oracle Linux 8+,            ####
####            Debian 11+, Ubuntu 20.04+                       ####
####                                                            ####
#### Author:  Phil Connor                                       ####
#### Contact: contact@mylinux.work                              ####
#### Website: https://mylinux.work                              ####
#### License: MIT                                               ####
#### Version: 2.0                                               ####
####                                                            ####
#### Usage:                                                     ####
####   Edit the User Configuration section below, then run:     ####
####   sudo ./install-code-server.sh                            ####
####                                                            ####
####   Or pass options on the command line:                     ####
####   sudo ./install-code-server.sh --server code.example.com  ####
####        --email admin@example.com --http nginx              ####
####        --password 'YourPassword'                           ####
####                                                            ####
####   sudo ./install-code-server.sh --help                     ####
####################################################################

set -euo pipefail

# ============================================================================
# USER CONFIGURATION -- edit these or override with command-line options
# ============================================================================

CODEDIR="/code"                       # Home directory for your code
EMAIL="admin@mydomain.com"            # Email for Let's Encrypt registration
HTTPTYPE="NGINX"                      # APACHE, NGINX, or CADDY
PASSWD=""                             # code-server password (prompted if empty)
UNAME=""                              # Username for Caddy basic auth (optional)
SERVDIR="/usr/local/code-server"      # code-server install directory
SERVERNAME=""                         # Server FQDN (required)
USRDIR="/var/lib/code-server"         # User data directory (settings, extensions)
INSTALL_CRON=0                        # Set to 1 to install auto-update cron
CRON_DAYS=25                          # Days between automatic updates

# ============================================================================
# INTERNAL VARIABLES
# ============================================================================

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'
VERSION="2.0"
ARCH=""

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

log()      { echo -e "${GREEN}[+]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[!]${NC} $1"; }
log_err()  { echo -e "${RED}[x]${NC} $1" >&2; }
log_info() { echo -e "${CYAN}[>]${NC} $1"; }

show_help() {
    cat <<EOF
Code-Server Install Script v${VERSION}

Usage: sudo $0 [OPTIONS]

Installs code-server with a reverse proxy (Nginx/Apache/Caddy),
Let's Encrypt TLS certificate, and systemd service.

OPTIONS:
    --server FQDN       Server fully qualified domain name (required)
    --email EMAIL       Email for Let's Encrypt (required)
    --http TYPE         Reverse proxy: NGINX, APACHE, or CADDY (default: NGINX)
    --password PASS     code-server password (prompted if not set)
    --username USER     Username for Caddy basic auth (optional)
    --codedir DIR       Code directory (default: /code)
    --installdir DIR    code-server install directory (default: /usr/local/code-server)
    --datadir DIR       User data directory (default: /var/lib/code-server)
    --install-cron      Install the update script and set up a cron job
                        for automatic updates (default: every 25 days)
    --cron-days DAYS    Days between automatic updates (default: 25)
    -h, --help          Show this help

EXAMPLES:
    $0 --server code.example.com --email admin@example.com
    $0 --server code.example.com --email admin@example.com --http apache
    $0 --server code.example.com --email admin@example.com --http caddy --username admin
    $0 --server code.example.com --email admin@example.com --install-cron

REQUIREMENTS:
    - Root privileges
    - RHEL/Rocky/Alma 8+, Oracle Linux 8+, Debian 11+, or Ubuntu 20.04+
    - A valid FQDN with DNS A record pointing to this server
    - At least 1 GB RAM and 20 GB disk

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --server)     SERVERNAME="$2"; shift 2 ;;
            --email)      EMAIL="$2"; shift 2 ;;
            --http)       HTTPTYPE=$(echo "$2" | tr '[:lower:]' '[:upper:]'); shift 2 ;;
            --password)   PASSWD="$2"; shift 2 ;;
            --username)   UNAME="$2"; shift 2 ;;
            --codedir)       CODEDIR="$2"; shift 2 ;;
            --installdir)    SERVDIR="$2"; shift 2 ;;
            --datadir)       USRDIR="$2"; shift 2 ;;
            --install-cron)  INSTALL_CRON=1; shift ;;
            --cron-days)     CRON_DAYS="$2"; shift 2 ;;
            -h|--help)       show_help ;;
            *) log_err "Unknown option: $1"; show_help ;;
        esac
    done
}

detect_arch() {
    local machine
    machine=$(uname -m)
    case "$machine" in
        x86_64)  ARCH="amd64" ;;
        aarch64) ARCH="arm64" ;;
        *)
            log_err "Unsupported architecture: $machine"
            exit 1
            ;;
    esac
}

detect_os() {
    if [ -f /etc/os-release ]; then
        # shellcheck disable=SC1091
        . /etc/os-release
        OS_ID=$(echo "$ID" | tr '[:upper:]' '[:lower:]')
        OS_VERSION="${VERSION_ID%%.*}"
    else
        log_err "Cannot detect OS -- /etc/os-release not found"
        exit 1
    fi

    # shellcheck disable=SC2034
    case "$OS_ID" in
        ubuntu|debian)
            PAKMGR="apt-get -y"
            PKG_TYPE="deb"
            ;;
        rhel|centos|rocky|almalinux|ol|oracle)
            PAKMGR="dnf -y"
            PKG_TYPE="rpm"
            ;;
        *)
            log_err "Unsupported OS: $OS_ID"
            exit 1
            ;;
    esac
}

preflight_checks() {
    # Root check
    if [[ $EUID -ne 0 ]]; then
        log_err "This script must be run as root"
        exit 1
    fi

    # Required tools
    for cmd in curl tar; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            log_err "Required command not found: $cmd"
            exit 1
        fi
    done

    # FQDN check
    if [[ -z "$SERVERNAME" ]]; then
        log_err "Server FQDN is required (--server or edit SERVERNAME in script)"
        exit 1
    fi

    # HTTP type check
    case "$HTTPTYPE" in
        APACHE|NGINX|CADDY) ;;
        *)
            log_err "Invalid HTTP type: $HTTPTYPE (must be APACHE, NGINX, or CADDY)"
            exit 1
            ;;
    esac

    # Password -- prompt if empty
    if [[ -z "$PASSWD" ]]; then
        echo -n "Enter code-server password: "
        read -rs PASSWD
        echo
        if [[ -z "$PASSWD" ]]; then
            log_err "Password cannot be empty"
            exit 1
        fi
    fi

    # DNS check
    if ! host "$SERVERNAME" >/dev/null 2>&1 && ! dig +short "$SERVERNAME" 2>/dev/null | grep -q .; then
        log_warn "DNS lookup for $SERVERNAME failed -- certbot may fail if DNS is not configured"
    fi
}

# ============================================================================
# SYSTEM UPDATE
# ============================================================================

update_system() {
    log "Updating system packages"
    if [[ "$PKG_TYPE" == "deb" ]]; then
        apt-get -y update
        apt-get -y upgrade
    else
        dnf -y update
    fi
}

# ============================================================================
# INSTALL CODE-SERVER
# ============================================================================

install_codeserver() {
    log "Installing code-server"

    # Get latest version
    local version
    version=$(curl -fsSL "https://api.github.com/repos/coder/code-server/releases/latest" | \
        grep '"tag_name"' | head -1 | grep -oP 'v\K[0-9]+\.[0-9]+\.[0-9]+')

    if [[ -z "$version" ]]; then
        log_err "Failed to determine latest code-server version"
        exit 1
    fi
    log_info "Latest version: $version"

    # Download
    local tarball="code-server-${version}-linux-${ARCH}.tar.gz"
    local url="https://github.com/coder/code-server/releases/download/v${version}/${tarball}"

    cd /tmp
    log_info "Downloading $tarball"
    curl -fsSL -o "$tarball" "$url"
    tar xzf "$tarball"

    # Install
    mkdir -p "$SERVDIR"
    cp -r "code-server-${version}-linux-${ARCH}"/* "$SERVDIR"/

    # Create symlink
    ln -sf "${SERVDIR}/bin/code-server" /usr/bin/code-server

    # Create directories
    mkdir -p "$CODEDIR"
    mkdir -p "$USRDIR"

    # Cleanup
    rm -rf "/tmp/$tarball" "/tmp/code-server-${version}-linux-${ARCH}"

    # Verify
    if ! command -v code-server >/dev/null 2>&1; then
        log_err "code-server installation failed"
        exit 1
    fi
    log "code-server $(code-server --version | head -1) installed to $SERVDIR"
}

# ============================================================================
# SYSTEMD SERVICE
# ============================================================================

create_service() {
    log "Creating systemd service"

    local after_service="network.target"
    case "$HTTPTYPE" in
        NGINX)  after_service="nginx.service" ;;
        APACHE) after_service="httpd.service" ;;
        CADDY)  after_service="caddy.service" ;;
    esac

    cat > /lib/systemd/system/code-server.service <<EOF
[Unit]
Description=code-server
After=$after_service

[Service]
Type=simple
Environment=PASSWORD=$PASSWD
ExecStart=/usr/bin/code-server --bind-addr 127.0.0.1:8080 --user-data-dir $USRDIR --auth password
Restart=always
RestartSec=5

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
    systemctl enable code-server
    systemctl start code-server
    log "code-server service started (listening on 127.0.0.1:8080)"
}

# ============================================================================
# REVERSE PROXY -- NGINX
# ============================================================================

install_nginx() {
    log "Installing Nginx"

    if [[ "$PKG_TYPE" == "deb" ]]; then
        apt-get -y install nginx
    else
        dnf -y install nginx
    fi

    # Determine config directory
    local confdir="/etc/nginx/conf.d"
    if [[ "$PKG_TYPE" == "deb" ]] && [[ -d /etc/nginx/sites-available ]]; then
        confdir="/etc/nginx/sites-available"
    fi

    cat > "${confdir}/code-server.conf" <<EOF
server {
    listen 80;
    listen [::]:80;
    server_name $SERVERNAME;

    location / {
        proxy_pass http://127.0.0.1:8080/;
        proxy_set_header Host \$host;
        proxy_set_header Upgrade \$http_upgrade;
        proxy_set_header Connection upgrade;
        proxy_set_header Accept-Encoding gzip;
        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto \$scheme;
    }
}
EOF

    # Disable default site
    if [[ -f "${confdir}/default.conf" ]]; then
        mv "${confdir}/default.conf" "${confdir}/default.conf.disabled"
    fi
    if [[ -L /etc/nginx/sites-enabled/default ]]; then
        rm -f /etc/nginx/sites-enabled/default
    fi

    # Enable site (Ubuntu/Debian sites-available pattern)
    if [[ "$confdir" == "/etc/nginx/sites-available" ]]; then
        ln -sf "${confdir}/code-server.conf" /etc/nginx/sites-enabled/code-server.conf
    fi

    nginx -t
    systemctl enable --now nginx
    systemctl reload nginx
    log "Nginx configured and running"
}

# ============================================================================
# REVERSE PROXY -- APACHE
# ============================================================================

install_apache() {
    log "Installing Apache"

    if [[ "$PKG_TYPE" == "deb" ]]; then
        apt-get -y install apache2
        # Enable required modules
        a2enmod proxy proxy_http proxy_wstunnel rewrite headers ssl
        local conffile="/etc/apache2/sites-available/code-server.conf"
    else
        dnf -y install httpd mod_ssl
        local conffile="/etc/httpd/conf.d/code-server.conf"
    fi

    cat > "$conffile" <<EOF
<VirtualHost *:80>
    ServerName $SERVERNAME

    RewriteEngine On
    RewriteCond %{HTTP:Upgrade} =websocket [NC]
    RewriteRule /(.*)           ws://127.0.0.1:8080/\$1 [P,L]
    RewriteCond %{HTTP:Upgrade} !=websocket [NC]
    RewriteRule /(.*)           http://127.0.0.1:8080/\$1 [P,L]

    ProxyRequests off
    ProxyPreserveHost On
    ProxyPass / http://127.0.0.1:8080/ nocanon
    ProxyPassReverse / http://127.0.0.1:8080/

    RequestHeader set X-Forwarded-Proto "http"
</VirtualHost>
EOF

    if [[ "$PKG_TYPE" == "deb" ]]; then
        a2ensite code-server.conf
        a2dissite 000-default.conf 2>/dev/null || true
        systemctl enable --now apache2
        systemctl reload apache2
    else
        systemctl enable --now httpd
        systemctl reload httpd
    fi

    log "Apache configured and running"
}

# ============================================================================
# REVERSE PROXY -- CADDY
# ============================================================================

install_caddy() {
    log "Installing Caddy"

    if [[ "$PKG_TYPE" == "deb" ]]; then
        apt-get -y install debian-keyring debian-archive-keyring apt-transport-https curl
        curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/gpg.key' | \
            gpg --dearmor -o /usr/share/keyrings/caddy-stable-archive-keyring.gpg
        curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/debian.deb.txt' | \
            tee /etc/apt/sources.list.d/caddy-stable.list
        apt-get -y update
        apt-get -y install caddy
    else
        dnf -y install 'dnf-command(copr)'
        dnf -y copr enable @caddy/caddy
        dnf -y install caddy
    fi

    # Caddy handles TLS automatically -- no certbot needed
    if [[ -n "$UNAME" ]]; then
        # With basic auth
        local caddy_hash
        caddy_hash=$(caddy hash-password --plaintext "$PASSWD" 2>/dev/null || \
            echo "$PASSWD" | caddy hash-password 2>/dev/null || echo "")

        if [[ -z "$caddy_hash" ]]; then
            log_warn "Could not generate Caddy password hash -- writing config without basic auth"
            cat > /etc/caddy/Caddyfile <<EOF
$SERVERNAME {
    reverse_proxy 127.0.0.1:8080
}
EOF
        else
            cat > /etc/caddy/Caddyfile <<EOF
$SERVERNAME {
    basicauth / {
        $UNAME $caddy_hash
    }
    reverse_proxy 127.0.0.1:8080
}
EOF
        fi
    else
        # Without basic auth (code-server handles its own auth)
        cat > /etc/caddy/Caddyfile <<EOF
$SERVERNAME {
    reverse_proxy 127.0.0.1:8080
}
EOF
    fi

    # Caddy handles its own auth -- disable code-server auth
    sed -i 's/--auth password/--auth none/' /lib/systemd/system/code-server.service
    systemctl daemon-reload
    systemctl restart code-server

    systemctl enable --now caddy
    log "Caddy configured and running (TLS handled automatically)"
}

# ============================================================================
# LET'S ENCRYPT (for Nginx and Apache only -- Caddy handles its own TLS)
# ============================================================================

install_certbot() {
    if [[ "$HTTPTYPE" == "CADDY" ]]; then
        log_info "Skipping certbot -- Caddy handles TLS automatically"
        return
    fi

    log "Installing certbot for Let's Encrypt"

    if [[ "$PKG_TYPE" == "deb" ]]; then
        apt-get -y install snapd
        snap install core 2>/dev/null; snap refresh core 2>/dev/null
        snap install --classic certbot 2>/dev/null
        ln -sf /snap/bin/certbot /usr/bin/certbot
    else
        dnf -y install epel-release
        dnf -y install certbot
        if [[ "$HTTPTYPE" == "NGINX" ]]; then
            dnf -y install python3-certbot-nginx
        elif [[ "$HTTPTYPE" == "APACHE" ]]; then
            dnf -y install python3-certbot-apache
        fi
    fi

    # Request certificate
    log_info "Requesting TLS certificate for $SERVERNAME"
    if [[ "$HTTPTYPE" == "NGINX" ]]; then
        certbot --non-interactive --redirect --agree-tos \
            --nginx -d "$SERVERNAME" -m "$EMAIL"
        systemctl reload nginx
    elif [[ "$HTTPTYPE" == "APACHE" ]]; then
        certbot --non-interactive --redirect --agree-tos \
            --apache -d "$SERVERNAME" -m "$EMAIL"
        if [[ "$PKG_TYPE" == "deb" ]]; then
            systemctl reload apache2
        else
            systemctl reload httpd
        fi
    fi

    # Certbot auto-renewal is handled by the systemd timer installed by certbot
    if systemctl is-enabled certbot.timer >/dev/null 2>&1; then
        log_info "Certbot auto-renewal timer is active"
    elif systemctl is-enabled snap.certbot.renew.timer >/dev/null 2>&1; then
        log_info "Certbot snap auto-renewal timer is active"
    else
        log_warn "No certbot auto-renewal timer found -- add a cron job for 'certbot renew'"
    fi

    log "TLS certificate installed"
}

# ============================================================================
# SELINUX (RHEL-based only)
# ============================================================================

configure_selinux() {
    if ! command -v getenforce >/dev/null 2>&1; then
        return
    fi

    if [[ "$(getenforce 2>/dev/null)" == "Enforcing" ]]; then
        log_info "Configuring SELinux for reverse proxy"
        setsebool -P httpd_can_network_connect 1 2>/dev/null || true
    fi
}

# ============================================================================
# AUTO-UPDATE CRON
# ============================================================================

install_update_cron() {
    if [[ $INSTALL_CRON -eq 0 ]]; then
        return
    fi

    log "Installing auto-update cron job"

    local update_url="https://mylinux.work/downloads/update-code-server.sh"
    local script_dest="/usr/local/bin/update-code-server.sh"

    curl -fsSL -o "$script_dest" "$update_url"
    chmod 700 "$script_dest"
    log_info "Update script installed to $script_dest"

    # Build cron entry -- run every N days at 3:00 AM
    local cron_line="0 3 */${CRON_DAYS} * * ${script_dest} 2>&1 | logger -t update-code-server"
    local cron_marker="# code-server auto-update"

    # Remove existing code-server cron entries
    local existing_cron
    existing_cron=$(crontab -l 2>/dev/null || true)
    local new_cron
    new_cron=$(echo "$existing_cron" | grep -v "$cron_marker" | grep -v "update-code-server" || true)

    # Add new entry
    echo "${new_cron}
${cron_line} ${cron_marker}" | crontab -

    log "Cron job installed -- updates every $CRON_DAYS days at 3:00 AM"
    log_info "View with: crontab -l"
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    parse_args "$@"

    echo ""
    echo "=============================================="
    echo "  Code-Server Install Script v${VERSION}"
    echo "=============================================="
    echo ""

    detect_os
    detect_arch
    preflight_checks

    log_info "OS: $OS_ID $OS_VERSION ($PKG_TYPE)"
    log_info "Architecture: $ARCH"
    log_info "Server: $SERVERNAME"
    log_info "Reverse proxy: $HTTPTYPE"
    log_info "Code directory: $CODEDIR"
    log_info "Install directory: $SERVDIR"
    log_info "Data directory: $USRDIR"
    echo ""

    update_system
    install_codeserver
    create_service

    case "$HTTPTYPE" in
        NGINX)  install_nginx ;;
        APACHE) install_apache ;;
        CADDY)  install_caddy ;;
    esac

    configure_selinux
    install_certbot
    install_update_cron

    echo ""
    echo "=============================================="
    echo -e "  ${GREEN}Installation complete${NC}"
    echo "=============================================="
    echo ""
    echo "  code-server is running at:"
    echo "    https://$SERVERNAME"
    echo ""
    echo "  Reverse proxy: $HTTPTYPE"
    echo "  Data directory: $USRDIR"
    echo "  Code directory: $CODEDIR"
    echo ""
    echo "  Manage the service:"
    echo "    systemctl status code-server"
    echo "    systemctl restart code-server"
    echo ""
    echo "=============================================="
}

main "$@"