#!/usr/bin/env bash
# Kloak Demo Setup Script (Lima + K3s)
# Creates a Lima VM with K3s and deploys the Kloak demo
# Certificates are auto-generated by the controller in auto-cert mode

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(dirname "$SCRIPT_DIR")"
KUBECONFIG_PATH="/tmp/kloak-k3s.yaml"
LIMA_INSTANCE="kloak"
DEMO_NAMESPACE="kloak-demo"

echo "================================"
echo " Kloak Demo Setup (Lima + K3s)"
echo "================================"
echo ""
echo "PATH: $PATH"

# Check prerequisites
check_prereqs() {
    echo "Checking prerequisites..."
    for cmd in limactl kubectl docker helm; do
        if ! command -v $cmd &> /dev/null; then
             echo "Debug: command -v $cmd failed"
             echo "Debug: which $cmd: $(which $cmd)"
            echo "Error: $cmd is required but not installed."
            exit 1
        fi
    done
    echo "✓ Prerequisites OK"
}

# Create/Start Lima VM with K3s
create_cluster() {
    echo ""
    echo "Starting Lima VM with K3s..."
    
    # Check if instance exists and has K3s
    if limactl list -q | grep -q "^${LIMA_INSTANCE}$"; then
        echo "Checking existing instance state..."
        # If K3s is missing, it's likely the old VM without port forwarding.
        if ! limactl shell "${LIMA_INSTANCE}" -- which k3s &>/dev/null; then
            echo "Existing instance lacks K3s configuration (or port forwarding). Recreating..."
            limactl delete -f "${LIMA_INSTANCE}"
        else
            echo "Instance '${LIMA_INSTANCE}' appears compatible."
        fi
    fi

    # Create/Start
    if ! limactl list -q | grep -q "^${LIMA_INSTANCE}$"; then
        echo "Creating Lima instance '${LIMA_INSTANCE}' from lima.yaml..."
        limactl start --name="${LIMA_INSTANCE}" "$ROOT_DIR/lima.yaml"
    else
        echo "Starting existing Lima instance '${LIMA_INSTANCE}'..."
        limactl start "${LIMA_INSTANCE}"
    fi

    # Ensure K3s is installed (in case provision script skipped or failed)
    echo "Ensuring K3s is installed..."
    limactl shell "${LIMA_INSTANCE}" -- sudo sh -c 'if [ ! -f /usr/local/bin/k3s ]; then curl -sfL https://get.k3s.io | sh -s - --write-kubeconfig-mode 644 --disable traefik; fi'

    # Wait for K3s service
    echo "Waiting for K3s service..."
    until limactl shell "${LIMA_INSTANCE}" -- sudo systemctl is-active k3s >/dev/null 2>&1; do
         sleep 5
         echo -n "."
    done

    # Wait for K3s kubeconfig
    echo "Waiting for K3s kubeconfig..."
    until limactl shell "${LIMA_INSTANCE}" -- sudo test -f /etc/rancher/k3s/k3s.yaml; do
        sleep 5
        echo -n "."
    done
    echo ""

    # Fetch and fix kubeconfig
    echo "Fetching kubeconfig..."
    limactl shell "${LIMA_INSTANCE}" -- sudo cat /etc/rancher/k3s/k3s.yaml > "$KUBECONFIG_PATH"
    # Ensure rights
    chmod 600 "$KUBECONFIG_PATH"
    # K3s uses 127.0.0.1:6443 which works via Lima port forwarding
    
    export KUBECONFIG="$KUBECONFIG_PATH"
    echo "✓ K3s ready. Kubeconfig: $KUBECONFIG_PATH"
    
    # Wait for node ready
    echo "Waiting for K3s node readiness..."
    kubectl wait --for=condition=Ready node --all --timeout=60s
}

