#!/bin/bash
# K3S 3-Node Cluster Deployment Script
# Deploys a K3S cluster with 1 server and 2 agent nodes

set -euo pipefail

# Temp file cleanup
TEMP_FILES=()
cleanup_temp_files() {
    for f in "${TEMP_FILES[@]}"; do
        rm -f "$f" 2>/dev/null || true
    done
}
trap cleanup_temp_files EXIT

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

# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(dirname "$SCRIPT_DIR")"
ANSIBLE_DIR="$ROOT_DIR/ansible"

# Source deploy-vm.sh functions
# shellcheck source=./deploy-vm.sh
source "$SCRIPT_DIR/deploy-vm.sh"

# Logging (all output to stderr to avoid interfering with function return values)
log_info() {
    echo -e "${GREEN}[INFO]${NC} $*" >&2
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $*" >&2
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $*" >&2
}

log_step() {
    echo -e "${BLUE}[STEP]${NC} $*" >&2
}

log_cluster() {
    echo -e "${CYAN}[K3S]${NC} $*" >&2
}

#=============================================================================
# VMID Management
#=============================================================================

# Get all existing VMIDs in the cluster
get_existing_vmids() {
    local vmids=()

    # Try cluster-wide lookup first
    if pvesh get /cluster/resources --type vm --output-format json 2>/dev/null; then
        mapfile -t vmids < <(pvesh get /cluster/resources --type vm --output-format json 2>/dev/null | \
            grep -oP '"vmid":\K[0-9]+' | sort -n)
    else
        # Fall back to local node
        mapfile -t vmids < <(qm list | awk 'NR>1 {print $1}' | sort -n)
    fi

    printf '%s\n' "${vmids[@]}"
}

# Check if VMID is available
is_vmid_available() {
    local vmid=$1
    local existing_vmids

    existing_vmids=$(get_existing_vmids)

    if echo "$existing_vmids" | grep -q "^${vmid}$"; then
        return 1  # VMID exists
    else
        return 0  # VMID available
    fi
}

# Calculate 3 consecutive available VMIDs
calculate_cluster_vmids() {
    local start_vmid=${1:-200}  # Default starting point
    local max_vmid=9999

    log_step "Finding 3 consecutive available VMIDs..."

    local existing_vmids
    existing_vmids=$(get_existing_vmids)

    # Find 3 consecutive available VMIDs
    for ((vmid = start_vmid; vmid <= max_vmid - 2; vmid++)); do
        local vmid1=$vmid
        local vmid2=$((vmid + 1))
        local vmid3=$((vmid + 2))

        # Check if all 3 are available
        if ! echo "$existing_vmids" | grep -q "^${vmid1}$" && \
           ! echo "$existing_vmids" | grep -q "^${vmid2}$" && \
           ! echo "$existing_vmids" | grep -q "^${vmid3}$"; then
            echo "$vmid1 $vmid2 $vmid3"
            return 0
        fi
    done

    log_error "Could not find 3 consecutive available VMIDs"
    return 1
}

#=============================================================================
# VM Deployment
#=============================================================================

# Deploy a single cluster node VM
deploy_cluster_node() {
    local template_id=$1
    local vm_id=$2
    local vm_name=$3
    local cores=$4
    local memory=$5
    local disk_size=$6
    local clone_type=${7:-linked}

    log_step "Deploying $vm_name (VMID $vm_id)..."

    # Clone VM
    if ! clone_vm "$template_id" "$vm_id" "$vm_name" "$clone_type"; then
        log_error "Failed to clone VM $vm_name"
        return 1
    fi

    # Resize disk
    if ! resize_disk "$vm_id" "$disk_size"; then
        log_error "Failed to resize disk for VM $vm_name"
        return 1
    fi

    # Configure resources
    configure_vm "$vm_id" "$cores" "$memory"

    # Configure network (DHCP)
    configure_network "$vm_id" "dhcp"

    log_info "✓ VM $vm_name configured"
}

