#!/usr/bin/env bash
# =============================================================================
# WireGuard VPS Tunnel — 1-Command Deployer
# Deploys both server (VPS) and client (home server) from any machine with
# SSH access to both. Generates all WireGuard keys locally, pre-populates
# configs, and starts containers — no manual key exchange needed.
# =============================================================================
set -euo pipefail

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

# ── Paths ───────────────────────────────────────────────────────────────────
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DEPLOY_TAR="/tmp/wg-vps-tunnel-deploy-$$.tar.gz"

# ── Defaults ────────────────────────────────────────────────────────────────
VPS_HOST=""
CLIENT_HOST=""
VPS_PASS=""
CLIENT_PASS=""
PORTS="80,443"
WG_PORT="51820"
KEEPALIVE="21"
FORCE=false
VPS_PUBLIC_IP=""

# ── Output ──────────────────────────────────────────────────────────────────
info()    { echo -e "${BLUE}[INFO]${NC}    $*"; }
success() { echo -e "${GREEN}[OK]${NC}      $*"; }
warn()    { echo -e "${YELLOW}[WARN]${NC}    $*"; }
error()   { echo -e "${RED}[ERROR]${NC}   $*"; }
header()  { echo -e "\n${BOLD}${CYAN}═══ $* ═══${NC}\n"; }

# ── Banner ──────────────────────────────────────────────────────────────────
banner() {
    echo -e "${CYAN}${BOLD}"
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║     WireGuard VPS Tunnel — 1-Command Deployer          ║"
    echo "║     Deploy server + client from any machine             ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

# ── Usage ───────────────────────────────────────────────────────────────────
usage() {
    echo -e "${BOLD}Usage:${NC} $0 --vps HOST --client HOST [OPTIONS]"
    echo ""
    echo -e "${BOLD}Required:${NC}"
    echo "  --vps HOST            VPS address (root@2.59.219.234)"
    echo "  --client HOST         Client address (root@192.168.2.44)"
    echo ""
    echo -e "${BOLD}Optional:${NC}"
    echo "  --vps-pass PASS       VPS SSH password"
    echo "  --client-pass PASS    Client SSH password"
    echo "  --ports PORT,...      Ports to forward (default: 80,443)"
    echo "  --wg-port PORT        WireGuard port (default: 51820)"
    echo "  --keepalive SEC       PersistentKeepalive seconds (default: 21)"
    echo "  --force               Stop existing containers and redeploy"
    echo "  --help, -h            Show this help"
    echo ""
    echo -e "${BOLD}Examples:${NC}"
    echo "  # Key-based SSH auth"
    echo "  $0 --vps root@2.59.219.234 --client root@192.168.2.44 --ports 80,443"
    echo ""
    echo "  # Password-based SSH auth"
    echo "  $0 --vps root@2.59.219.234 --vps-pass wumN8inGSTiFk3Jcy5 \\"
    echo "     --client root@192.168.2.44 --client-pass retrowest --ports 80,443"
    exit 0
}

# ── Parse arguments ─────────────────────────────────────────────────────────
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --vps)          VPS_HOST="$2"; shift 2 ;;
            --client)       CLIENT_HOST="$2"; shift 2 ;;
            --vps-pass)     VPS_PASS="$2"; shift 2 ;;
            --client-pass)  CLIENT_PASS="$2"; shift 2 ;;
            --ports)        PORTS="$2"; shift 2 ;;
            --wg-port)      WG_PORT="$2"; shift 2 ;;
            --keepalive)    KEEPALIVE="$2"; shift 2 ;;
            --force)        FORCE=true; shift ;;
            --help|-h)      usage ;;
            *)
                error "Unknown argument: $1"
                echo "Use --help for usage information."
                exit 1
                ;;
        esac
    done
}

# ── Validate arguments ──────────────────────────────────────────────────────
validate_args() {
    if [[ -z "$VPS_HOST" ]]; then
        error "--vps is required (e.g. root@2.59.219.234)"
        exit 1
    fi
    if [[ -z "$CLIENT_HOST" ]]; then
        error "--client is required (e.g. root@192.168.2.44)"
        exit 1
    fi
}

# ── SSH helpers ─────────────────────────────────────────────────────────────
ssh_remote() {
    local host="$1" pass="$2" cmd="$3"
    if [[ -n "$pass" ]]; then
        sshpass -p "$pass" ssh -o StrictHostKeyChecking=no -o ConnectTimeout=10 -o BatchMode=no "$host" "$cmd"
    else
        ssh -o StrictHostKeyChecking=no -o ConnectTimeout=10 "$host" "$cmd"
    fi
}