# Build images and import to K3s
build_images() {
    echo ""
    echo "Building images..."
    
    # Build demo Python app
    docker build -t kloak-demo-python:latest "$SCRIPT_DIR/demo-python"

    # Build demo Go app
    docker build -t kloak-demo-go:latest "$SCRIPT_DIR/demo-go"

    # Build demo JS app
    docker build -t kloak-demo-js:latest "$SCRIPT_DIR/demo-js"

    # Build demo Python raw TLS app
    docker build -t kloak-demo-python-raw-tls:latest "$SCRIPT_DIR/demo-python-raw-tls"

    # Build Kloak controller
    docker build -t kloak:latest "$ROOT_DIR"

    echo "Importing images into K3s (this may take a moment)..."
    # We pipeline docker save -> lima -> k3s ctr import
    
    echo "Importing kloak-demo-python..."
    docker save kloak-demo-python:latest | limactl shell "${LIMA_INSTANCE}" -- sudo k3s ctr images import -

    echo "Importing kloak-demo-go..."
    docker save kloak-demo-go:latest | limactl shell "${LIMA_INSTANCE}" -- sudo k3s ctr images import -

    echo "Importing kloak-demo-js..."
    docker save kloak-demo-js:latest | limactl shell "${LIMA_INSTANCE}" -- sudo k3s ctr images import -

    echo "Importing kloak-demo-python-raw-tls..."
    docker save kloak-demo-python-raw-tls:latest | limactl shell "${LIMA_INSTANCE}" -- sudo k3s ctr images import -

    echo "Importing kloak..."
    docker save kloak:latest | limactl shell "${LIMA_INSTANCE}" -- sudo k3s ctr images import -

    echo "✓ Images built and imported"
}

# NOTE: Certificates are now auto-generated by the controller in auto-cert mode.
# The generate_certs function is no longer needed.

# Deploy Kloak components using Helm chart
deploy_kloak() {
    export KUBECONFIG="$KUBECONFIG_PATH"
    echo ""
    echo "Deploying Kloak using Helm chart (auto-cert mode)..."

    # Install or upgrade Kloak via Helm
    helm upgrade --install kloak "$ROOT_DIR/charts/kloak" \
        -n kloak-system --create-namespace \
        --set image.repository=kloak \
        --set image.tag=latest \
        --set image.pullPolicy=Never \
        --set log.level=debug \
        --set log.dev=true \
        --wait --timeout 120s

    until kubectl get secret kloak-webhook-certs -n kloak-system 2>/dev/null; do
        echo -n "."
        sleep 2
    done
    echo ""
    echo "✓ Webhook certificate generated"

    # Remove default namespace label if it exists (cleanup)
    kubectl label namespace default getkloak.io/enabled- --overwrite 2>/dev/null || true

    echo "✓ Kloak deployed (controller DaemonSet, webhook)"
}

# Wait for Kloak pods
wait_for_kloak() {
    export KUBECONFIG="$KUBECONFIG_PATH"
    echo ""
    echo "Waiting for Kloak pods to be ready..."
    
    # Use rollout status instead of kubectl wait.
    # After deploy_kloak does "rollout restart", the old pods go into Terminating
    # state with Ready=False. "kubectl wait --for=condition=Ready" matches those
    # dying pods and blocks until timeout. "rollout status" properly tracks the
    # rolling update and returns as soon as new pods are ready.
    kubectl rollout status daemonset/kloak-controller -n kloak-system --timeout=120s
    kubectl rollout status deployment/kloak-webhook -n kloak-system --timeout=120s
    
    echo ""
    echo "Kloak pods status:"
    kubectl get pods -n kloak-system
}