# Deploy all 3 cluster VMs
deploy_cluster_vms() {
    local template_id=$1
    local base_name=$2
    local server_vmid=$3
    local agent1_vmid=$4
    local agent2_vmid=$5
    local cores=$6
    local memory=$7
    local disk_size=$8
    local clone_type=$9

    echo
    echo "========================================"
    echo "K3S Cluster VM Deployment"
    echo "========================================"
    echo "Template ID:     $template_id"
    echo "Base Name:       $base_name"
    echo "Clone Type:      $clone_type"
    echo
    echo "Server Node:     ${base_name}-server   (VMID $server_vmid)"
    echo "Agent Node 1:    ${base_name}-agent1   (VMID $agent1_vmid)"
    echo "Agent Node 2:    ${base_name}-agent2   (VMID $agent2_vmid)"
    echo
    echo "Resources/Node:  ${cores} cores, ${memory}MB RAM, ${disk_size} disk"
    echo "========================================"
    echo

    # Deploy server node
    log_step "[1/3] Deploying K3S server node..."
    if ! deploy_cluster_node "$template_id" "$server_vmid" "${base_name}-server" \
        "$cores" "$memory" "$disk_size" "$clone_type"; then
        return 1
    fi

    # Deploy agent node 1
    log_step "[2/3] Deploying K3S agent node 1..."
    if ! deploy_cluster_node "$template_id" "$agent1_vmid" "${base_name}-agent1" \
        "$cores" "$memory" "$disk_size" "$clone_type"; then
        return 1
    fi

    # Deploy agent node 3
    log_step "[3/3] Deploying K3S agent node 2..."
    if ! deploy_cluster_node "$template_id" "$agent2_vmid" "${base_name}-agent2" \
        "$cores" "$memory" "$disk_size" "$clone_type"; then
        return 1
    fi

    log_info "✓ All 3 VMs deployed successfully"
}

# Start all cluster VMs and get their IPs
start_cluster_vms() {
    local server_vmid=$1
    local agent1_vmid=$2
    local agent2_vmid=$3

    echo
    log_step "Starting all cluster VMs..."

    # Start all VMs
    log_info "Starting server node (VMID $server_vmid)..."
    qm start "$server_vmid"

    log_info "Starting agent node 1 (VMID $agent1_vmid)..."
    qm start "$agent1_vmid"

    log_info "Starting agent node 2 (VMID $agent2_vmid)..."
    qm start "$agent2_vmid"

    echo
    log_step "Waiting for all VMs to boot..."

    # Wait for QEMU agents
    log_info "[1/3] Waiting for server node QEMU agent..."
    if ! start_and_wait_vm "$server_vmid" 180; then
        log_error "Server node failed to start"
        return 1
    fi

    log_info "[2/3] Waiting for agent node 1 QEMU agent..."
    if ! start_and_wait_vm "$agent1_vmid" 180; then
        log_error "Agent node 1 failed to start"
        return 1
    fi

    log_info "[3/3] Waiting for agent node 2 QEMU agent..."
    if ! start_and_wait_vm "$agent2_vmid" 180; then
        log_error "Agent node 2 failed to start"
        return 1
    fi

    echo
    log_step "Getting IP addresses for all nodes..."

    # Get IPs
    local server_ip
    if ! server_ip=$(get_vm_ip "$server_vmid"); then
        log_error "Failed to get server IP"
        return 1
    fi
    log_info "Server IP: $server_ip"

    local agent1_ip
    if ! agent1_ip=$(get_vm_ip "$agent1_vmid"); then
        log_error "Failed to get agent1 IP"
        return 1
    fi
    log_info "Agent1 IP: $agent1_ip"

    local agent2_ip
    if ! agent2_ip=$(get_vm_ip "$agent2_vmid"); then
        log_error "Failed to get agent2 IP"
        return 1
    fi
    log_info "Agent2 IP: $agent2_ip"

    # Return IPs as space-separated string
    echo "$server_ip $agent1_ip $agent2_ip"
}

#=============================================================================
# K3S Configuration
#=============================================================================

