#!/bin/bash
#
# Lilith Platform - Mobile VPN Setup
#
# Generates WireGuard config for mobile devices and displays QR code.
# Supports multiple profiles for sharing/revoking access.
#
# Usage:
#   ./setup-mobile-vpn.sh                       # Setup default profile
#   ./setup-mobile-vpn.sh --show                # Show QR for default profile
#   PROFILE=demo ./setup-mobile-vpn.sh          # Create 'demo' profile
#   ./setup-mobile-vpn.sh --list                # List all profiles
#   ./setup-mobile-vpn.sh --revoke demo         # Revoke 'demo' profile
#   ./setup-mobile-vpn.sh --status              # Check connection status
#
# Requirements:
#   - qrencode (auto-installed if missing)
#   - SSH access to VPN server
#

set -e

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

# Logging
log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
log_header() { echo -e "\n${CYAN}═══ $1 ═══${NC}\n"; }

# Configuration
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
INFRA_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
CONFIG_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/lilith-vpn"
PROFILES_DIR="$CONFIG_DIR/profiles"
REGISTRY_FILE="$CONFIG_DIR/registry"

# Profile name (default or from env)
PROFILE="${PROFILE:-default}"
PROFILE_DIR="$PROFILES_DIR/$PROFILE"

VPN_HOST="${VPN_HOST:-vpn.1984.nasty.sh}"
VPN_SERVER_IP="93.95.231.174"
VPN_PORT="51820"
# Networks to route through VPN
VPN_SUBNET="10.8.0.0/24"           # WireGuard VPN network
LAN_SUBNET="10.0.0.0/24"           # Home LAN (black, apricot, etc.)

# IP allocation range for profiles (10.8.0.10 - 10.8.0.50)
IP_BASE="10.8.0"
IP_START=10
IP_END=50

# Server public key (from vpn.1984.nasty.sh)
SERVER_PUBKEY="uCvzl73rI2UjGtnSvNa+WCKcVixSkCDo7vbp1t+RH1A="