# Deploy demo app
deploy_demo() {
    export KUBECONFIG="$KUBECONFIG_PATH"
    echo ""
    echo "Deploying demo application to $DEMO_NAMESPACE..."

    # Create Demo Namespace
    echo "Creating demo namespace: $DEMO_NAMESPACE"
    kubectl create namespace "$DEMO_NAMESPACE" --dry-run=client -o yaml | kubectl apply -f -

    # Label demo namespace to enable webhook and CA sync
    kubectl label namespace "$DEMO_NAMESPACE" getkloak.io/enabled=true --overwrite
    
    # Delete any existing demo deployment/pod first
    kubectl delete deployment demo-python -n "$DEMO_NAMESPACE" --ignore-not-found 2>/dev/null || true
    kubectl delete pod demo-python -n "$DEMO_NAMESPACE" --ignore-not-found 2>/dev/null || true
    kubectl delete deployment demo-go -n "$DEMO_NAMESPACE" --ignore-not-found 2>/dev/null || true
    kubectl delete pod demo-go -n "$DEMO_NAMESPACE" --ignore-not-found 2>/dev/null || true
    kubectl delete deployment demo-js -n "$DEMO_NAMESPACE" --ignore-not-found 2>/dev/null || true
    kubectl delete pod demo-js -n "$DEMO_NAMESPACE" --ignore-not-found 2>/dev/null || true
    kubectl delete deployment demo-python-raw-tls -n "$DEMO_NAMESPACE" --ignore-not-found 2>/dev/null || true
    kubectl delete service tls -n "$DEMO_NAMESPACE" --ignore-not-found 2>/dev/null || true
    
    # Create SECRET 1: Allowed for httpbin.org
    echo "Creating secret-allowed (hosts=httpbin.org)..."
    kubectl create secret generic secret-allowed \
        --from-literal=api-key="REAL-ALLOWED-KEY-12345" \
        -n "$DEMO_NAMESPACE" --dry-run=client -o yaml | \
        kubectl label -f - getkloak.io/enabled="true" --local -o yaml | \
        kubectl annotate -f - getkloak.io/hosts="httpbin.org" --local -o yaml | \
        kubectl apply -f -
    
    # Create SECRET for raw TLS demo: Allowed for tls.kloak-demo (short FQDN fits MAX_HOST_LEN=32)
    echo "Creating secret-rawtls-allowed (hosts=tls.kloak-demo.svc.cluster.local)..."
    kubectl create secret generic secret-rawtls-allowed \
        --from-literal=api-key="REAL-ALLOWED-KEY-12345" \
        -n "$DEMO_NAMESPACE" --dry-run=client -o yaml | \
        kubectl label -f - getkloak.io/enabled="true" --local -o yaml | \
        kubectl annotate -f - getkloak.io/hosts="tls.kloak-demo.svc.cluster.local" --local -o yaml | \
        kubectl apply -f -

    # Create SECRET for raw TLS demo: Blocked (only allowed for example.com)
    echo "Creating secret-rawtls-blocked (hosts=example.com)..."
    kubectl create secret generic secret-rawtls-blocked \
        --from-literal=api-key="REAL-BLOCKED-KEY-67890" \
        -n "$DEMO_NAMESPACE" --dry-run=client -o yaml | \
        kubectl label -f - getkloak.io/enabled="true" --local -o yaml | \
        kubectl annotate -f - getkloak.io/hosts="example.com" --local -o yaml | \
        kubectl apply -f -

    # Create SECRET 2: Blocked for httpbin.org (only allowed for example.com)
    echo "Creating secret-blocked (hosts=example.com)..."
    kubectl create secret generic secret-blocked \
        --from-literal=api-key="REAL-BLOCKED-KEY-67890" \
        -n "$DEMO_NAMESPACE" --dry-run=client -o yaml | \
        kubectl label -f - getkloak.io/enabled="true" --local -o yaml | \
        kubectl annotate -f - getkloak.io/hosts="example.com" --local -o yaml | \
        kubectl apply -f -
    
    # Wait a moment for cleanup
    sleep 2
    
    # Apply the demo deployments (webhook should mutate pods)
    kubectl apply -f "$SCRIPT_DIR/demo-python/deployment.yaml" -n "$DEMO_NAMESPACE"
    kubectl apply -f "$SCRIPT_DIR/demo-go/deployment.yaml" -n "$DEMO_NAMESPACE"
    kubectl apply -f "$SCRIPT_DIR/demo-js/deployment.yaml" -n "$DEMO_NAMESPACE"
    kubectl apply -f "$SCRIPT_DIR/demo-python-raw-tls/deployment.yaml" -n "$DEMO_NAMESPACE"

    echo "✓ Demo applications deployed (Python + Go + JS)"
}

# Wait for demo deployment
wait_for_demo() {
    export KUBECONFIG="$KUBECONFIG_PATH"
    echo ""
    echo "Waiting for demo deployment to be ready..."
    
    kubectl rollout status deployment/demo-python -n "$DEMO_NAMESPACE" --timeout=120s || {
        echo "Warning: Python demo deployment may not be ready"
    }
    
    kubectl rollout status deployment/demo-go -n "$DEMO_NAMESPACE" --timeout=120s || {
        echo "Warning: Go demo deployment may not be ready"
    }

    kubectl rollout status deployment/demo-js -n "$DEMO_NAMESPACE" --timeout=120s || {
        echo "Warning: JS demo deployment may not be ready"
    }

    echo ""
    echo "Demo pods status:"
    kubectl get pods -n "$DEMO_NAMESPACE" -o wide

    echo ""
    echo "Containers in Python demo pod:"
    kubectl get pods -l app=demo-python -n "$DEMO_NAMESPACE" -o jsonpath='{.items[0].spec.containers[*].name}' && echo ""

    echo "Containers in Go demo pod:"
    kubectl get pods -l app=demo-go -n "$DEMO_NAMESPACE" -o jsonpath='{.items[0].spec.containers[*].name}' && echo ""

    echo "Containers in JS demo pod:"
    kubectl get pods -l app=demo-js -n "$DEMO_NAMESPACE" -o jsonpath='{.items[0].spec.containers[*].name}' && echo ""
}