# Install K3S server node
install_k3s_server() {
    local server_ip=$1
    local k3s_version=${2:-stable}

    log_cluster "Installing K3S server on $server_ip..."

    # Source defaults.conf for cloud-init credentials
    local config_file="$ROOT_DIR/config/defaults.conf"
    if [[ -f "$config_file" ]]; then
        if grep -Pv '^\s*$|^\s*#|^[A-Za-z_][A-Za-z0-9_]*=' "$config_file" | grep -q .; then
            log_error "defaults.conf contains invalid lines."
            return 1
        fi
        # shellcheck source=/dev/null
        source "$config_file"
    else
        log_error "Configuration file not found: $config_file"
        return 1
    fi

    # Export password for Ansible
    export ANSIBLE_PASSWORD="${DEFAULT_CLOUDINIT_PASSWORD}"

    # Set timestamped Ansible log file
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local ansible_log="$ROOT_DIR/logs/k3s-cluster-server-${timestamp}.log"
    export ANSIBLE_LOG_PATH="$ansible_log"

    log_info "Ansible logs: $ansible_log"

    cd "$ANSIBLE_DIR" || return 1

    # Wait for SSH
    if ! wait_for_ssh "$server_ip" 60; then
        log_warn "SSH not ready, but continuing anyway"
    fi

    # Set environment variables for Ansible
    export ANSIBLE_HOST_KEY_CHECKING=False

    # Verify we have an IP
    if [[ -z "$server_ip" ]]; then
        log_error "Server IP is empty!"
        return 1
    fi

    # Create temporary inventory file
    local temp_inventory=$(mktemp)
    chmod 600 "$temp_inventory"
    TEMP_FILES+=("$temp_inventory")
    cat > "$temp_inventory" <<EOF
[k3s_server]
${server_ip} ansible_user=${DEFAULT_CLOUDINIT_USER} ansible_password=${DEFAULT_CLOUDINIT_PASSWORD} ansible_become_password=${DEFAULT_CLOUDINIT_PASSWORD} ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
EOF

    log_info "Created temporary inventory for server"

    # Run K3S server playbook
    if ansible-playbook -i "$temp_inventory" \
        --extra-vars "k3s_version=${k3s_version}" \
        "playbooks/k3s-cluster-server.yml"; then
        rm -f "$temp_inventory"
        log_cluster "✓ K3S server installed successfully"
    else
        log_error "K3S server installation failed"
        return 1
    fi

    cd "$ROOT_DIR" || return 1
}

# Get K3S token from server
get_k3s_token() {
    local server_ip=$1

    log_cluster "Retrieving K3S token from server..."

    # Source defaults.conf for cloud-init credentials
    local config_file="$ROOT_DIR/config/defaults.conf"
    if [[ -f "$config_file" ]]; then
        if grep -Pv '^\s*$|^\s*#|^[A-Za-z_][A-Za-z0-9_]*=' "$config_file" | grep -q .; then
            log_error "defaults.conf contains invalid lines."
            return 1
        fi
        # shellcheck source=/dev/null
        source "$config_file"
    fi

    # Get token from server using Ansible
    export ANSIBLE_HOST_KEY_CHECKING=False

    # Create temporary inventory file with SSH args to bypass host key checking
    local temp_inventory=$(mktemp)
    chmod 600 "$temp_inventory"
    TEMP_FILES+=("$temp_inventory")
    echo "${server_ip} ansible_user=${DEFAULT_CLOUDINIT_USER} ansible_password=${DEFAULT_CLOUDINIT_PASSWORD} ansible_become_password=${DEFAULT_CLOUDINIT_PASSWORD} ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'" > "$temp_inventory"

    local token
    token=$(ansible -i "$temp_inventory" all \
        -m shell \
        -a "sudo cat /var/lib/rancher/k3s/server/node-token" \
        2>/dev/null | grep -v "^${server_ip}" | grep -v "CHANGED" | xargs)

    rm -f "$temp_inventory"

    if [[ -z "$token" ]]; then
        log_error "Failed to retrieve K3S token"
        return 1
    fi

    log_cluster "✓ K3S token retrieved"
    echo "$token"
}