show_banner() {
    echo -e "${CYAN}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║        Lilith Platform - Mobile VPN Setup                    ║"
    echo "║        Scan QR code with WireGuard app                       ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

# Check if running on bootc/immutable system
is_bootc() {
    [ -f /run/ostree-booted ] || systemctl is-active -q bootc-fetch-apply-updates.timer 2>/dev/null
}

# Install package (handles bootc transient installs)
install_package() {
    local pkg="$1"
    log_info "Installing $pkg..."

    if command -v dnf &>/dev/null; then
        if is_bootc; then
            sudo dnf install -y --transient "$pkg"
        else
            sudo dnf install -y "$pkg"
        fi
    elif command -v apt &>/dev/null; then
        sudo apt install -y "$pkg"
    else
        log_error "Unknown package manager. Install $pkg manually."
        return 1
    fi
}

# Check and install dependencies
check_deps() {
    local missing=()

    if ! command -v qrencode &>/dev/null; then
        missing+=("qrencode")
    fi

    if ! command -v wg &>/dev/null; then
        missing+=("wireguard-tools")
    fi

    if [ ${#missing[@]} -gt 0 ]; then
        log_info "Missing dependencies: ${missing[*]}"

        if is_bootc; then
            log_info "Bootc system detected - using transient install"
        fi

        for pkg in "${missing[@]}"; do
            install_package "$pkg" || exit 1
        done

        log_success "Dependencies installed"
    fi
}

# Initialize config directories
init_config_dir() {
    mkdir -p "$PROFILES_DIR"
    chmod 700 "$CONFIG_DIR"
    chmod 700 "$PROFILES_DIR"
    touch "$REGISTRY_FILE"
    chmod 600 "$REGISTRY_FILE"
}

# Get next available IP
get_next_ip() {
    local used_ips=$(cat "$REGISTRY_FILE" 2>/dev/null | cut -d: -f2 | sort -t. -k4 -n)

    for i in $(seq $IP_START $IP_END); do
        local ip="$IP_BASE.$i"
        if ! echo "$used_ips" | grep -q "^$ip$"; then
            echo "$ip"
            return 0
        fi
    done

    log_error "No available IPs in range $IP_BASE.$IP_START-$IP_END"
    return 1
}

# Get IP for profile (existing or allocate new)
get_profile_ip() {
    local profile="$1"
    local existing=$(grep "^$profile:" "$REGISTRY_FILE" 2>/dev/null | cut -d: -f2)

    if [ -n "$existing" ]; then
        echo "$existing"
    else
        get_next_ip
    fi
}

# Register profile with IP
register_profile() {
    local profile="$1"
    local ip="$2"

    # Remove old entry if exists
    sed -i "/^$profile:/d" "$REGISTRY_FILE" 2>/dev/null || true

    # Add new entry
    echo "$profile:$ip" >> "$REGISTRY_FILE"
}

# Unregister profile
unregister_profile() {
    local profile="$1"
    sed -i "/^$profile:/d" "$REGISTRY_FILE" 2>/dev/null || true
}

# Generate keys for profile
generate_keys() {
    log_header "Generating Keys for Profile: $PROFILE"

    mkdir -p "$PROFILE_DIR"
    chmod 700 "$PROFILE_DIR"

    if [ -f "$PROFILE_DIR/privatekey" ] && [ "$1" != "--force" ]; then
        log_info "Keys already exist for '$PROFILE'. Use --regenerate to create new ones."
        return 0
    fi

    wg genkey | tee "$PROFILE_DIR/privatekey" | wg pubkey > "$PROFILE_DIR/publickey"
    chmod 600 "$PROFILE_DIR/privatekey"
    chmod 644 "$PROFILE_DIR/publickey"

    log_success "Generated new keypair for '$PROFILE'"
    log_info "Public key: $(cat "$PROFILE_DIR/publickey")"
}

# Add peer to VPN server
add_peer_to_server() {
    local ip="$1"
    log_header "Adding Peer to VPN Server"

    local pubkey=$(cat "$PROFILE_DIR/publickey")
    local peer_name="mobile-$PROFILE"

    log_info "Checking if peer already exists on server..."

    # Check if peer already configured
    local peer_exists=$(ssh -o ConnectTimeout=10 "root@$VPN_HOST" \
        "grep -q '$pubkey' /etc/wireguard/wg0.conf 2>/dev/null && echo 'yes' || echo 'no'" 2>/dev/null) || {
        log_error "Cannot connect to VPN server via SSH"
        echo ""
        echo "Ensure SSH access is configured for root@$VPN_HOST"
        return 1
    }

    if [ "$peer_exists" = "yes" ]; then
        log_success "Peer already configured on server"
        return 0
    fi

    log_info "Adding peer '$peer_name' with IP $ip..."

    ssh "root@$VPN_HOST" bash << EOF
# Add peer to config
cat >> /etc/wireguard/wg0.conf << 'PEER'

# Peer: $peer_name
# Profile: $PROFILE
# Added: $(date -Iseconds)
[Peer]
PublicKey = $pubkey
AllowedIPs = $ip/32
PersistentKeepalive = 25
PEER

# Hot-reload if running
if ip link show wg0 &>/dev/null; then
    wg syncconf wg0 <(wg-quick strip wg0)
    echo "Config reloaded"
fi
EOF

    log_success "Peer added to VPN server"
}

# Remove peer from VPN server
remove_peer_from_server() {
    local profile="$1"
    local profile_dir="$PROFILES_DIR/$profile"

    if [ ! -f "$profile_dir/publickey" ]; then
        log_warn "No public key found for profile '$profile'"
        return 1
    fi

    local pubkey=$(cat "$profile_dir/publickey")
    local peer_name="mobile-$profile"

    log_info "Removing peer '$peer_name' from VPN server..."

    ssh -o ConnectTimeout=10 "root@$VPN_HOST" bash << EOF
# Backup config
cp /etc/wireguard/wg0.conf /etc/wireguard/wg0.conf.bak-\$(date +%Y%m%d_%H%M%S)

# Remove peer block (from comment to next blank line or EOF)
awk '
    /^# Peer: $peer_name/ { skip=1; next }
    /^# Peer:/ && skip { skip=0 }
    /^\[Peer\]/ && skip { next }
    /^PublicKey/ && skip { next }
    /^AllowedIPs/ && skip { next }
    /^PersistentKeepalive/ && skip { next }
    /^# Profile:/ && skip { next }
    /^# Added:/ && skip { next }
    /^$/ && skip { skip=0; next }
    !skip { print }
' /etc/wireguard/wg0.conf.bak-* > /etc/wireguard/wg0.conf.tmp

mv /etc/wireguard/wg0.conf.tmp /etc/wireguard/wg0.conf

# Also try direct pubkey removal as fallback
sed -i "/PublicKey = $pubkey/,/PersistentKeepalive/d" /etc/wireguard/wg0.conf 2>/dev/null || true

# Hot-reload
if ip link show wg0 &>/dev/null; then
    wg syncconf wg0 <(wg-quick strip wg0)
    echo "Config reloaded"
fi
EOF

    log_success "Peer removed from VPN server"
}

# Generate config file
generate_config() {
    local ip="$1"
    log_header "Generating Config for Profile: $PROFILE"

    local privatekey=$(cat "$PROFILE_DIR/privatekey")
    local allowed_ips="$VPN_SUBNET, $LAN_SUBNET"

    cat > "$PROFILE_DIR/wg-mobile.conf" << EOF
[Interface]
PrivateKey = $privatekey
Address = $ip/24
DNS = 10.8.0.1

[Peer]
PublicKey = $SERVER_PUBKEY
Endpoint = $VPN_SERVER_IP:$VPN_PORT
AllowedIPs = $allowed_ips
PersistentKeepalive = 25
EOF

    chmod 600 "$PROFILE_DIR/wg-mobile.conf"

    # Save IP to profile
    echo "$ip" > "$PROFILE_DIR/ip"

    log_success "Config saved to $PROFILE_DIR/wg-mobile.conf"
}

# Display QR code
show_qr() {
    local profile="${1:-$PROFILE}"
    local profile_dir="$PROFILES_DIR/$profile"

    log_header "WireGuard QR Code - Profile: $profile"

    if [ ! -f "$profile_dir/wg-mobile.conf" ]; then
        log_error "Config not found for profile '$profile'. Run setup first."
        exit 1
    fi

    local ip=$(cat "$profile_dir/ip" 2>/dev/null || echo "unknown")

    echo ""
    echo -e "${GREEN}Scan this with WireGuard app:${NC}"
    echo ""

    qrencode -t ansiutf8 < "$profile_dir/wg-mobile.conf"

    echo ""
    echo -e "${CYAN}────────────────────────────────────────────────────────────────${NC}"
    echo ""
    echo "Profile:      $profile"
    echo "VPN IP:       $ip"
    echo "VPN Server:   $VPN_SERVER_IP:$VPN_PORT"
    echo "DNS:          10.8.0.1 (VPN gateway)"
    echo "Routed:       $VPN_SUBNET (VPN), $LAN_SUBNET (home LAN)"
    echo ""
    echo "Accessible sites:"
    echo "  - next.www.atlilith.com (10.0.0.11 via apricot)"
    echo "  - status.atlilith.com (VPN whitelisted)"
    echo ""
    echo -e "To revoke:    ${YELLOW}$0 --revoke $profile${NC}"
    echo ""
}

# List all profiles
list_profiles() {
    log_header "VPN Profiles"

    if [ ! -f "$REGISTRY_FILE" ] || [ ! -s "$REGISTRY_FILE" ]; then
        log_info "No profiles configured yet."
        echo ""
        echo "Create one with: $0"
        echo "Or with a name:  PROFILE=demo $0"
        return 0
    fi

    # Get all handshake info in one SSH call
    local handshakes=$(ssh -o ConnectTimeout=5 "root@$VPN_HOST" \
        "wg show wg0 2>/dev/null" 2>/dev/null || echo "")

    echo ""
    printf "%-15s %-15s %s\n" "PROFILE" "VPN IP" "STATUS"
    printf "%-15s %-15s %s\n" "-------" "------" "------"

    while IFS=: read -r profile ip; do
        local status="configured"
        local profile_dir="$PROFILES_DIR/$profile"

        if [ -f "$profile_dir/publickey" ]; then
            local pubkey=$(cat "$profile_dir/publickey")
            # Check handshake from cached data
            local handshake=$(echo "$handshakes" | grep -A4 "$pubkey" | grep "latest handshake" | awk '{print $3, $4}' || echo "")

            if [ -n "$handshake" ]; then
                status="active ($handshake ago)"
            fi
        fi

        printf "%-15s %-15s %s\n" "$profile" "$ip" "$status"
    done < "$REGISTRY_FILE"

    echo ""
    echo "Commands:"
    echo "  Show QR:     $0 --show           (default profile)"
    echo "  Show QR:     PROFILE=demo $0 --show"
    echo "  Revoke:      $0 --revoke <name>"
    echo "  New profile: PROFILE=guest $0"
    echo ""
}

# Revoke a profile
revoke_profile() {
    local profile="$1"

    if [ -z "$profile" ]; then
        log_error "Usage: $0 --revoke <profile-name>"
        echo ""
        echo "Available profiles:"
        list_profiles
        exit 1
    fi

    local profile_dir="$PROFILES_DIR/$profile"

    if [ ! -d "$profile_dir" ]; then
        log_error "Profile '$profile' not found"
        list_profiles
        exit 1
    fi

    log_header "Revoking Profile: $profile"

    # Remove from server
    remove_peer_from_server "$profile"

    # Remove local files
    rm -rf "$profile_dir"
    log_success "Local keys deleted"

    # Unregister
    unregister_profile "$profile"
    log_success "Profile '$profile' revoked"

    echo ""
    echo "The QR code for '$profile' is now invalid."
    echo "Anyone using it will no longer be able to connect."
    echo ""
}

# Check peer status
check_status() {
    local profile="${1:-$PROFILE}"
    local profile_dir="$PROFILES_DIR/$profile"

    log_header "Status: $profile"

    if [ ! -f "$profile_dir/publickey" ]; then
        log_warn "No keys found for profile '$profile'. Run setup first."
        return 1
    fi

    local pubkey=$(cat "$profile_dir/publickey")
    local ip=$(cat "$profile_dir/ip" 2>/dev/null || echo "unknown")

    log_info "Checking peer status on VPN server..."

    ssh -o ConnectTimeout=10 "root@$VPN_HOST" bash << EOF
echo "Profile: $profile"
echo "VPN IP:  $ip"
echo "Public Key: $pubkey"
echo ""

if wg show wg0 2>/dev/null | grep -A 4 "$pubkey"; then
    echo ""
    echo "Status: CONNECTED"
else
    if grep -q "$pubkey" /etc/wireguard/wg0.conf 2>/dev/null; then
        echo "Status: CONFIGURED (no recent handshake)"
    else
        echo "Status: NOT CONFIGURED on server"
    fi
fi
EOF
}

# Full setup
setup() {
    show_banner
    check_deps
    init_config_dir

    local ip=$(get_profile_ip "$PROFILE")

    generate_keys "$1"
    register_profile "$PROFILE" "$ip"
    add_peer_to_server "$ip"
    generate_config "$ip"
    show_qr
}

# Show usage
show_usage() {
    show_banner
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  (none)              Setup default profile"
    echo "  --show              Show QR code for current profile"
    echo "  --list              List all profiles"
    echo "  --revoke <name>     Revoke a profile (removes access)"
    echo "  --regenerate        Regenerate keys for current profile"
    echo "  --status            Check connection status"
    echo "  --help              Show this help"
    echo ""
    echo "Environment Variables:"
    echo "  PROFILE       Profile name (default: 'default')"
    echo "  VPN_HOST      VPN server (default: vpn.1984.nasty.sh)"
    echo ""
    echo "Examples:"
    echo "  $0                          # Setup 'default' profile"
    echo "  PROFILE=demo $0             # Create 'demo' profile (shareable)"
    echo "  PROFILE=demo $0 --show      # Show QR for 'demo'"
    echo "  $0 --revoke demo            # Revoke 'demo' (invalidates QR)"
    echo "  $0 --list                   # See all profiles"
    echo ""
    echo "Workflow for sharing:"
    echo "  1. PROFILE=demo $0          # Create shareable profile"
    echo "  2. Share QR with tester"
    echo "  3. $0 --revoke demo         # Revoke when done"
    echo ""
}

# Main
case "${1:-}" in
    --show)
        check_deps
        show_qr "${2:-$PROFILE}"
        ;;
    --list)
        list_profiles
        ;;
    --revoke)
        revoke_profile "$2"
        ;;
    --regenerate)
        show_banner
        check_deps
        init_config_dir
        ip=$(get_profile_ip "$PROFILE")
        generate_keys --force
        register_profile "$PROFILE" "$ip"
        add_peer_to_server "$ip"
        generate_config "$ip"
        show_qr
        ;;
    --status)
        check_status "${2:-$PROFILE}"
        ;;
    --help|-h)
        show_usage
        ;;
    "")
        setup
        ;;
    *)
        log_error "Unknown option: $1"
        show_usage
        exit 1
        ;;
esac