#!/bin/bash

set -e
set -o pipefail

# Get WC_ROOT (where this script and templates live)
WC_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.."; pwd)"
export WC_ROOT

# Set up cloud directory (WC_HOME is where user's cloud will be)
WC_HOME="$(pwd)"
export WC_HOME

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

# Helper functions
print_header() {
    echo -e "\n${BLUE}=== $1 ===${NC}\n"
}

print_info() {
    echo -e "${BLUE}INFO:${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}WARNING:${NC} $1"
}

print_success() {
    echo -e "${GREEN}SUCCESS:${NC} $1"
}

print_error() {
    echo -e "${RED}ERROR:${NC} $1"
}

# Function to get current config value safely
get_current_config() {
    local key="$1"
    if [ -f "${WC_HOME}/config.yaml" ]; then
        set +e
        result=$(wild-config "${key}" 2>/dev/null)
        set -e
        echo "${result}"
    else
        echo ""
    fi
}

# Usage function
usage() {
    echo "Usage: wild-cluster-node-machine-config-generate <node-ip>"
    echo ""
    echo "Generate Talos machine configuration for a specific registered node."
    echo ""
    echo "Arguments:"
    echo "  node-ip             IP address of the registered node"
    echo ""
    echo "Options:"
    echo "  -h, --help          Show this help message"
    echo ""
    echo "Examples:"
    echo "  wild-cluster-node-machine-config-generate 192.168.1.91"
    echo "  wild-cluster-node-machine-config-generate 192.168.1.100"
    echo ""
    echo "This script will:"
    echo "  - Generate initial cluster secrets if not present"
    echo "  - Use patch templates from the wild-cloud repository"
    echo "  - Create machine configuration for the specified node"
    echo "  - Generate patched config with node-specific hardware settings"
    echo "  - Update talosctl context with the node"
    echo ""
    echo "Requirements:"
    echo "  - Must be run from a wild-cloud directory"
    echo "  - Node must be registered (hardware detected) first"
    echo "  - Basic cluster configuration must be completed"
    echo "  - Patch templates must exist in WC_ROOT/setup/cluster-nodes/"
}

# Parse arguments
NODE_IP=""
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            usage
            exit 0
            ;;
        -*)
            echo "Unknown option $1"
            usage
            exit 1
            ;;
        *)
            if [ -z "$NODE_IP" ]; then
                NODE_IP="$1"
            else
                echo "Unexpected argument: $1"
                usage
                exit 1
            fi
            shift
            ;;
    esac
done

# Check if node IP was provided
if [ -z "$NODE_IP" ]; then
    echo "Error: Node IP address is required"
    usage
    exit 1
fi

# Check if we're in a wild-cloud directory
if [ ! -d ".wildcloud" ]; then
    print_error "You must run this script from a wild-cloud directory"
    print_info "Run 'wild-setup' or 'wild-init' first to initialize a wild-cloud project"
    exit 1
fi

# Check required configuration
if [ -z "$(get_current_config "cluster.name")" ]; then
    print_error "Basic cluster configuration is missing"
    print_info "Run 'wild-setup' or 'wild-init' first to configure your cluster"
    exit 1
fi

# Function to ensure required directories exist in WC_HOME
ensure_required_directories() {
    # Create output directories in WC_HOME for patch and final configs
    mkdir -p "${WC_HOME}/setup/cluster-nodes/patch"
    mkdir -p "${WC_HOME}/setup/cluster-nodes/final"
    
    # Ensure the generated directory exists (for cluster secrets)
    mkdir -p "${WC_HOME}/setup/cluster-nodes/generated"
}

# =============================================================================
# MACHINE CONFIG GENERATION
# =============================================================================

print_header "Talos Machine Config Generation"

# Ensure required directories exist in WC_HOME
ensure_required_directories

# Define directories
TEMPLATE_SOURCE_DIR="${WC_ROOT}/setup/cluster-nodes"
NODE_SETUP_DIR="${WC_HOME}/setup/cluster-nodes"

# Check if cluster has been initialized
if [ ! -f "${NODE_SETUP_DIR}/generated/secrets.yaml" ]; then
    print_error "Cluster not initialized. Base cluster configuration is required."
    print_info "Run 'wild-cluster-config-generate' first to generate cluster secrets and base configs"
    exit 1
fi