# Install K3S agent node
install_k3s_agent() {
    local agent_ip=$1
    local server_ip=$2
    local k3s_token=$3
    local k3s_version=${4:-stable}
    local node_name=$5

    log_cluster "Installing K3S agent on $agent_ip ($node_name)..."

    # Source defaults.conf for cloud-init credentials
    local config_file="$ROOT_DIR/config/defaults.conf"
    if [[ -f "$config_file" ]]; then
        if grep -Pv '^\s*$|^\s*#|^[A-Za-z_][A-Za-z0-9_]*=' "$config_file" | grep -q .; then
            log_error "defaults.conf contains invalid lines."
            return 1
        fi
        # shellcheck source=/dev/null
        source "$config_file"
    else
        log_error "Configuration file not found: $config_file"
        return 1
    fi

    # Export password for Ansible
    export ANSIBLE_PASSWORD="${DEFAULT_CLOUDINIT_PASSWORD}"

    # Set timestamped Ansible log file
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local ansible_log="$ROOT_DIR/logs/k3s-cluster-agent-${node_name}-${timestamp}.log"
    export ANSIBLE_LOG_PATH="$ansible_log"

    log_info "Ansible logs: $ansible_log"

    cd "$ANSIBLE_DIR" || return 1

    # Wait for SSH
    if ! wait_for_ssh "$agent_ip" 60; then
        log_warn "SSH not ready, but continuing anyway"
    fi

    # Set environment variables for Ansible
    export ANSIBLE_HOST_KEY_CHECKING=False

    # Create temporary inventory file
    local temp_inventory=$(mktemp)
    chmod 600 "$temp_inventory"
    TEMP_FILES+=("$temp_inventory")
    cat > "$temp_inventory" <<EOF
[k3s_agent]
${agent_ip} ansible_user=${DEFAULT_CLOUDINIT_USER} ansible_password=${DEFAULT_CLOUDINIT_PASSWORD} ansible_become_password=${DEFAULT_CLOUDINIT_PASSWORD} ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
EOF

    log_info "Created temporary inventory for $node_name: $temp_inventory"

    # Run K3S agent playbook
    if ansible-playbook -i "$temp_inventory" \
        --extra-vars "k3s_server_ip=${server_ip}" \
        --extra-vars "k3s_token=${k3s_token}" \
        --extra-vars "k3s_version=${k3s_version}" \
        "playbooks/k3s-cluster-agent.yml"; then
        rm -f "$temp_inventory"
        log_cluster "✓ K3S agent installed successfully on $node_name"
    else
        log_error "K3S agent installation failed on $node_name"
        return 1
    fi

    cd "$ROOT_DIR" || return 1
}

# Verify cluster health
verify_k3s_cluster() {
    local server_ip=$1

    log_cluster "Verifying K3S cluster health..."

    # Source defaults.conf for cloud-init credentials
    local config_file="$ROOT_DIR/config/defaults.conf"
    if [[ -f "$config_file" ]]; then
        if grep -Pv '^\s*$|^\s*#|^[A-Za-z_][A-Za-z0-9_]*=' "$config_file" | grep -q .; then
            log_error "defaults.conf contains invalid lines."
            return 1
        fi
        # shellcheck source=/dev/null
        source "$config_file"
    fi

    # Get node status from server
    export ANSIBLE_HOST_KEY_CHECKING=False

    # Create temporary inventory file with SSH args to bypass host key checking
    local temp_inventory=$(mktemp)
    chmod 600 "$temp_inventory"
    TEMP_FILES+=("$temp_inventory")
    echo "${server_ip} ansible_user=${DEFAULT_CLOUDINIT_USER} ansible_password=${DEFAULT_CLOUDINIT_PASSWORD} ansible_become_password=${DEFAULT_CLOUDINIT_PASSWORD} ansible_ssh_common_args='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'" > "$temp_inventory"

    local node_status
    node_status=$(ansible -i "$temp_inventory" all \
        -m shell \
        -a "sudo kubectl get nodes --kubeconfig=/etc/rancher/k3s/k3s.yaml" \
        2>/dev/null | grep -A 10 "NAME")

    rm -f "$temp_inventory"

    if [[ -z "$node_status" ]]; then
        log_error "Failed to get cluster status"
        return 1
    fi

    echo
    log_cluster "Cluster Status:"
    echo "$node_status"
    echo

    # Check if all nodes are Ready
    local ready_count
    ready_count=$(echo "$node_status" | grep -c "Ready" || true)

    if [[ $ready_count -ge 3 ]]; then
        log_cluster "✓ All 3 nodes are Ready"
        return 0
    else
        log_warn "Only $ready_count node(s) are Ready (expected 3)"
        return 1
    fi
}