scp_remote() {
    local host="$1" pass="$2" src="$3" dst="$4"
    if [[ -n "$pass" ]]; then
        sshpass -p "$pass" scp -o StrictHostKeyChecking=no -o ConnectTimeout=10 "$src" "${host}:${dst}"
    else
        scp -o StrictHostKeyChecking=no -o ConnectTimeout=10 "$src" "${host}:${dst}"
    fi
}

# ── Check dependencies ──────────────────────────────────────────────────────
check_deps() {
    header "Checking dependencies"

    local missing=()

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

    if [[ -n "$VPS_PASS" ]] || [[ -n "$CLIENT_PASS" ]]; then
        if ! command -v sshpass &>/dev/null; then
            warn "sshpass not found. Attempting to install..."
            if command -v apt-get &>/dev/null; then
                apt-get update -qq && apt-get install -y -qq sshpass 2>/dev/null || missing+=("sshpass")
            elif command -v yum &>/dev/null; then
                yum install -y sshpass 2>/dev/null || missing+=("sshpass")
            else
                missing+=("sshpass")
            fi
        fi
    fi

    if [[ ${#missing[@]} -gt 0 ]]; then
        error "Missing dependencies: ${missing[*]}"
        exit 1
    fi

    success "All dependencies available"
}

# ── Generate WireGuard keys ─────────────────────────────────────────────────
generate_all_keys() {
    header "Generating WireGuard keys"

    if command -v wg &>/dev/null; then
        info "Using host wireguard-tools"
        SERVER_PRIVATE_KEY=$(wg genkey)
        SERVER_PUBLIC_KEY=$(echo "$SERVER_PRIVATE_KEY" | wg pubkey)
        CLIENT_PRIVATE_KEY=$(wg genkey)
        CLIENT_PUBLIC_KEY=$(echo "$CLIENT_PRIVATE_KEY" | wg pubkey)
    else
        info "wireguard-tools not found on host, using Docker fallback..."
        local keys
        keys=$(docker run --rm alpine:3.20 sh -c "
            apk add --no-cache wireguard-tools >/dev/null 2>&1
            spk=\$(wg genkey); echo \"SERVER_PRIV=\$spk\"; echo \"SERVER_PUB=\$(echo \$spk | wg pubkey)\"
            cpk=\$(wg genkey); echo \"CLIENT_PRIV=\$cpk\"; echo \"CLIENT_PUB=\$(echo \$cpk | wg pubkey)\"
        " 2>/dev/null) || {
            error "Docker key generation failed. Install wireguard-tools: apt-get install wireguard-tools"
            exit 1
        }
        SERVER_PRIVATE_KEY=$(echo "$keys" | grep '^SERVER_PRIV=' | cut -d= -f2-)
        SERVER_PUBLIC_KEY=$(echo "$keys"  | grep '^SERVER_PUB='  | cut -d= -f2-)
        CLIENT_PRIVATE_KEY=$(echo "$keys" | grep '^CLIENT_PRIV=' | cut -d= -f2-)
        CLIENT_PUBLIC_KEY=$(echo "$keys"  | grep '^CLIENT_PUB='  | cut -d= -f2-)
    fi

    success "Server key: ${SERVER_PUBLIC_KEY:0:12}..."
    success "Client key: ${CLIENT_PUBLIC_KEY:0:12}..."
}

# ── Detect VPS info ────────────────────────────────────────────────────────
detect_vps_info() {
    header "Detecting VPS network info"

    if [[ -z "$VPS_PUBLIC_IP" ]]; then
        info "Detecting VPS public IP..."
        VPS_PUBLIC_IP=$(ssh_remote "$VPS_HOST" "$VPS_PASS" \
            "curl -4 -s --connect-timeout 5 ifconfig.me 2>/dev/null || curl -4 -s --connect-timeout 5 icanhazip.com 2>/dev/null || echo ''")
        if [[ -z "$VPS_PUBLIC_IP" ]]; then
            error "Could not detect VPS public IP. Provide it via VPS_PUBLIC_IP env or check connectivity."
            exit 1
        fi
    fi
    success "VPS public IP: ${VPS_PUBLIC_IP}"

    info "Detecting VPS public interface..."
    VPS_IFACE=$(ssh_remote "$VPS_HOST" "$VPS_PASS" \
        "ip route get 8.8.8.8 2>/dev/null | awk '{print \$5; exit}'")
    if [[ -z "$VPS_IFACE" ]]; then
        warn "Could not detect VPS interface. Container will auto-detect."
        VPS_IFACE=""
    else
        success "VPS public interface: ${VPS_IFACE}"
    fi
}

# ── Create deployment archive ────────────────────────────────────────────────
create_archive() {
    header "Creating deployment archive"
    tar czf "$DEPLOY_TAR" \
        --exclude='.git' \
        --exclude='deploy.sh' \
        --exclude='config' \
        --exclude='.env' \
        -C "$SCRIPT_DIR" .
    success "Archive created: ${DEPLOY_TAR} ($(du -h "$DEPLOY_TAR" | cut -f1))"
}

# ── Deploy to remote host ───────────────────────────────────────────────────
deploy_to_host() {
    local host="$1" pass="$2" role="$3"
    local target="/opt/wireguard-vps-tunnel"

    header "Deploying to ${role}: ${host}"

    info "Copying archive..."
    scp_remote "$host" "$pass" "$DEPLOY_TAR" "/tmp/wg-deploy.tar.gz"

    info "Extracting to ${target}..."
    ssh_remote "$host" "$pass" "mkdir -p ${target} && rm -rf ${target}/* 2>/dev/null; tar xzf /tmp/wg-deploy.tar.gz -C ${target} && rm -f /tmp/wg-deploy.tar.gz"

    info "Creating config directory with keys..."
    ssh_remote "$host" "$pass" "
        mkdir -p ${target}/config
        echo '${SERVER_PRIVATE_KEY}' > ${target}/config/server_private.key
        echo '${SERVER_PUBLIC_KEY}' > ${target}/config/server_public.key
        echo '${CLIENT_PRIVATE_KEY}' > ${target}/config/client_private.key
        echo '${CLIENT_PUBLIC_KEY}' > ${target}/config/client_public.key
        chmod 600 ${target}/config/*.key
    "

    if [[ "$role" == "vps" ]]; then
        info "Creating server .env..."
        ssh_remote "$host" "$pass" "
            cat > ${target}/.env << 'DOTENV'
# WireGuard VPS Tunnel — Server (VPS)
CLIENT_PUBLIC_KEY=${CLIENT_PUBLIC_KEY}
SERVER_PUBLIC_IP=${VPS_PUBLIC_IP}
SERVER_PUBLIC_IFACE=${VPS_IFACE}
FORWARD_PORTS=${PORTS}
SERVER_WG_PORT=${WG_PORT}
DOTENV
            chmod 600 ${target}/.env
        "
    else
        info "Creating client .env..."
        ssh_remote "$host" "$pass" "
            cat > ${target}/.env << 'DOTENV'
# WireGuard VPS Tunnel — Client (home server)
VPS_PUBLIC_IP=${VPS_PUBLIC_IP}
SERVER_PUBLIC_KEY=${SERVER_PUBLIC_KEY}
PERSISTENT_KEEPALIVE=${KEEPALIVE}
VPS_WG_PORT=${WG_PORT}
DOTENV
            chmod 600 ${target}/.env
        "
    fi

    success "Deployed to ${role}"
}

# ── Stop host-level WireGuard ───────────────────────────────────────────────
stop_host_wg() {
    local host="$1" pass="$2"

    ssh_remote "$host" "$pass" "
        if systemctl is-active --quiet wg-quick@wg0 2>/dev/null; then
            echo '[WARN] Stopping host-level wg-quick@wg0...'
            systemctl stop wg-quick@wg0 2>/dev/null || true
            systemctl disable wg-quick@wg0 2>/dev/null || true
            ip link del wg0 2>/dev/null || true
        fi
        if systemctl is-active --quiet wg-watchdog 2>/dev/null; then
            systemctl stop wg-watchdog 2>/dev/null || true
            systemctl disable wg-watchdog 2>/dev/null || true
        fi
    " 2>/dev/null || true
}

# ── Enable IP forwarding on VPS ─────────────────────────────────────────────
enable_ip_forwarding() {
    header "Enabling IP forwarding on VPS"

    ssh_remote "$VPS_HOST" "$VPS_PASS" "
        echo 1 > /proc/sys/net/ipv4/ip_forward 2>/dev/null || true
        echo 1 > /proc/sys/net/ipv6/conf/all/forwarding 2>/dev/null || true

        if ! grep -q 'net.ipv4.ip_forward' /etc/sysctl.conf 2>/dev/null; then
            echo 'net.ipv4.ip_forward = 1' >> /etc/sysctl.conf
        else
            sed -i 's/.*net.ipv4.ip_forward.*/net.ipv4.ip_forward = 1/' /etc/sysctl.conf
        fi

        if ! grep -q 'net.ipv6.conf.all.forwarding' /etc/sysctl.conf 2>/dev/null; then
            echo 'net.ipv6.conf.all.forwarding = 1' >> /etc/sysctl.conf
        else
            sed -i 's/.*net.ipv6.conf.all.forwarding.*/net.ipv6.conf.all.forwarding = 1/' /etc/sysctl.conf
        fi

        sysctl -p 2>/dev/null || true
    "

    success "IP forwarding enabled on VPS"
}

# ── Start container ─────────────────────────────────────────────────────────
start_container() {
    local host="$1" pass="$2" role="$3"
    local target="/opt/wireguard-vps-tunnel"
    local compose_file

    if [[ "$role" == "vps" ]]; then
        compose_file="docker-compose.server.yml"
    else
        compose_file="docker-compose.client.yml"
    fi

    header "Starting ${role} container: ${host}"

    # Remove stale wg0 interface (leftover from previous container with host network)
    ssh_remote "$host" "$pass" "ip link del wg0 2>/dev/null || true"

    ssh_remote "$host" "$pass" "
        cd ${target}
        docker compose -f ${compose_file} up -d --remove-orphans 2>&1
    "

    success "${role} container started"
}

# ── Force stop existing containers ──────────────────────────────────────────
force_stop() {
    header "Force-stopping existing containers"

    info "Stopping VPS container and cleaning up..."
    ssh_remote "$VPS_HOST" "$VPS_PASS" "
        cd /opt/wireguard-vps-tunnel 2>/dev/null || exit 0
        docker compose -f docker-compose.server.yml down --remove-orphans 2>/dev/null || true
        ip link del wg0 2>/dev/null || true
    " || true

    info "Stopping client container and cleaning up..."
    ssh_remote "$CLIENT_HOST" "$CLIENT_PASS" "
        cd /opt/wireguard-vps-tunnel 2>/dev/null || exit 0
        docker compose -f docker-compose.client.yml down --remove-orphans 2>/dev/null || true
        ip link del wg0 2>/dev/null || true
    " || true

    success "Existing containers stopped and interfaces cleaned"
}

# ── Wait for healthy containers ─────────────────────────────────────────────
wait_for_healthy() {
    header "Waiting for containers to be healthy (max 60s)"

    local vps_healthy=false
    local client_healthy=false

    for i in $(seq 1 12); do
        if [[ "$vps_healthy" != true ]]; then
            local vps_status
            vps_status=$(ssh_remote "$VPS_HOST" "$VPS_PASS" \
                "docker inspect wireguard-server --format='{{.State.Health.Status}}' 2>/dev/null || echo 'missing'")
            if [[ "$vps_status" == "healthy" ]]; then
                success "VPS container healthy"
                vps_healthy=true
            elif [[ "$vps_status" == "missing" ]]; then
                warn "VPS container not found yet..."
            else
                info "VPS container status: ${vps_status} (${i}/12)"
            fi
        fi

        if [[ "$client_healthy" != true ]]; then
            local client_status
            client_status=$(ssh_remote "$CLIENT_HOST" "$CLIENT_PASS" \
                "docker inspect wireguard-client --format='{{.State.Health.Status}}' 2>/dev/null || echo 'missing'")
            if [[ "$client_status" == "healthy" ]]; then
                success "Client container healthy"
                client_healthy=true
            elif [[ "$client_status" == "missing" ]]; then
                warn "Client container not found yet..."
            else
                info "Client container status: ${client_status} (${i}/12)"
            fi
        fi

        if [[ "$vps_healthy" == true ]] && [[ "$client_healthy" == true ]]; then
            break
        fi

        sleep 5
    done

    if [[ "$vps_healthy" != true ]]; then
        warn "VPS container may not be fully healthy. Check logs: docker compose -f docker-compose.server.yml logs"
    fi
    if [[ "$client_healthy" != true ]]; then
        warn "Client container may not be fully healthy. Check logs: docker compose -f docker-compose.client.yml logs"
    fi
}

# ── Verify tunnel ───────────────────────────────────────────────────────────
verify_tunnel() {
    header "Verifying tunnel"

    info "Ping from VPS (10.0.0.1) to client (10.0.0.2)..."
    if ssh_remote "$VPS_HOST" "$VPS_PASS" "ping -c 3 -W 2 10.0.0.2" 2>/dev/null; then
        success "VPS → client ping OK"
    else
        warn "VPS → client ping failed. Tunnel may need more time to establish."
    fi

    info "Ping from client (10.0.0.2) to VPS (10.0.0.1)..."
    if ssh_remote "$CLIENT_HOST" "$CLIENT_PASS" "ping -c 3 -W 2 10.0.0.1" 2>/dev/null; then
        success "Client → VPS ping OK"
    else
        warn "Client → VPS ping failed. Tunnel may need more time to establish."
    fi

    info "Checking DNAT on VPS (curl localhost)..."
    local http_code
    http_code=$(ssh_remote "$VPS_HOST" "$VPS_PASS" \
        "curl -s -o /dev/null -w '%{http_code}' --connect-timeout 3 http://localhost 2>/dev/null || echo '000'")
    if [[ "$http_code" != "000" ]]; then
        success "DNAT responding (HTTP ${http_code})"
    else
        warn "DNAT check: no response (expected if no web server on client yet)"
    fi
}

# ── Print summary ───────────────────────────────────────────────────────────
print_summary() {
    header "Deployment Complete"

    echo -e "${BOLD}${GREEN}WireGuard VPS Tunnel deployed successfully!${NC}"
    echo ""
    echo -e "${BOLD}VPS (Server):${NC}         ${VPS_HOST}"
    echo -e "${BOLD}Client:${NC}              ${CLIENT_HOST}"
    echo -e "${BOLD}VPS Public IP:${NC}       ${VPS_PUBLIC_IP}"
    echo -e "${BOLD}WireGuard Network:${NC}   10.0.0.0/24"
    echo -e "${BOLD}  Server WG IP:${NC}      10.0.0.1"
    echo -e "${BOLD}  Client WG IP:${NC}      10.0.0.2"
    echo -e "${BOLD}WG Port:${NC}             ${WG_PORT}/udp"
    echo -e "${BOLD}Forwarded Ports:${NC}     ${PORTS} → 10.0.0.2"
    echo -e "${BOLD}Keepalive:${NC}           ${KEEPALIVE}s"
    echo ""
    echo -e "${BOLD}Keys:${NC}"
    echo -e "  Server Public:  ${SERVER_PUBLIC_KEY}"
    echo -e "  Client Public:  ${CLIENT_PUBLIC_KEY}"
    echo ""
    echo -e "${BOLD}Useful commands:${NC}"
    echo "  # Check tunnel status on VPS"
    echo "  ssh ${VPS_HOST} docker exec wireguard-server wg show"
    echo ""
    echo "  # Check tunnel status on client"
    echo "  ssh ${CLIENT_HOST} docker exec wireguard-client wg show"
    echo ""
    echo "  # View logs"
    echo "  ssh ${VPS_HOST} 'docker compose -f /opt/wireguard-vps-tunnel/docker-compose.server.yml logs -f'"
    echo "  ssh ${CLIENT_HOST} 'docker compose -f /opt/wireguard-vps-tunnel/docker-compose.client.yml logs -f'"
    echo ""
    echo -e "${BOLD}Config locations:${NC}"
    echo "  VPS:  /opt/wireguard-vps-tunnel/"
    echo "  Client: /opt/wireguard-vps-tunnel/"
}

# ── Cleanup ─────────────────────────────────────────────────────────────────
cleanup() {
    rm -f "$DEPLOY_TAR"
}
trap cleanup EXIT

# ── Main ────────────────────────────────────────────────────────────────────
main() {
    banner
    parse_args "$@"
    validate_args
    check_deps

    if [[ "$FORCE" == true ]]; then
        force_stop
    fi

    generate_all_keys
    detect_vps_info
    create_archive

    deploy_to_host "$VPS_HOST" "$VPS_PASS" "vps"
    deploy_to_host "$CLIENT_HOST" "$CLIENT_PASS" "client"

    stop_host_wg "$VPS_HOST" "$VPS_PASS"
    stop_host_wg "$CLIENT_HOST" "$CLIENT_PASS"
    enable_ip_forwarding

    start_container "$VPS_HOST" "$VPS_PASS" "vps"
    start_container "$CLIENT_HOST" "$CLIENT_PASS" "client"

    wait_for_healthy
    verify_tunnel
    print_summary
}

main "$@"