#!/bin/bash
source ~/.bashrc
# Azure Arc Resource Bridge Setup Script
# Deploys Azure Arc Resource Bridge for Proxmox VM lifecycle management
# This uses a K3s-based approach for the Resource Bridge

set -e

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

# Azure configuration
RESOURCE_GROUP="${RESOURCE_GROUP:-HC-Stack}"
TENANT_ID="${TENANT_ID:-}"
LOCATION="${LOCATION:-eastus}"
SUBSCRIPTION_ID="${SUBSCRIPTION_ID:-}"
CLUSTER_NAME="${CLUSTER_NAME:-proxmox-arc-bridge}"

# K3s configuration
K3S_NODE_IP="${K3S_NODE_IP:-}"
K3S_USER="${K3S_USER:-root}"
SSH_KEY="${SSH_KEY:-}"

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

validate_config() {
    if [ -z "$TENANT_ID" ] || [ -z "$SUBSCRIPTION_ID" ] || [ -z "$RESOURCE_GROUP" ]; then
        log_error "Required Azure configuration missing"
        exit 1
    fi
    
    if [ -z "$K3S_NODE_IP" ]; then
        log_error "K3S_NODE_IP must be set (IP of node where K3s will run)"
        exit 1
    fi
    
    if ! command -v az &> /dev/null; then
        log_error "Azure CLI not found"
        exit 1
    fi
    
    if ! command -v kubectl &> /dev/null; then
        log_error "kubectl not found"
        exit 1
    fi
}

check_k3s_installed() {
    log_info "Checking K3s installation on $K3S_NODE_IP..."
    
    if [ -n "$SSH_KEY" ]; then
        SSH_CMD="ssh -i $SSH_KEY -o StrictHostKeyChecking=no $K3S_USER@$K3S_NODE_IP"
    else
        SSH_CMD="ssh -o StrictHostKeyChecking=no $K3S_USER@$K3S_NODE_IP"
    fi
    
    if $SSH_CMD "command -v k3s &>/dev/null"; then
        log_info "K3s is installed"
        $SSH_CMD "k3s --version"
        return 0
    else
        log_warn "K3s not found. Please install K3s first using k3s-install.sh"
        return 1
    fi
}

get_k3s_kubeconfig() {
    log_info "Retrieving K3s kubeconfig..."
    
    # Get kubeconfig from remote K3s node
    if [ -n "$SSH_KEY" ]; then
        ssh -i "$SSH_KEY" -o StrictHostKeyChecking=no "$K3S_USER@$K3S_NODE_IP" \
            "sudo cat /etc/rancher/k3s/k3s.yaml" > /tmp/k3s-kubeconfig.yaml
    else
        ssh -o StrictHostKeyChecking=no "$K3S_USER@$K3S_NODE_IP" \
            "sudo cat /etc/rancher/k3s/k3s.yaml" > /tmp/k3s-kubeconfig.yaml
    fi
    
    # Update server URL to use node IP
    sed -i "s/127.0.0.1/$K3S_NODE_IP/g" /tmp/k3s-kubeconfig.yaml
    
    export KUBECONFIG=/tmp/k3s-kubeconfig.yaml
    
    # Verify connection
    if kubectl cluster-info &>/dev/null; then
        log_info "Successfully connected to K3s cluster"
        kubectl get nodes
    else
        log_error "Failed to connect to K3s cluster"
        exit 1
    fi
}

onboard_k8s_to_arc() {
    log_info "Onboarding Kubernetes cluster to Azure Arc..."
    
    # Check if already onboarded
    if az arc kubernetes show \
        --resource-group "$RESOURCE_GROUP" \
        --name "$CLUSTER_NAME" &>/dev/null; then
        log_warn "Cluster already onboarded to Azure Arc"
        return
    fi
    
    # Install Azure Arc extensions for Kubernetes
    log_info "Installing Azure Arc extensions..."
    az extension add --name connectedk8s --upgrade || true
    az extension add --name k8s-extension --upgrade || true
    
    # Connect cluster to Azure Arc
    log_info "Connecting cluster to Azure Arc..."
    az connectedk8s connect \
        --resource-group "$RESOURCE_GROUP" \
        --name "$CLUSTER_NAME" \
        --location "$LOCATION" \
        --tags "type=proxmox-resource-bridge"
    
    log_info "Waiting for cluster to be connected..."
    sleep 30
    
    # Verify connection
    if az arc kubernetes show \
        --resource-group "$RESOURCE_GROUP" \
        --name "$CLUSTER_NAME" \
        --query "connectivityStatus" -o tsv | grep -q "Connected"; then
        log_info "Cluster successfully connected to Azure Arc"
    else
        log_error "Cluster connection failed or still pending"
        log_info "Check status: az arc kubernetes show -g $RESOURCE_GROUP -n $CLUSTER_NAME"
    fi
}

install_gitops_extension() {
    log_info "Installing GitOps extension for Azure Arc Kubernetes..."
    
    # Install GitOps extension
    az k8s-extension create \
        --resource-group "$RESOURCE_GROUP" \
        --cluster-name "$CLUSTER_NAME" \
        --cluster-type connectedClusters \
        --extension-type microsoft.flux \
        --name flux \
        --scope cluster \
        --release-namespace flux-system
    
    log_info "GitOps extension installed"
    log_info "This may take a few minutes to complete. Check status with:"
    log_info "  az k8s-extension show -g $RESOURCE_GROUP -c $CLUSTER_NAME -t connectedClusters -n flux"
}

create_custom_location() {
    log_info "Creating custom location for Resource Bridge..."
    
    CUSTOM_LOCATION_NAME="${CLUSTER_NAME}-location"
    
    # Get cluster ID
    CLUSTER_ID=$(az arc kubernetes show \
        --resource-group "$RESOURCE_GROUP" \
        --name "$CLUSTER_NAME" \
        --query "id" -o tsv)
    
    # Create custom location
    az customlocation create \
        --resource-group "$RESOURCE_GROUP" \
        --name "$CUSTOM_LOCATION_NAME" \
        --host-resource-id "$CLUSTER_ID" \
        --namespace arc-resource-bridge \
        --location "$LOCATION"
    
    log_info "Custom location created: $CUSTOM_LOCATION_NAME"
}

main() {
    log_info "Starting Azure Arc Resource Bridge setup..."
    validate_config
    
    if ! check_k3s_installed; then
        log_error "K3s must be installed first. Run k3s-install.sh"
        exit 1
    fi
    
    get_k3s_kubeconfig
    onboard_k8s_to_arc
    install_gitops_extension
    create_custom_location
    
    log_info "Azure Arc Resource Bridge setup completed!"
    log_info "Next steps:"
    log_info "  1. Configure Proxmox custom provider for VM lifecycle control"
    log_info "  2. Set up GitOps repository for declarative deployments"
    log_info "  3. View cluster in Azure Portal:"
    log_info "     https://portal.azure.com/#view/Microsoft_Azure_HybridCompute/KubernetesBlade"
}

main "$@"