#=============================================================================
# Main Deployment Function
#=============================================================================

deploy_k3s_cluster() {
    local template_id=$1
    local base_name=$2
    local cores=${3:-2}
    local memory=${4:-4096}
    local disk_size=${5:-20G}
    local clone_type=${6:-linked}
    local k3s_version=${7:-stable}
    local start_vmid=${8:-200}

    echo
    echo "=========================================="
    echo "K3S 3-Node Cluster Deployment"
    echo "=========================================="

    # Calculate VMIDs
    local vmids
    if ! vmids=$(calculate_cluster_vmids "$start_vmid"); then
        log_error "Failed to calculate VMIDs"
        return 1
    fi

    read -r server_vmid agent1_vmid agent2_vmid <<< "$vmids"

    log_info "Assigned VMIDs: Server=$server_vmid, Agent1=$agent1_vmid, Agent2=$agent2_vmid"

    # Deploy VMs
    if ! deploy_cluster_vms "$template_id" "$base_name" \
        "$server_vmid" "$agent1_vmid" "$agent2_vmid" \
        "$cores" "$memory" "$disk_size" "$clone_type"; then
        log_error "Failed to deploy cluster VMs"
        return 1
    fi

    # Start VMs and get IPs
    local ips
    if ! ips=$(start_cluster_vms "$server_vmid" "$agent1_vmid" "$agent2_vmid"); then
        log_error "Failed to start cluster VMs"
        return 1
    fi

    # Get only the last line (which contains the IPs)
    ips=$(echo "$ips" | tail -1)

    # Debug: Show what we captured
    log_info "IP addresses from last line: [$ips]"

    read -r server_ip agent1_ip agent2_ip <<< "$ips"

    # Debug: Show what we parsed
    log_info "Parsed IPs: server=$server_ip, agent1=$agent1_ip, agent2=$agent2_ip"

    echo
    log_info "All VMs are running with IPs assigned"

    # Configure K3S cluster
    echo
    echo "========================================"
    echo "K3S Cluster Configuration"
    echo "========================================"

    # Step 1: Install K3S server
    if ! install_k3s_server "$server_ip" "$k3s_version"; then
        log_error "Failed to install K3S server"
        return 1
    fi

    # Step 2: Get K3S token
    local k3s_token
    if ! k3s_token=$(get_k3s_token "$server_ip"); then
        log_error "Failed to get K3S token"
        return 1
    fi

    # Step 3: Install K3S agent 1
    if ! install_k3s_agent "$agent1_ip" "$server_ip" "$k3s_token" "$k3s_version" "agent1"; then
        log_error "Failed to install K3S agent 1"
        return 1
    fi

    # Step 4: Install K3S agent 2
    if ! install_k3s_agent "$agent2_ip" "$server_ip" "$k3s_token" "$k3s_version" "agent2"; then
        log_error "Failed to install K3S agent 2"
        return 1
    fi

    # Wait a bit for nodes to fully join
    log_cluster "Waiting for nodes to fully join cluster..."
    sleep 10

    # Step 5: Verify cluster
    if ! verify_k3s_cluster "$server_ip"; then
        log_warn "Cluster verification reported issues, but cluster may still be functional"
    fi

    # Final summary
    echo
    echo "=========================================="
    log_info "K3S Cluster Deployment Complete!"
    echo "=========================================="
    echo
    echo "Server Node:     ${base_name}-server  (VMID $server_vmid) - $server_ip"
    echo "Agent Node 1:    ${base_name}-agent1  (VMID $agent1_vmid) - $agent1_ip"
    echo "Agent Node 2:    ${base_name}-agent2  (VMID $agent2_vmid) - $agent2_ip"
    echo
    echo "Kubeconfig:      /etc/rancher/k3s/k3s.yaml (on server)"
    echo "K3S Version:     $k3s_version"
    echo
    echo "Access cluster:"
    echo "  ssh admin@$server_ip"
    echo "  export KUBECONFIG=/etc/rancher/k3s/k3s.yaml"
    echo "  kubectl get nodes"
    echo
    echo "Remote access:"
    echo "  scp admin@${server_ip}:/etc/rancher/k3s/k3s.yaml ~/.kube/config"
    echo "  sed -i 's/127.0.0.1/${server_ip}/' ~/.kube/config"
    echo "  kubectl get nodes"
    echo "=========================================="
    echo

    return 0
}