# Get cluster configuration from config.yaml
CLUSTER_NAME=$(wild-config cluster.name)
VIP=$(wild-config cluster.nodes.control.vip)

print_info "Generating machine configuration for node: $NODE_IP"
print_info "Cluster: $CLUSTER_NAME"

# Check if the specified node is registered
NODE_INTERFACE=$(yq eval ".cluster.nodes.active.\"${NODE_IP}\".interface" "${WC_HOME}/config.yaml" 2>/dev/null)
NODE_DISK=$(yq eval ".cluster.nodes.active.\"${NODE_IP}\".disk" "${WC_HOME}/config.yaml" 2>/dev/null)
IS_CONTROL=$(yq eval ".cluster.nodes.active.\"${NODE_IP}\".control" "${WC_HOME}/config.yaml" 2>/dev/null)

if [ -z "$NODE_INTERFACE" ] || [ "$NODE_INTERFACE" = "null" ]; then
    print_error "Node $NODE_IP is not registered in config.yaml"
    print_info "Please register the node first by running node hardware detection:"
    print_info "  wild-node-detect $NODE_IP"
    print_info "Or run 'wild-setup' to register nodes interactively"
    exit 1
fi

# Determine node type
if [ "$IS_CONTROL" = "true" ]; then
    NODE_TYPE="control"
    print_success "Registered control plane node: $NODE_IP"
else
    NODE_TYPE="worker" 
    print_success "Registered worker node: $NODE_IP"
fi

print_info "Node details:"
print_info "  - Interface: $NODE_INTERFACE"
print_info "  - Disk: $NODE_DISK"
print_info "  - Type: $NODE_TYPE"

# Compile patch template for the specified node
print_info "Compiling patch template for $NODE_TYPE node $NODE_IP..."

if [ "$NODE_TYPE" = "control" ]; then
    TEMPLATE_FILE="${TEMPLATE_SOURCE_DIR}/patch.templates/controlplane.yaml"
    BASE_CONFIG="${NODE_SETUP_DIR}/generated/controlplane.yaml"
else
    TEMPLATE_FILE="${TEMPLATE_SOURCE_DIR}/patch.templates/worker.yaml" 
    BASE_CONFIG="${NODE_SETUP_DIR}/generated/worker.yaml"
fi

# Use IP as the patch name and output config name
PATCH_FILE="${NODE_SETUP_DIR}/patch/${NODE_IP}.yaml"
OUTPUT_CONFIG="${NODE_SETUP_DIR}/final/${NODE_IP}.yaml"

# Check if the patch template exists
if [ ! -f "$TEMPLATE_FILE" ]; then
    print_error "Patch template not found: $TEMPLATE_FILE"
    print_info "Make sure the wild-cloud repository is properly set up"
    exit 1
fi

# Create a temporary template with the node IP for gomplate processing
TEMP_TEMPLATE="/tmp/${NODE_IP//\//_}-$(date +%s).yaml"
sed "s/{{NODE_IP}}/${NODE_IP}/g" "$TEMPLATE_FILE" > "$TEMP_TEMPLATE"
cat "$TEMP_TEMPLATE" | wild-compile-template > "$PATCH_FILE"
rm -f "$TEMP_TEMPLATE"

# Generate final machine config for the specified node
print_info "Generating final machine configuration..."
talosctl machineconfig patch "$BASE_CONFIG" --patch @"$PATCH_FILE" -o "$OUTPUT_CONFIG"

# Update talosctl context with this node
print_info "Updating talosctl context..."
talosctl config node "$NODE_IP"

print_success "Machine configuration generated successfully!"
echo ""
print_info "Generated files:"
print_info "  - Patch: $PATCH_FILE"
print_info "  - Final config: $OUTPUT_CONFIG"
echo ""
print_info "Template used: ${TEMPLATE_FILE}"

echo ""
print_info "Next steps:"
echo "  1. Apply configuration to the node:"
echo "     talosctl apply-config -i -n $NODE_IP -f $OUTPUT_CONFIG"
echo ""
if [ "$NODE_TYPE" = "control" ]; then
    echo "  2. If this is your first control plane node, bootstrap it:"
    echo "     talosctl bootstrap -n $NODE_IP"
    echo ""
    echo "  3. Get kubeconfig when cluster is ready:"
    echo "     talosctl kubeconfig"
else
    echo "  2. Node will join the cluster automatically after applying config"
fi

print_success "Machine config generation completed!"