# Verify webhook mutation
verify_mutation() {
    export KUBECONFIG="$KUBECONFIG_PATH"
    echo ""
    echo "Verifying webhook mutation..."
    
    for app in demo-python demo-go demo-js; do
        echo "Checking $app..."
        ANNOTATION=$(kubectl get pods -l app=$app -n "$DEMO_NAMESPACE" -o jsonpath='{.items[0].metadata.annotations.getkloak\.io/enabled}' 2>/dev/null || echo "")
        
        if [ "$ANNOTATION" = "true" ]; then
            echo "✓ [$app] Pod successfully mutated by webhook!"
        else
            echo "✗ [$app] Pod NOT mutated"
        fi
    done
    
    # Show webhook logs if any mutation failed
    if [ "$(kubectl get pods -l app=demo-python -n "$DEMO_NAMESPACE" -o jsonpath='{.items[0].metadata.annotations.getkloak\.io/enabled}' 2>/dev/null)" != "true" ] || \
       [ "$(kubectl get pods -l app=demo-go -n "$DEMO_NAMESPACE" -o jsonpath='{.items[0].metadata.annotations.getkloak\.io/enabled}' 2>/dev/null)" != "true" ] || \
       [ "$(kubectl get pods -l app=demo-js -n "$DEMO_NAMESPACE" -o jsonpath='{.items[0].metadata.annotations.getkloak\.io/enabled}' 2>/dev/null)" != "true" ]; then
        echo ""
        echo "Checking webhook logs..."
        kubectl logs -n kloak-system -l app.kubernetes.io/component=webhook --tail=10 2>/dev/null || true
    fi
}

# Show summary
show_summary() {
    echo ""
    echo "================================"
    echo " Demo Setup Complete!"
    echo "================================"
    echo ""
    echo "Architecture:"
    echo "  Controller:          DaemonSet (runs on each node for eBPF)"
    echo "  Webhook:             Deployment"
    echo ""
    echo "Environment:"
    echo "  Kubeconfig:          $KUBECONFIG_PATH"
    echo "  Helm chart:          $ROOT_DIR/charts/kloak"
    echo ""
    echo "Commands:"
    echo "  Use kubectl:         export KUBECONFIG=$KUBECONFIG_PATH"
    echo "  View Python logs:    kubectl logs -f -l app=demo-python -n $DEMO_NAMESPACE -c demo-app"
    echo "  View Go logs:        kubectl logs -f -l app=demo-go -n $DEMO_NAMESPACE -c demo-app"
    echo "  View JS logs:        kubectl logs -f -l app=demo-js -n $DEMO_NAMESPACE -c demo-app"
    echo "  View Raw TLS logs:   kubectl logs -f -l app=demo-python-raw-tls -n $DEMO_NAMESPACE -c demo-app"
    echo "  View webhook logs:   kubectl logs -n kloak-system -l app.kubernetes.io/component=webhook"
    echo "  View controller logs: kubectl logs -n kloak-system -l app.kubernetes.io/component=controller"
    echo ""
    echo "Verification (for demo-python, demo-go, or demo-js):"
    echo "  1. Check demo logs:  kubectl logs -l app=demo-python -n $DEMO_NAMESPACE -c demo-app | head -30"
    echo "     (Should show 'kloak:...' UUIDs for both secrets)"
    echo "  2. Check response:   kubectl logs -l app=demo-python -n $DEMO_NAMESPACE -c demo-app | grep -A20 'Response headers'"
    echo "     X-Secret-Allowed: Should show 'REAL-ALLOWED-KEY-12345' (replaced)"
    echo "     X-Secret-Blocked: Should show 'kloak:...' UUID (NOT replaced - wrong host)"
    echo ""
}

# Main
main() {
    check_prereqs
    create_cluster
    build_images
    # Certs are auto-generated by controller in auto-cert mode
    deploy_kloak
    wait_for_kloak
    deploy_demo
    wait_for_demo
    verify_mutation
    show_summary
}

main "$@"