#=============================================================================
# Script Entry Point
#=============================================================================

# Usage
usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Deploy a 3-node K3S cluster (1 server + 2 agents) from a Proxmox template.

Required Arguments:
  -t TEMPLATE_ID    Proxmox template ID to clone from
  -n BASE_NAME      Base name for cluster nodes (e.g., "k3s-lab")

Optional Arguments:
  -c CORES          CPU cores per node (default: 2)
  -m MEMORY         Memory per node in MB (default: 4096)
  -d DISK_SIZE      Disk size per node (default: 20G)
  -k CLONE_TYPE     Clone type: 'linked' or 'full' (default: linked)
  -v K3S_VERSION    K3S version channel: 'stable', 'latest', or specific version (default: stable)
  -s START_VMID     Starting VMID to search from (default: 200)
  -h                Show this help message

Examples:
  # Deploy with defaults
  $0 -t 9001 -n k3s-lab

  # Deploy with custom resources
  $0 -t 9001 -n k3s-prod -c 4 -m 8192 -d 40G

  # Deploy with specific K3S version
  $0 -t 9001 -n k3s-test -v v1.28.5+k3s1

EOF
}

# Main
main() {
    local template_id=""
    local base_name=""
    local cores=2
    local memory=4096
    local disk_size="20G"
    local clone_type="linked"
    local k3s_version="stable"
    local start_vmid=200

    # Parse arguments
    while getopts "t:n:c:m:d:k:v:s:h" opt; do
        case $opt in
            t) template_id="$OPTARG" ;;
            n) base_name="$OPTARG" ;;
            c) cores="$OPTARG" ;;
            m) memory="$OPTARG" ;;
            d) disk_size="$OPTARG" ;;
            k) clone_type="$OPTARG" ;;
            v) k3s_version="$OPTARG" ;;
            s) start_vmid="$OPTARG" ;;
            h) usage; exit 0 ;;
            *) usage; exit 1 ;;
        esac
    done

    # Validate required arguments
    if [[ -z "$template_id" ]] || [[ -z "$base_name" ]]; then
        log_error "Missing required arguments"
        usage
        exit 1
    fi

    # Validate cluster name
    if [[ ! "$base_name" =~ ^[a-zA-Z0-9._-]+$ ]]; then
        log_error "Invalid cluster name: $base_name (only alphanumeric, dots, hyphens, underscores allowed)"
        exit 1
    fi

    # Validate template exists
    if ! qm status "$template_id" &>/dev/null; then
        log_error "Template $template_id does not exist"
        exit 1
    fi

    # Deploy cluster
    if deploy_k3s_cluster "$template_id" "$base_name" \
        "$cores" "$memory" "$disk_size" "$clone_type" "$k3s_version" "$start_vmid"; then
        exit 0
    else
        log_error "Cluster deployment failed"
        exit 1
    fi
}

# Run main if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi
