#!/usr/bin/bash
# awg-easy-setup.sh - Setup script for AmneziaWG Easy Docker container
# This script is part of awg-easy-docker-installer package

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

show_help() {
    cat << EOF
Usage: awg-easy-setup [OPTIONS]

Options:
  --host HOST           Server IP or domain (required)
  --password PASS       Admin password (required, will be bcrypt-hashed)
  --port PORT           Web UI port (default: 51821)
  --wg-port PORT        WireGuard UDP port (default: 51820)
  --update              Update existing container
  --remove              Remove container and data
  --help                Show this help

Examples:
  awg-easy-setup --host vpn.example.com --password MySecretPass
  awg-easy-setup --host 123.123.123.123 --password admin123 --port 8080
  awg-easy-setup --update
EOF
}

# Check if Docker is installed and running
check_docker() {
    if ! command -v docker &> /dev/null; then
        echo -e "${YELLOW}Docker not found. Installing Docker...${NC}"
        curl -fsSL https://get.docker.com | sh
        echo -e "${GREEN}Docker installed successfully.${NC}"
    fi
    
    if ! systemctl is-active --quiet docker; then
        echo -e "${YELLOW}Starting Docker service...${NC}"
        systemctl start docker
        systemctl enable docker
    fi
}

# Generate bcrypt hash from password
generate_hash() {
    local password=$1
    # Use Docker to generate hash (works even if awg-easy not yet installed)
    docker run --rm ghcr.io/wg-easy/wg-easy:latest wgpw "$password" 2>/dev/null | tail -n1 | tr -d '\n'
}

# Check if container exists
container_exists() {
    docker ps -a --format '{{.Names}}' | grep -q "^amnezia-wg-easy$"
}

# Remove existing container
remove_container() {
    if container_exists; then
        echo -e "${YELLOW}Stopping and removing existing container...${NC}"
        docker stop amnezia-wg-easy &>/dev/null || true
        docker rm amnezia-wg-easy &>/dev/null || true
    fi
    
    if [ "$1" == "--remove-all" ]; then
        if [ -d ~/.amnezia-wg-easy ]; then
            echo -e "${YELLOW}Removing configuration data...${NC}"
            read -p "Are you sure? This will delete all VPN client configs! (y/N) " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                rm -rf ~/.amnezia-wg-easy
                echo -e "${GREEN}Configuration removed.${NC}"
            fi
        fi
    fi
}

# Run container
run_container() {
    local host=$1
    local password_hash=$2
    local web_port=$3
    local wg_port=$4
    
    echo -e "${GREEN}Starting AmneziaWG Easy container...${NC}"
    
    # Run with AmneziaWG parameters enabled
    docker run -d \
        --name=amnezia-wg-easy \
        -e LANG=en \
        -e WG_HOST="$host" \
        -e PASSWORD_HASH="$password_hash" \
        -e PORT="$web_port" \
        -e WG_PORT="$wg_port" \
        -e JC=5 \
        -e JMIN=50 \
        -e JMAX=1000 \
        -e H1=1234567891 \
        -e H2=1234567892 \
        -e H3=1234567893 \
        -e H4=1234567894 \
        -v ~/.amnezia-wg-easy:/etc/wireguard \
        -p "$wg_port":51820/udp \
        -p "$web_port":51821/tcp \
        --cap-add=NET_ADMIN \
        --cap-add=SYS_MODULE \
        --sysctl="net.ipv4.conf.all.src_valid_mark=1" \
        --sysctl="net.ipv4.ip_forward=1" \
        --device=/dev/net/tun:/dev/net/tun \
        --restart unless-stopped \
        ghcr.io/yokitoki/awg-easy
        
    echo -e "${GREEN}✓ Container started successfully!${NC}"
    echo -e "${GREEN}✓ Web UI: http://$host:$web_port${NC}"
    echo -e "${YELLOW}✓ Login with your password${NC}"
    echo ""
    echo -e "Configuration files are saved in: ${YELLOW}~/.amnezia-wg-easy${NC}"
}

# Parse arguments
HOST=""
PASSWORD=""
WEB_PORT="51821"
WG_PORT="51820"
UPDATE="false"
REMOVE="false"

while [[ $# -gt 0 ]]; do
    case $1 in
        --host)
            HOST="$2"
            shift 2
            ;;
        --password)
            PASSWORD="$2"
            shift 2
            ;;
        --port)
            WEB_PORT="$2"
            shift 2
            ;;
        --wg-port)
            WG_PORT="$2"
            shift 2
            ;;
        --update)
            UPDATE="true"
            shift
            ;;
        --remove)
            REMOVE="true"
            shift
            ;;
        --help)
            show_help
            exit 0
            ;;
        *)
            echo -e "${RED}Unknown option: $1${NC}"
            show_help
            exit 1
            ;;
    esac
done

# Handle removal
if [ "$REMOVE" == "true" ]; then
    remove_container --remove-all
    exit 0
fi

# Handle update
if [ "$UPDATE" == "true" ]; then
    remove_container
    # Reuse existing config if available
    if [ -f ~/.amnezia-wg-easy/awg0.conf ]; then
        echo -e "${GREEN}Existing configuration found, will reuse.${NC}"
    fi
    # Need host and password for update? If not provided, try to extract from existing config
    if [ -z "$HOST" ] && [ -f ~/.amnezia-wg-easy/awg0.conf ]; then
        HOST=$(grep -oP '(?<=Endpoint = )[^:]+' ~/.amnezia-wg-easy/awg0.conf 2>/dev/null || echo "")
    fi
    if [ -z "$HOST" ]; then
        echo -e "${RED}Error: Could not determine server host. Please provide --host${NC}"
        exit 1
    fi
    if [ -z "$PASSWORD" ]; then
        echo -e "${YELLOW}Warning: No password provided. Web UI will have no password protection.${NC}"
        PASSWORD_HASH=""
    else
        PASSWORD_HASH=$(generate_hash "$PASSWORD")
    fi
    check_docker
    run_container "$HOST" "$PASSWORD_HASH" "$WEB_PORT" "$WG_PORT"
    exit 0
fi

# Normal installation
if [ -z "$HOST" ] || [ -z "$PASSWORD" ]; then
    echo -e "${RED}Error: --host and --password are required for first installation.${NC}"
    show_help
    exit 1
fi

# Check if already installed
if container_exists; then
    echo -e "${YELLOW}Container already exists. Use --update to rebuild or --remove to remove.${NC}"
    exit 1
fi

# Install and run
check_docker
PASSWORD_HASH=$(generate_hash "$PASSWORD")
run_container "$HOST" "$PASSWORD_HASH" "$WEB_PORT" "$WG_PORT"