#!/bin/bash
# =============================================================================
# Klyro Edge Installer
# https://get.klyro.security/install.sh
#
# Usage:
#   Production: curl -sfL https://get.klyro.security | sh -s -- <INSTALL_TOKEN>
#   Local:      ./install.sh inst_test --local
#
# Get your install token from: https://app.klyro.security/agents/new
#
# Local Mode (--local):
#   Sets up a complete local environment with:
#   - k3d cluster with local registry
#   - Mock Cloud gRPC server for agent registration
#   - Temporal for workflow orchestration
#   - Agent connected to mock cloud and Temporal
#   - Worker/Operator polling Temporal for tasks
#   - SQLite database for agent state
# =============================================================================
set -e

KLYRO_API_URL="${KLYRO_API_URL:-https://api.klyro.security}"
NAMESPACE="klyro-system"
RELEASE_NAME="klyro-edge"

# Script and project directories
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="${SCRIPT_DIR}/.."

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

log_info() { printf "${BLUE}[INFO]${NC} %s\n" "$1"; }
log_success() { printf "${GREEN}[OK]${NC} %s\n" "$1"; }
log_warn() { printf "${YELLOW}[WARN]${NC} %s\n" "$1"; }
log_error() { printf "${RED}[ERROR]${NC} %s\n" "$1" >&2; }
log_step() { printf "${GREEN}>>>${NC} ${BLUE}%s${NC}\n" "$1"; }
log_local() { printf "${CYAN}[LOCAL]${NC} %s\n" "$1"; }

# Spinner for long-running operations
spin() {
    local pid=$1
    local msg=$2
    local spinstr='⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'
    local i=0
    printf "${BLUE}[...]${NC} %s " "$msg"
    while kill -0 "$pid" 2>/dev/null; do
        i=$(( (i + 1) % 10 ))
        printf "\r${BLUE}[%s]${NC} %s " "$(echo $spinstr | cut -c$((i+1)))" "$msg"
        sleep 0.1
    done
    wait "$pid"
    return $?
}

# Run command with spinner
run_with_spinner() {
    local msg="$1"
    shift
    if [ "${KLYRO_DEBUG:-}" = "true" ]; then
        printf "${BLUE}[...]${NC} %s\n" "$msg"
        "$@"
        return $?
    else
        "$@" >/dev/null 2>&1 &
        local pid=$!
        spin $pid "$msg"
        local status=$?
        if [ $status -eq 0 ]; then
            printf "\r${GREEN}[OK]${NC} %s    \n" "$msg"
        else
            printf "\r${RED}[FAIL]${NC} %s    \n" "$msg"
        fi
        return $status
    fi
}

show_banner() {
    echo ""
    printf "${BLUE}╔═══════════════════════════════════════════════════════════════╗${NC}\n"
    printf "${BLUE}║${NC}   ${GREEN}Klyro Edge Agent Installer${NC}                                ${BLUE}║${NC}\n"
    printf "${BLUE}║${NC}   AI-powered identity provisioning for Kubernetes            ${BLUE}║${NC}\n"
    printf "${BLUE}╚═══════════════════════════════════════════════════════════════╝${NC}\n"
    echo ""
}

usage() {
    cat <<EOF
Klyro Edge Installer

Usage:
  Production: OPENAI_API_KEY=sk-... curl -sfL https://get.klyro.security | sh -s -- <INSTALL_TOKEN>
  Local:      ./install.sh <TOKEN> --local
  Local+Cloud: ./install.sh <TOKEN> --local --grpc-endpoint grpc.klyro.security:443

Arguments:
  INSTALL_TOKEN       Your one-time install token from the Klyro dashboard
  --local             Enable full local mode with mock cloud, Temporal, etc.
  --grpc-endpoint     Custom gRPC endpoint (skips mock cloud setup)
  --api-endpoint      Custom API endpoint (optional, used with --grpc-endpoint)

Local Mode (--local):
  Sets up a complete local environment:
  - k3d cluster with local registry
  - Mock Cloud gRPC server for agent registration
  - Temporal for workflow orchestration
  - Agent + Worker/Operator fully connected
  - OpenAI credentials (dummy key for testing)

Local Mode with Custom Endpoints (--local --grpc-endpoint):
  Sets up local k3d cluster but connects to real/external cloud:
  - k3d cluster with local registry
  - Temporal for workflow orchestration
  - Agent connects to specified gRPC endpoint (no mock cloud)
  - Skips mock cloud connectivity tests

Production Mode:
  Requires the following environment variables:
  - OPENAI_API_KEY (mandatory) - OpenAI API key for the AI worker

Get your install token:
  1. Log in to https://app.klyro.security
  2. Navigate to Agents → Add Agent
  3. Copy the install command

Environment Variables:
  OPENAI_API_KEY   (mandatory for prod) OpenAI API key for AI-powered provisioning
  KLYRO_API_URL    Override API endpoint (default: https://api.klyro.security)
  KLYRO_DEBUG      Enable debug output (set to 'true')

Examples:
  # Full local mode with mock cloud
  ./install.sh inst_test --local

  # Local k3d cluster connecting to real cloud
  ./install.sh inst_abc123 --local --grpc-endpoint grpc.klyro.security:443

  # Production installation
  OPENAI_API_KEY=sk-... ./install.sh inst_abc123

EOF
    exit 0
}

# =============================================================================
# LOCAL MODE FUNCTIONS
# =============================================================================

check_local_prerequisites() {
    log_local "Checking local mode prerequisites..."

    local MISSING=0

    check_command "docker" "brew install --cask docker" || MISSING=1
    check_command "k3d" "brew install k3d" || MISSING=1
    check_command "go" "brew install go" || MISSING=1

    if [ "$MISSING" -eq 1 ]; then
        echo ""
        log_error "Missing local mode tools. Please install them and try again."
        exit 1
    fi

    # Check Docker daemon is running
    if ! docker info >/dev/null 2>&1; then
        log_error "Docker daemon is not running. Start Docker Desktop."
        exit 1
    fi

    log_success "Local mode prerequisites verified"
}

create_k3d_cluster() {
    local cluster_name="klyro-local"

    log_local "Setting up k3d cluster..."

    # Check if cluster already exists
    if k3d cluster list 2>/dev/null | grep -q "$cluster_name"; then
        log_info "Cluster '$cluster_name' already exists"
        k3d kubeconfig merge "$cluster_name" --kubeconfig-switch-context >/dev/null 2>&1
        export KUBECONFIG=$(k3d kubeconfig write "$cluster_name")
        return 0
    fi

    # Create cluster with local registry
    log_info "Creating k3d cluster with local registry..."
    k3d cluster create "$cluster_name" \
        --registry-create klyro-registry:0.0.0.0:5000 \
        --port "8080:80@loadbalancer" \
        --port "8443:443@loadbalancer" \
        --port "7233:7233@loadbalancer" \
        --agents 1 \
        --wait

    # Set KUBECONFIG
    export KUBECONFIG=$(k3d kubeconfig write "$cluster_name")

    # Wait for cluster to be ready
    log_info "Waiting for cluster to be ready..."
    kubectl wait --for=condition=Ready nodes --all --timeout=120s

    log_success "k3d cluster '$cluster_name' created"
}

build_local_images() {
    log_local "Building Docker images..."

    local registry="localhost:5000"

    # Generate proto if needed
    if [ ! -d "${PROJECT_ROOT}/gen/go" ]; then
        log_info "Generating proto files..."
        cd "${PROJECT_ROOT}"
        make proto || {
            log_warn "Proto generation failed - checking if gen/ exists..."
            if [ ! -d "${PROJECT_ROOT}/gen/go" ]; then
                log_error "Proto files not found. Run 'make proto' first."
                exit 1
            fi
        }
    fi

    # Build agent
    log_info "Building agent image..."
    docker build -t ${registry}/klyro-agent:local \
        -f "${PROJECT_ROOT}/services/agent/Dockerfile" \
        "${PROJECT_ROOT}" || {
        log_error "Failed to build agent image"
        exit 1
    }
    docker push ${registry}/klyro-agent:local
    log_success "Agent image built and pushed"

    # Build worker/operator
    log_info "Building operator/worker image..."
    docker build -t ${registry}/klyro-operator:local \
        -f "${PROJECT_ROOT}/services/operator/Dockerfile" \
        "${PROJECT_ROOT}/services/operator" || {
        log_error "Failed to build operator image"
        exit 1
    }
    docker push ${registry}/klyro-operator:local
    log_success "Operator image built and pushed"

    # Build mock-cloud
    log_info "Building mock-cloud image..."
    docker build -t ${registry}/mock-cloud:local \
        -f "${PROJECT_ROOT}/services/agent/cmd/mock-cloud/Dockerfile" \
        "${PROJECT_ROOT}" || {
        log_error "Failed to build mock-cloud image"
        exit 1
    }
    docker push ${registry}/mock-cloud:local
    log_success "Mock-cloud image built and pushed"

    log_success "All images built and pushed to local registry"
}

start_mock_cloud() {
    log_local "Starting Mock Cloud gRPC server..."

    # Check if already running
    if kubectl get deployment mock-cloud -n "$NAMESPACE" >/dev/null 2>&1; then
        log_info "Mock Cloud already deployed"
        return 0
    fi

    # Create mock-cloud deployment
    kubectl apply -f - <<EOF
apiVersion: apps/v1
kind: Deployment
metadata:
  name: mock-cloud
  namespace: $NAMESPACE
spec:
  replicas: 1
  selector:
    matchLabels:
      app: mock-cloud
  template:
    metadata:
      labels:
        app: mock-cloud
    spec:
      containers:
      - name: mock-cloud
        image: klyro-registry:5000/mock-cloud:local
        imagePullPolicy: Always
        ports:
        - containerPort: 50051
          name: grpc
        - containerPort: 50052
          name: http
        resources:
          requests:
            cpu: 50m
            memory: 64Mi
          limits:
            cpu: 200m
            memory: 128Mi
---
apiVersion: v1
kind: Service
metadata:
  name: mock-cloud
  namespace: $NAMESPACE
spec:
  selector:
    app: mock-cloud
  ports:
  - name: grpc
    port: 50051
    targetPort: 50051
  - name: http
    port: 50052
    targetPort: 50052
EOF

    # Wait for mock-cloud to be ready
    log_info "Waiting for Mock Cloud to be ready..."
    kubectl wait --for=condition=Available deployment/mock-cloud -n "$NAMESPACE" --timeout=120s || {
        log_warn "Mock Cloud not ready yet, continuing..."
    }

    log_success "Mock Cloud deployed"
}

start_temporal() {
    log_local "Starting Temporal..."

    # Check if already running
    if kubectl get pods -n "$NAMESPACE" -l app=temporal 2>/dev/null | grep -q "Running"; then
        log_info "Temporal already running"
        return 0
    fi

    # Deploy Temporal using auto-setup image (simpler than Helm for local dev)
    log_info "Deploying Temporal with PostgreSQL..."

    # Deploy PostgreSQL first
    kubectl apply -f - <<EOF
---
apiVersion: v1
kind: Service
metadata:
  name: temporal-postgresql
  namespace: $NAMESPACE
spec:
  selector:
    app: temporal-postgresql
  ports:
  - port: 5432
    targetPort: 5432
---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: temporal-postgresql
  namespace: $NAMESPACE
spec:
  replicas: 1
  selector:
    matchLabels:
      app: temporal-postgresql
  template:
    metadata:
      labels:
        app: temporal-postgresql
    spec:
      containers:
      - name: postgresql
        image: postgres:15-alpine
        ports:
        - containerPort: 5432
        env:
        - name: POSTGRES_USER
          value: temporal
        - name: POSTGRES_PASSWORD
          value: temporal
        - name: POSTGRES_DB
          value: temporal
        resources:
          requests:
            cpu: 100m
            memory: 256Mi
          limits:
            cpu: 500m
            memory: 512Mi
---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: temporal
  namespace: $NAMESPACE
spec:
  replicas: 1
  selector:
    matchLabels:
      app: temporal
  template:
    metadata:
      labels:
        app: temporal
    spec:
      containers:
      - name: temporal
        image: temporalio/auto-setup:1.24.2
        ports:
        - containerPort: 7233
          name: grpc
        env:
        - name: DB
          value: postgres12
        - name: DB_PORT
          value: "5432"
        - name: POSTGRES_USER
          value: temporal
        - name: POSTGRES_PWD
          value: temporal
        - name: POSTGRES_SEEDS
          value: temporal-postgresql
        resources:
          requests:
            cpu: 100m
            memory: 256Mi
          limits:
            cpu: 500m
            memory: 512Mi
---
apiVersion: v1
kind: Service
metadata:
  name: temporal-frontend
  namespace: $NAMESPACE
spec:
  selector:
    app: temporal
  ports:
  - name: grpc
    port: 7233
    targetPort: 7233
---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: temporal-ui
  namespace: $NAMESPACE
spec:
  replicas: 1
  selector:
    matchLabels:
      app: temporal-ui
  template:
    metadata:
      labels:
        app: temporal-ui
    spec:
      containers:
      - name: temporal-ui
        image: temporalio/ui:2.26.2
        ports:
        - containerPort: 8080
        env:
        - name: TEMPORAL_ADDRESS
          value: temporal-frontend:7233
        resources:
          requests:
            cpu: 50m
            memory: 64Mi
          limits:
            cpu: 200m
            memory: 256Mi
---
apiVersion: v1
kind: Service
metadata:
  name: temporal-web
  namespace: $NAMESPACE
spec:
  selector:
    app: temporal-ui
  ports:
  - port: 8088
    targetPort: 8080
EOF

    # Wait for PostgreSQL to be ready first
    log_info "Waiting for PostgreSQL..."
    kubectl wait --for=condition=Available deployment/temporal-postgresql \
        -n "$NAMESPACE" --timeout=120s || {
        log_warn "PostgreSQL deployment not ready yet"
    }

    # Wait for Temporal to be ready
    log_info "Waiting for Temporal (this may take a minute)..."
    sleep 10  # Give PostgreSQL time to initialize

    kubectl wait --for=condition=Available deployment/temporal \
        -n "$NAMESPACE" --timeout=300s || {
        log_warn "Temporal not ready yet, continuing..."
        kubectl get pods -n "$NAMESPACE" -l app=temporal
    }

    log_success "Temporal deployed"
}

fetch_credentials_local() {
    log_local "Using local mock credentials..."

    # For local mode, we use test credentials
    # Ensure INSTALL_TOKEN has inst_ prefix for mock cloud compatibility
    if [[ ! "$INSTALL_TOKEN" =~ ^inst_ ]]; then
        INSTALL_TOKEN="inst_${INSTALL_TOKEN}"
        log_info "Prefixed token with 'inst_' for mock cloud compatibility"
    fi

    AGENT_ID="klyro_local_test"
    ACCESS_TOKEN="kat_local_test_token_${RANDOM}"
    CHART_VERSION="local"
    ECR_TOKEN=""
    DATADOG_API_KEY=""

    # OAuth credentials for local testing
    OAUTH2_CLIENT_ID="klyro-edge-local-test"
    OAUTH2_CLIENT_SECRET="local_test_secret_${RANDOM}"

    echo ""
    echo "    ┌─────────────────────────────────────────────────────┐"
    echo "    │  Install Token: $INSTALL_TOKEN"
    echo "    │  Agent ID:      $AGENT_ID"
    echo "    │  Mode:          LOCAL (mock cloud)"
    echo "    │  gRPC:          mock-cloud:50051"
    echo "    │  Temporal:      temporal-frontend:7233"
    echo "    └─────────────────────────────────────────────────────┘"
}

create_local_secrets() {
    log_local "Creating local test secrets..."

    # Create dummy OpenAI credentials (required by worker)
    kubectl create secret generic openai-credentials \
        --from-literal=api-key="sk-local-test-key-not-real" \
        --namespace="$NAMESPACE" \
        --dry-run=client -o yaml | kubectl apply -f - >/dev/null

    printf "    ${GREEN}[OK]${NC} Secret: openai-credentials (dummy for local testing)\n"
}

install_helm_chart_local() {
    log_local "Installing Helm chart in local mode..."

    local CHART_PATH="${PROJECT_ROOT}/charts/klyro-agent"

    if [ ! -d "$CHART_PATH" ]; then
        log_error "Chart not found at: $CHART_PATH"
        exit 1
    fi

    # Determine gRPC endpoint and TLS settings
    local grpc_endpoint="mock-cloud:50051"
    local grpc_insecure="true"

    if [ -n "$CUSTOM_GRPC_ENDPOINT" ]; then
        grpc_endpoint="$CUSTOM_GRPC_ENDPOINT"
        # If endpoint contains standard TLS ports or https, use secure connection
        if [[ "$grpc_endpoint" == *":443"* ]] || [[ "$grpc_endpoint" == "grpc."* ]]; then
            grpc_insecure="false"
        fi
        log_info "Using custom gRPC endpoint: $grpc_endpoint (insecure=$grpc_insecure)"
    fi

    # Determine API endpoint
    local api_endpoint="https://api.klyro.security"
    if [ -n "$CUSTOM_API_ENDPOINT" ]; then
        api_endpoint="$CUSTOM_API_ENDPOINT"
        log_info "Using custom API endpoint: $api_endpoint"
    fi

    # Build Helm arguments for local mode
    # NOTE: We use full image paths for our images instead of global.imageRegistry
    # This prevents the registry override from affecting subchart images (postgresql, etc.)
    HELM_ARGS="--namespace $NAMESPACE"
    HELM_ARGS="$HELM_ARGS --set global.imagePullSecrets=null"

    # Agent configuration - use full image path to local registry
    HELM_ARGS="$HELM_ARGS --set agent.image.repository=klyro-registry:5000/klyro-agent"
    HELM_ARGS="$HELM_ARGS --set agent.image.tag=local"
    HELM_ARGS="$HELM_ARGS --set agent.image.pullPolicy=Always"
    HELM_ARGS="$HELM_ARGS --set agent.persistence.enabled=true"
    HELM_ARGS="$HELM_ARGS --set agent.persistence.storageClass=local-path"

    # Klyro Cloud connection configuration
    HELM_ARGS="$HELM_ARGS --set klyroCloud.endpoint=$api_endpoint"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.grpc.endpoint=$grpc_endpoint"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.grpc.insecure=$grpc_insecure"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.apiKeySecretName=klyro-cloud-credentials"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.apiKeySecretKey=access-token"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.agentIdSecretKey=agent-id"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.heartbeat.enabled=true"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.heartbeat.intervalSeconds=10"

    # Worker/Operator configuration - use full image path to local registry
    HELM_ARGS="$HELM_ARGS --set worker.enabled=true"
    HELM_ARGS="$HELM_ARGS --set worker.image.repository=klyro-registry:5000/klyro-operator"
    HELM_ARGS="$HELM_ARGS --set worker.image.tag=local"
    HELM_ARGS="$HELM_ARGS --set worker.image.pullPolicy=Always"
    HELM_ARGS="$HELM_ARGS --set worker.temporal.hostPort=temporal-frontend:7233"
    HELM_ARGS="$HELM_ARGS --set worker.temporal.namespace=default"
    HELM_ARGS="$HELM_ARGS --set worker.temporal.taskQueue=klyro-tasks"

    # Disable ALL subcharts - we deploy Temporal separately in local mode
    HELM_ARGS="$HELM_ARGS --set temporal.enabled=false"
    HELM_ARGS="$HELM_ARGS --set temporal.mysql.enabled=false"
    HELM_ARGS="$HELM_ARGS --set temporal-mysql.enabled=false"

    # OpenTelemetry - disable for local
    HELM_ARGS="$HELM_ARGS --set opentelemetry.enabled=false"
    HELM_ARGS="$HELM_ARGS --set opentelemetry.clusterName=$AGENT_ID"
    HELM_ARGS="$HELM_ARGS --set opentelemetry.collector.enabled=false"

    # ECR token refresher not needed for local
    HELM_ARGS="$HELM_ARGS --set ecr.enabled=false"
    HELM_ARGS="$HELM_ARGS --set imagePullSecrets.create=false"

    HELM_ARGS="$HELM_ARGS --wait --timeout 5m"

    log_info "Running Helm install..."
    helm upgrade --install "$RELEASE_NAME" "$CHART_PATH" $HELM_ARGS || {
        log_error "Helm installation failed"
        echo ""
        echo "Debug commands:"
        echo "  kubectl get pods -n $NAMESPACE"
        echo "  kubectl describe pods -n $NAMESPACE"
        exit 1
    }

    log_success "Helm chart installed"
}

verify_local_installation() {
    log_local "Verifying local installation..."

    echo ""

    # Check Mock Cloud
    printf "    Checking Mock Cloud... "
    if kubectl get pods -n "$NAMESPACE" -l app=mock-cloud 2>/dev/null | grep -q "Running"; then
        echo -e "${GREEN}OK${NC}"
    else
        echo -e "${YELLOW}PENDING${NC}"
    fi

    # Check Temporal
    printf "    Checking Temporal... "
    if kubectl get pods -n "$NAMESPACE" -l app.kubernetes.io/component=frontend 2>/dev/null | grep -q "Running"; then
        echo -e "${GREEN}OK${NC}"
    else
        echo -e "${YELLOW}PENDING${NC}"
    fi

    # Check Agent
    printf "    Checking Agent... "
    if kubectl get pods -n "$NAMESPACE" -l app.kubernetes.io/component=agent 2>/dev/null | grep -q "Running"; then
        echo -e "${GREEN}OK${NC}"
    else
        echo -e "${YELLOW}PENDING${NC}"
    fi

    # Check Worker
    printf "    Checking Worker... "
    if kubectl get pods -n "$NAMESPACE" -l app.kubernetes.io/component=worker 2>/dev/null | grep -q "Running"; then
        echo -e "${GREEN}OK${NC}"
    else
        echo -e "${YELLOW}PENDING${NC}"
    fi

    echo ""

    # Wait for agent deployment
    log_info "Waiting for agent to be ready..."
    kubectl rollout status deployment/${RELEASE_NAME}-klyro-agent-agent \
        -n "$NAMESPACE" --timeout=120s 2>/dev/null || {
        log_warn "Agent deployment taking longer than expected"
    }

    # Test agent health
    sleep 3
    local POD_NAME=$(kubectl get pods -n "$NAMESPACE" \
        -l app.kubernetes.io/component=agent \
        -o jsonpath='{.items[0].metadata.name}' 2>/dev/null)

    if [ -n "$POD_NAME" ]; then
        log_info "Testing agent health..."
        local HEALTH=$(kubectl exec -n "$NAMESPACE" "$POD_NAME" -- \
            wget -qO- http://localhost:8080/health 2>/dev/null || echo "")

        if echo "$HEALTH" | grep -qi "ok"; then
            printf "    ${GREEN}[OK]${NC} Agent health check passed\n"
        else
            printf "    ${YELLOW}[WARN]${NC} Agent health check inconclusive\n"
        fi

        # Check agent logs for gRPC connection
        log_info "Checking agent connectivity..."
        local LOGS=$(kubectl logs -n "$NAMESPACE" "$POD_NAME" --tail=50 2>/dev/null)

        if echo "$LOGS" | grep -qi "connected.*grpc\|heartbeat.*started\|control stream"; then
            printf "    ${GREEN}[OK]${NC} Agent connected to gRPC cloud\n"
        else
            printf "    ${YELLOW}[INFO]${NC} Agent gRPC connection pending\n"
        fi

        if echo "$LOGS" | grep -qi "temporal\|connector.*initialized"; then
            printf "    ${GREEN}[OK]${NC} Agent connected to Temporal\n"
        else
            printf "    ${YELLOW}[INFO]${NC} Agent Temporal connection pending\n"
        fi
    fi

    # Show pod status
    echo ""
    log_info "Pod status:"
    kubectl get pods -n "$NAMESPACE" -o wide 2>/dev/null | sed 's/^/    /'
}

run_connectivity_tests() {
    log_local "Running connectivity tests..."

    echo ""

    local AGENT_POD=$(kubectl get pods -n "$NAMESPACE" \
        -l app.kubernetes.io/component=agent \
        -o jsonpath='{.items[0].metadata.name}' 2>/dev/null)

    if [ -z "$AGENT_POD" ]; then
        log_warn "Agent pod not found, skipping tests"
        return 1
    fi

    # Test 1: Agent Health
    printf "    Test 1: Agent Health... "
    if kubectl exec -n "$NAMESPACE" "$AGENT_POD" -- wget -qO- http://localhost:8080/health 2>/dev/null | grep -qi "ok"; then
        echo -e "${GREEN}PASS${NC}"
    else
        echo -e "${RED}FAIL${NC}"
    fi

    # Test 2: Agent Status API
    printf "    Test 2: Agent Status API... "
    if kubectl exec -n "$NAMESPACE" "$AGENT_POD" -- wget -qO- http://localhost:8080/api/v1/status 2>/dev/null | grep -qi "agent"; then
        echo -e "${GREEN}PASS${NC}"
    else
        echo -e "${RED}FAIL${NC}"
    fi

    # Test 3: SQLite Database
    printf "    Test 3: SQLite Database... "
    if kubectl exec -n "$NAMESPACE" "$AGENT_POD" -- ls /data/klyro.db 2>/dev/null; then
        echo -e "${GREEN}PASS${NC}"
    else
        echo -e "${YELLOW}PENDING${NC} (DB will be created on first write)"
    fi

    # Test 4: Temporal Frontend Reachability
    printf "    Test 4: Temporal Frontend... "
    if kubectl exec -n "$NAMESPACE" "$AGENT_POD" -- \
        wget -qO- --timeout=5 http://temporal-frontend:7233 2>&1 | grep -qi "error\|refused"; then
        echo -e "${YELLOW}PENDING${NC}"
    else
        echo -e "${GREEN}REACHABLE${NC}"
    fi

    # Test 5: Mock Cloud gRPC Reachability
    printf "    Test 5: Mock Cloud gRPC... "
    if kubectl exec -n "$NAMESPACE" "$AGENT_POD" -- \
        wget -qO- --timeout=5 http://mock-cloud:50052/agents 2>/dev/null; then
        echo -e "${GREEN}REACHABLE${NC}"
    else
        echo -e "${YELLOW}PENDING${NC}"
    fi

    echo ""
    log_success "Connectivity tests completed"
}

# =============================================================================
# Prerequisites Check
# =============================================================================

check_command() {
    if ! command -v "$1" >/dev/null 2>&1; then
        log_error "$1 not found"
        echo "  Install: $2"
        return 1
    fi
    return 0
}

check_prerequisites() {
    log_info "Checking prerequisites..."

    MISSING=0

    check_command "curl" "https://curl.se/download.html" || MISSING=1
    check_command "kubectl" "https://kubernetes.io/docs/tasks/tools/" || MISSING=1
    check_command "helm" "https://helm.sh/docs/intro/install/" || MISSING=1

    if [ "$MISSING" -eq 1 ]; then
        echo ""
        log_error "Missing required tools. Please install them and try again."
        exit 1
    fi
    log_success "Required tools installed"

    # Check Kubernetes connectivity
    log_info "Checking Kubernetes connectivity..."
    if ! kubectl cluster-info >/dev/null 2>&1; then
        log_error "Cannot connect to Kubernetes cluster"
        echo ""
        echo "Please ensure:"
        echo "  - You have a valid kubeconfig (~/.kube/config)"
        echo "  - Your cluster is running and accessible"
        echo "  - Run: kubectl cluster-info"
        exit 1
    fi
    log_success "Kubernetes cluster accessible"

    # Check RBAC permissions
    log_info "Checking RBAC permissions..."
    if ! kubectl auth can-i create namespace >/dev/null 2>&1; then
        log_error "Insufficient permissions to create namespace"
        echo ""
        echo "The installer needs cluster-admin or equivalent permissions."
        exit 1
    fi

    if ! kubectl auth can-i create secret --namespace="$NAMESPACE" >/dev/null 2>&1; then
        log_error "Insufficient permissions to create secrets"
        exit 1
    fi
    log_success "RBAC permissions verified"

    # Check Helm version
    log_info "Checking Helm version..."
    HELM_VERSION=$(helm version --short 2>/dev/null | grep -oE 'v[0-9]+\.[0-9]+' | head -1)
    HELM_MAJOR=$(echo "$HELM_VERSION" | cut -d. -f1 | tr -d 'v')
    if [ "$HELM_MAJOR" -lt 3 ]; then
        log_error "Helm 3.x required (found: $HELM_VERSION)"
        exit 1
    fi
    log_success "Helm version: $HELM_VERSION"
}

# =============================================================================
# Token Validation & Credential Fetch
# =============================================================================
#
# API Contract:
#   GET ${KLYRO_API_URL}/v1/agents/install?token=<INSTALL_TOKEN>
#
#   Response (200 OK):
#   {
#     "agent_id": "{tenant_id}_{site_id}",
#     "access_token": "kat_...",
#     "chart_version": "0.2.0",
#     "ecr_token": "eyJ...",
#     "datadog_api_key": "dd_..."          # Optional: Klyro's DD API key
#   }
#
#   agent_id format: {tenant_id}_{site_id}
#   Examples: "acme_corp_us-east-prod", "globex_hq", "initech_dc1"
#   NOTE: agent_id is used as opentelemetry.clusterName to identify this
#         tenant in Datadog dashboards (all tenants share Klyro's DD account)
#
#   access_token: Long-lived token for gRPC auth (valid for agent lifecycle)
#   ecr_token: Docker auth token for ECR image pulls
#   datadog_api_key: Klyro's Datadog API key for centralized telemetry
# =============================================================================

# Hardcoded infrastructure values
# TODO: Re-enable ECR when ready for production
# ECR_REGISTRY="123456789.dkr.ecr.us-east-1.amazonaws.com"
LOCAL_REGISTRY="localhost:5000"
LOCAL_IMAGE="localhost:5000/klyro-agent:dev"
GRPC_ENDPOINT="grpc.klyro.security:443"

# Simple JSON value extractor (no jq dependency)
json_value() {
    echo "$1" | grep -o "\"$2\"[[:space:]]*:[[:space:]]*\"[^\"]*\"" | sed "s/\"$2\"[[:space:]]*:[[:space:]]*\"//" | sed 's/"$//'
}

fetch_credentials() {
    INSTALL_TOKEN="$1"

    RESPONSE_FILE=$(mktemp)
    # Clean up temp file on exit
    trap "rm -f $RESPONSE_FILE" EXIT

    HTTP_CODE=$(curl -sfL -w "%{http_code}" \
        -o "$RESPONSE_FILE" \
        "${KLYRO_API_URL}/api/v1/agents/install?token=${INSTALL_TOKEN}" \
        -H "User-Agent: klyro-installer/1.0" \
        2>/dev/null) || HTTP_CODE="000"

    case "$HTTP_CODE" in
        200|201)
            log_success "Install token validated"
            ;;
        400)
            log_error "Invalid install token format"
            echo "Please copy the complete token from: https://app.klyro.security/agents/new"
            exit 1
            ;;
        401|403)
            log_error "Install token expired or invalid"
            echo "Generate a new token from: https://app.klyro.security/agents/new"
            exit 1
            ;;
        404)
            log_error "Install token not found"
            exit 1
            ;;
        409)
            log_error "Install token has already been used"
            echo "Generate a new token from: https://app.klyro.security/agents/new"
            exit 1
            ;;
        429)
            log_error "Rate limited - please wait and try again"
            exit 1
            ;;
        000)
            log_error "Cannot connect to Klyro Cloud API"
            echo "Endpoint: ${KLYRO_API_URL}"
            exit 1
            ;;
        *)
            log_error "API error (HTTP $HTTP_CODE)"
            [ "${KLYRO_DEBUG:-}" = "true" ] && cat "$RESPONSE_FILE"
            exit 1
            ;;
    esac

    RESPONSE=$(cat "$RESPONSE_FILE")

    # Extract required fields from API response
    AGENT_ID=$(json_value "$RESPONSE" "agent_id")
    TENANT_ID=$(json_value "$RESPONSE" "tenant_id")
    SITE_ID=$(json_value "$RESPONSE" "site_id")
    CHART_VERSION=$(json_value "$RESPONSE" "chart_version")
    OAUTH2_CLIENT_ID=$(json_value "$RESPONSE" "oauth2_client_id")
    OAUTH2_CLIENT_SECRET=$(json_value "$RESPONSE" "oauth2_client_secret")
    GRPC_ENDPOINT=$(json_value "$RESPONSE" "grpc_endpoint")
    AUTH_ENDPOINT=$(json_value "$RESPONSE" "auth_endpoint")

    # Extract optional fields
    ECR_TOKEN=$(json_value "$RESPONSE" "ecr_token")
    DATADOG_API_KEY=$(json_value "$RESPONSE" "datadog_api_key")

    # Validate required fields
    if [ -z "$AGENT_ID" ] || [ -z "$OAUTH2_CLIENT_ID" ] || [ -z "$OAUTH2_CLIENT_SECRET" ]; then
        log_error "Incomplete response from API - missing required credentials"
        [ "${KLYRO_DEBUG:-}" = "true" ] && echo "Response: $RESPONSE"
        exit 1
    fi

    echo ""
    echo "    ┌─────────────────────────────────────────────────────┐"
    echo "    │  Agent ID:      $AGENT_ID"
    echo "    │  Tenant ID:     $TENANT_ID"
    echo "    │  Site ID:       $SITE_ID"
    echo "    │  Chart Version: $CHART_VERSION"
    echo "    │  Telemetry:     $([ -n "$DATADOG_API_KEY" ] && echo "Datadog enabled" || echo "Disabled")"
    echo "    └─────────────────────────────────────────────────────┘"
}

# =============================================================================
# Cleanup Existing Installation
# =============================================================================

cleanup_existing_installation() {
    # Check if Helm release exists
    if helm status "$RELEASE_NAME" -n "$NAMESPACE" >/dev/null 2>&1; then
        log_warn "Found existing Helm release: $RELEASE_NAME"

        # Uninstall with spinner
        helm uninstall "$RELEASE_NAME" -n "$NAMESPACE" --wait >/dev/null 2>&1 &
        UNINSTALL_PID=$!

        local spinstr='⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'
        local i=0
        printf "    ${BLUE}[%s]${NC} Uninstalling existing release..." "$(echo $spinstr | cut -c1)"
        while kill -0 "$UNINSTALL_PID" 2>/dev/null; do
            i=$(( (i + 1) % 10 ))
            printf "\r    ${BLUE}[%s]${NC} Uninstalling existing release..." "$(echo $spinstr | cut -c$((i+1)))"
            sleep 0.1
        done
        wait "$UNINSTALL_PID"
        UNINSTALL_STATUS=$?

        if [ $UNINSTALL_STATUS -eq 0 ]; then
            printf "\r    ${GREEN}[OK]${NC} Existing Helm release uninstalled    \n"
        else
            printf "\r    ${RED}[FAIL]${NC} Failed to uninstall existing release\n"
            exit 1
        fi
    else
        printf "    ${GREEN}[OK]${NC} No existing Helm release found\n"
    fi

    # Check if namespace exists and has resources
    if kubectl get namespace "$NAMESPACE" >/dev/null 2>&1; then
        log_info "Cleaning up namespace resources..."

        # Delete secrets
        kubectl delete secret klyro-cloud-credentials -n "$NAMESPACE" --ignore-not-found >/dev/null 2>&1
        kubectl delete secret klyro-registry-secret -n "$NAMESPACE" --ignore-not-found >/dev/null 2>&1
        kubectl delete secret datadog-secrets -n "$NAMESPACE" --ignore-not-found >/dev/null 2>&1
        printf "    ${GREEN}[OK]${NC} Secrets cleaned up\n"

        # Delete any remaining pods
        kubectl delete pods -n "$NAMESPACE" --all --ignore-not-found >/dev/null 2>&1

        # Wait for pods to terminate with spinner
        kubectl wait --for=delete pod --all -n "$NAMESPACE" --timeout=60s >/dev/null 2>&1 &
        WAIT_PID=$!

        local spinstr='⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'
        local i=0
        printf "    ${BLUE}[%s]${NC} Waiting for pods to terminate..." "$(echo $spinstr | cut -c1)"
        while kill -0 "$WAIT_PID" 2>/dev/null; do
            i=$(( (i + 1) % 10 ))
            printf "\r    ${BLUE}[%s]${NC} Waiting for pods to terminate..." "$(echo $spinstr | cut -c$((i+1)))"
            sleep 0.1
        done
        wait "$WAIT_PID" || true
        printf "\r    ${GREEN}[OK]${NC} Pods terminated                  \n"
    else
        printf "    ${GREEN}[OK]${NC} No existing namespace found\n"
    fi
}

# =============================================================================
# Kubernetes Resources Creation
# =============================================================================

create_namespace() {
    kubectl create namespace "$NAMESPACE" --dry-run=client -o yaml | kubectl apply -f - >/dev/null
    printf "    ${GREEN}[OK]${NC} Namespace: $NAMESPACE\n"
}

# TODO: Re-enable when ECR is ready for production
# create_registry_secret() {
#     log_info "Creating ECR registry secret..."
#
#     # Create docker-registry secret for image pulls
#     kubectl create secret docker-registry klyro-registry-secret \
#         --docker-server="$ECR_REGISTRY" \
#         --docker-username="AWS" \
#         --docker-password="$ECR_TOKEN" \
#         --namespace="$NAMESPACE" \
#         --dry-run=client -o yaml | kubectl apply -f - >/dev/null
#
#     log_success "Registry secret created"
# }


create_cloud_credentials_secret() {
    kubectl create secret generic klyro-cloud-credentials \
        --from-literal=agent-id="$AGENT_ID" \
        --from-literal=access-token="$ACCESS_TOKEN" \
        --from-literal=install-token="$INSTALL_TOKEN" \
        --namespace="$NAMESPACE" \
        --dry-run=client -o yaml | kubectl apply -f - >/dev/null

    printf "    ${GREEN}[OK]${NC} Secret: klyro-cloud-credentials\n"
}

create_oauth_credentials_secret() {
    # Create OAuth secret for JWT-based authentication
    # Platform API returns these during agent installation
    kubectl create secret generic klyro-edge-oauth-credentials \
        --from-literal=client-id="$OAUTH2_CLIENT_ID" \
        --from-literal=client-secret="$OAUTH2_CLIENT_SECRET" \
        --namespace="$NAMESPACE" \
        --dry-run=client -o yaml | kubectl apply -f - >/dev/null

    printf "    ${GREEN}[OK]${NC} Secret: klyro-edge-oauth-credentials\n"
}

create_datadog_secret() {
    if [ -z "$DATADOG_API_KEY" ]; then
        printf "    ${BLUE}[SKIP]${NC} Datadog telemetry (no API key)\n"
        return 0
    fi

    kubectl create secret generic datadog-secrets \
        --from-literal=api-key="$DATADOG_API_KEY" \
        --namespace="$NAMESPACE" \
        --dry-run=client -o yaml | kubectl apply -f - >/dev/null

    printf "    ${GREEN}[OK]${NC} Secret: datadog-secrets\n"
}

create_openai_secret() {
    # OpenAI API key is MANDATORY for worker/operator to function
    if [ -z "$OPENAI_API_KEY" ]; then
        log_error "OpenAI API key is required for the worker to function"
        echo ""
        echo "Please provide the OpenAI API key via environment variable:"
        echo "  export OPENAI_API_KEY=sk-..."
        echo "  ./install.sh <TOKEN>"
        echo ""
        echo "Or set it inline:"
        echo "  OPENAI_API_KEY=sk-... ./install.sh <TOKEN>"
        exit 1
    fi

    kubectl create secret generic openai-credentials \
        --from-literal=api-key="$OPENAI_API_KEY" \
        --namespace="$NAMESPACE" \
        --dry-run=client -o yaml | kubectl apply -f - >/dev/null

    printf "    ${GREEN}[OK]${NC} Secret: openai-credentials\n"
}

# =============================================================================
# Helm Installation
# =============================================================================

install_helm_chart() {
    # TODO: Re-enable ECR authentication when ready for production
    # # Login to ECR for Helm chart pull
    # log_info "Authenticating Helm with ECR..."
    # echo "$ECR_TOKEN" | helm registry login "$ECR_REGISTRY" \
    #     --username AWS \
    #     --password-stdin >/dev/null 2>&1
    #
    # log_success "Helm authenticated with ECR"

    # Get the script directory to find the local chart
    SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
    CHART_PATH="${SCRIPT_DIR}/../charts/klyro-agent"

    if [ ! -d "$CHART_PATH" ]; then
        log_error "Local chart not found at: $CHART_PATH"
        exit 1
    fi
    log_info "Using local chart: $CHART_PATH"

    # Build Helm arguments - production mode configuration
    # Should match local mode behavior but pointing to real cloud infrastructure
    HELM_ARGS="--namespace $NAMESPACE"
    HELM_ARGS="$HELM_ARGS --set global.imagePullSecrets=null"

    # Agent configuration
    HELM_ARGS="$HELM_ARGS --set agent.image.repository=${LOCAL_REGISTRY}/klyro-agent"
    HELM_ARGS="$HELM_ARGS --set agent.image.tag=dev"
    HELM_ARGS="$HELM_ARGS --set agent.image.pullPolicy=Always"
    HELM_ARGS="$HELM_ARGS --set agent.persistence.enabled=true"

    # Klyro Cloud gRPC configuration
    HELM_ARGS="$HELM_ARGS --set klyroCloud.grpc.endpoint=$GRPC_ENDPOINT"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.apiKeySecretName=klyro-cloud-credentials"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.apiKeySecretKey=access-token"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.agentIdSecretKey=agent-id"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.heartbeat.enabled=true"
    HELM_ARGS="$HELM_ARGS --set klyroCloud.heartbeat.intervalSeconds=30"

    # Worker/Operator configuration (enabled by default, uses openai-credentials secret)
    HELM_ARGS="$HELM_ARGS --set worker.enabled=true"
    HELM_ARGS="$HELM_ARGS --set worker.image.repository=${LOCAL_REGISTRY}/klyro-operator"
    HELM_ARGS="$HELM_ARGS --set worker.image.tag=dev"
    HELM_ARGS="$HELM_ARGS --set worker.image.pullPolicy=Always"
    # Temporal settings use defaults from values.yaml (temporal-frontend:7233, default namespace, klyro-tasks queue)

    # Disable Temporal subchart - customer provides their own Temporal or uses hosted
    HELM_ARGS="$HELM_ARGS --set temporal.enabled=false"

    # OpenTelemetry configuration
    HELM_ARGS="$HELM_ARGS --set opentelemetry.clusterName=$AGENT_ID"
    if [ -n "$DATADOG_API_KEY" ]; then
        HELM_ARGS="$HELM_ARGS --set opentelemetry.enabled=true"
        HELM_ARGS="$HELM_ARGS --set opentelemetry.collector.enabled=true"
        HELM_ARGS="$HELM_ARGS --set opentelemetry.collector.exporters.datadog.enabled=true"
    else
        HELM_ARGS="$HELM_ARGS --set opentelemetry.enabled=false"
        HELM_ARGS="$HELM_ARGS --set opentelemetry.collector.enabled=false"
    fi

    # ECR not needed for local dev images
    HELM_ARGS="$HELM_ARGS --set ecr.enabled=false"
    HELM_ARGS="$HELM_ARGS --set imagePullSecrets.create=false"

    HELM_ARGS="$HELM_ARGS --wait --timeout 10m"

    # TODO: Re-enable OCI chart pull when ECR is ready
    # helm upgrade --install "$RELEASE_NAME" \
    #     "oci://${ECR_REGISTRY}/charts/klyro-agent" \
    #     $HELM_ARGS

    # Create temp file for helm output
    HELM_OUTPUT=$(mktemp)
    trap "rm -f $HELM_OUTPUT" EXIT

    log_info "Running Helm upgrade/install (this may take a few minutes)..."
    echo ""

    # Run helm with spinner
    if [ "${KLYRO_DEBUG:-}" = "true" ]; then
        helm upgrade --install "$RELEASE_NAME" "$CHART_PATH" $HELM_ARGS 2>&1 | tee "$HELM_OUTPUT"
        HELM_STATUS=${PIPESTATUS[0]}
    else
        helm upgrade --install "$RELEASE_NAME" "$CHART_PATH" $HELM_ARGS >"$HELM_OUTPUT" 2>&1 &
        HELM_PID=$!

        # Show progress while helm runs
        local spinstr='⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'
        local i=0
        printf "    ${BLUE}[%s]${NC} Deploying Klyro Edge..." "$(echo $spinstr | cut -c1)"
        while kill -0 "$HELM_PID" 2>/dev/null; do
            i=$(( (i + 1) % 10 ))
            printf "\r    ${BLUE}[%s]${NC} Deploying Klyro Edge..." "$(echo $spinstr | cut -c$((i+1)))"
            sleep 0.1
        done
        wait "$HELM_PID"
        HELM_STATUS=$?

        if [ $HELM_STATUS -eq 0 ]; then
            printf "\r    ${GREEN}[OK]${NC} Deploying Klyro Edge      \n"
        else
            printf "\r    ${RED}[FAIL]${NC} Deploying Klyro Edge      \n"
        fi
    fi

    if [ $HELM_STATUS -ne 0 ]; then
        log_error "Helm installation failed"
        echo ""
        echo "Helm output:"
        cat "$HELM_OUTPUT"
        echo ""
        echo "Debug commands:"
        echo "  kubectl get pods -n $NAMESPACE"
        echo "  kubectl describe pods -n $NAMESPACE"
        echo "  kubectl logs -n $NAMESPACE -l app.kubernetes.io/name=klyro-agent"
        exit 1
    fi

    log_success "Helm chart installed successfully"
}

# =============================================================================
# Verification
# =============================================================================

verify_installation() {
    echo ""

    # Wait for agent deployment with spinner
    kubectl rollout status deployment/${RELEASE_NAME}-klyro-agent-agent \
        -n "$NAMESPACE" --timeout=120s >/dev/null 2>&1 &
    ROLLOUT_PID=$!

    local spinstr='⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'
    local i=0
    printf "    ${BLUE}[%s]${NC} Waiting for agent deployment..." "$(echo $spinstr | cut -c1)"
    while kill -0 "$ROLLOUT_PID" 2>/dev/null; do
        i=$(( (i + 1) % 10 ))
        printf "\r    ${BLUE}[%s]${NC} Waiting for agent deployment..." "$(echo $spinstr | cut -c$((i+1)))"
        sleep 0.1
    done
    wait "$ROLLOUT_PID"
    ROLLOUT_STATUS=$?

    if [ $ROLLOUT_STATUS -eq 0 ]; then
        printf "\r    ${GREEN}[OK]${NC} Agent deployment ready        \n"
    else
        printf "\r    ${YELLOW}[WARN]${NC} Agent deployment taking longer than expected\n"
        log_info "You can check status with: kubectl get pods -n $NAMESPACE"
    fi

    # Quick health check
    sleep 2

    POD_NAME=$(kubectl get pods -n "$NAMESPACE" \
        -l app.kubernetes.io/component=agent \
        -o jsonpath='{.items[0].metadata.name}' 2>/dev/null)

    if [ -n "$POD_NAME" ]; then
        HEALTH=$(kubectl exec -n "$NAMESPACE" "$POD_NAME" -- \
            wget -qO- http://localhost:8080/health 2>/dev/null || echo "")

        if echo "$HEALTH" | grep -qi "ok"; then
            printf "    ${GREEN}[OK]${NC} Health check passed\n"
        else
            printf "    ${YELLOW}[WARN]${NC} Health check inconclusive (agent may still be starting)\n"
        fi
    else
        printf "    ${YELLOW}[WARN]${NC} Could not find agent pod for health check\n"
    fi

    # Show pod status
    echo ""
    log_info "Pod status:"
    kubectl get pods -n "$NAMESPACE" -o wide 2>/dev/null | sed 's/^/    /'
}

# =============================================================================
# Main
# =============================================================================

main_local() {
    show_banner
    printf "${CYAN}════════════════════════════════════════════════════════════════════${NC}\n"
    printf "${CYAN}                     LOCAL MODE INSTALLATION                        ${NC}\n"
    printf "${CYAN}════════════════════════════════════════════════════════════════════${NC}\n"
    echo ""

    # Determine if using custom endpoints (skip mock cloud)
    local use_custom_endpoints=false
    if [ -n "$CUSTOM_GRPC_ENDPOINT" ]; then
        use_custom_endpoints=true
        log_info "Using custom gRPC endpoint: $CUSTOM_GRPC_ENDPOINT"
        if [ -n "$CUSTOM_API_ENDPOINT" ]; then
            log_info "Using custom API endpoint: $CUSTOM_API_ENDPOINT"
        fi
        echo ""
    fi

    # Determine number of steps based on mode
    local total_steps=8
    if [ "$use_custom_endpoints" = "true" ]; then
        total_steps=7  # Skip mock cloud step
    fi

    local step=1

    # Step 1: Prerequisites
    echo ""
    printf "${CYAN}Step ${step}/${total_steps}${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Checking local prerequisites"
    check_local_prerequisites
    ((step++))

    # Step 2: k3d cluster
    echo ""
    printf "${CYAN}Step ${step}/${total_steps}${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Setting up k3d cluster"
    create_k3d_cluster
    ((step++))

    # Step 3: Build images
    echo ""
    printf "${CYAN}Step ${step}/${total_steps}${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Building Docker images"
    build_local_images
    ((step++))

    # Step 4: Create namespace
    echo ""
    printf "${CYAN}Step ${step}/${total_steps}${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Creating Kubernetes namespace"
    create_namespace
    ((step++))

    # Step 5: Start Temporal
    echo ""
    printf "${CYAN}Step ${step}/${total_steps}${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Starting Temporal"
    start_temporal
    ((step++))

    # Step 6: Start Mock Cloud (skip if using custom endpoints)
    if [ "$use_custom_endpoints" = "false" ]; then
        echo ""
        printf "${CYAN}Step ${step}/${total_steps}${NC} ─────────────────────────────────────────────────────────\n"
        log_step "Starting Mock Cloud gRPC server"
        start_mock_cloud
        ((step++))
    fi

    # Step 7: Credentials and secrets
    echo ""
    printf "${CYAN}Step ${step}/${total_steps}${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Setting up credentials"
    fetch_credentials_local
    create_cloud_credentials_secret
    create_oauth_credentials_secret
    create_local_secrets
    ((step++))

    # Step 8: Install Helm chart
    echo ""
    printf "${CYAN}Step ${step}/${total_steps}${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Installing Klyro Edge"
    install_helm_chart_local

    # Verification
    echo ""
    printf "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}\n"
    log_step "Verifying installation"
    verify_local_installation

    # Connectivity tests (skip if using custom endpoints)
    if [ "$use_custom_endpoints" = "false" ]; then
        echo ""
        printf "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}\n"
        run_connectivity_tests
    fi

    # Success message
    echo ""
    echo ""
    printf "${GREEN}╔═══════════════════════════════════════════════════════════════════╗${NC}\n"
    printf "${GREEN}║${NC}                                                                   ${GREEN}║${NC}\n"
    printf "${GREEN}║${NC}   ${GREEN}✓ Klyro Edge LOCAL installation complete!${NC}                     ${GREEN}║${NC}\n"
    printf "${GREEN}║${NC}                                                                   ${GREEN}║${NC}\n"
    printf "${GREEN}╚═══════════════════════════════════════════════════════════════════╝${NC}\n"
    echo ""
    echo "  Agent ID:      ${GREEN}$AGENT_ID${NC}"
    echo "  Namespace:     ${GREEN}$NAMESPACE${NC}"
    echo "  Release:       ${GREEN}$RELEASE_NAME${NC}"
    echo "  Mode:          ${CYAN}LOCAL${NC}"
    if [ "$use_custom_endpoints" = "true" ]; then
        echo "  gRPC Endpoint: ${CYAN}$CUSTOM_GRPC_ENDPOINT${NC}"
        [ -n "$CUSTOM_API_ENDPOINT" ] && echo "  API Endpoint:  ${CYAN}$CUSTOM_API_ENDPOINT${NC}"
    fi
    echo ""
    printf "${BLUE}Components running:${NC}\n"
    if [ "$use_custom_endpoints" = "true" ]; then
        echo "  - Agent:       Connected to custom cloud endpoints"
    else
        echo "  - Agent:       Connected to mock cloud & Temporal"
        echo "  - Mock Cloud:  gRPC server for agent registration"
    fi
    echo "  - Worker:      Polling Temporal for tasks"
    echo "  - Temporal:    Workflow orchestration"
    echo ""
    printf "${BLUE}Useful commands:${NC}\n"
    echo "  kubectl get pods -n $NAMESPACE"
    echo "  kubectl logs -n $NAMESPACE -l app.kubernetes.io/component=agent -f"
    if [ "$use_custom_endpoints" = "false" ]; then
        echo "  kubectl logs -n $NAMESPACE -l app=mock-cloud -f"
    fi
    echo "  kubectl port-forward -n $NAMESPACE svc/temporal-web 8088:8088"
    echo ""
    printf "${BLUE}To clean up:${NC}\n"
    echo "  k3d cluster delete klyro-local"
    echo ""
}

main_production() {
    show_banner

    # Parse arguments
    INSTALL_TOKEN="$1"

    if [ -z "$INSTALL_TOKEN" ]; then
        log_error "Missing install token"
        echo ""
        usage
        exit 1
    fi

    # Validate token format
    case "$INSTALL_TOKEN" in
        inst_*)
            ;;
        *)
            log_warn "Token format looks unusual (expected: inst_xxxx...)"
            ;;
    esac

    echo ""
    printf "${BLUE}Step 1/6${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Checking prerequisites"
    check_prerequisites

    echo ""
    printf "${BLUE}Step 2/6${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Cleaning up existing installation"
    cleanup_existing_installation

    echo ""
    printf "${BLUE}Step 3/6${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Fetching credentials"
    fetch_credentials "$INSTALL_TOKEN"

    echo ""
    printf "${BLUE}Step 4/6${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Creating Kubernetes resources"
    create_namespace
    # TODO: Re-enable when ECR is ready
    # create_registry_secret
    create_cloud_credentials_secret
    create_oauth_credentials_secret
    create_openai_secret
    create_datadog_secret

    echo ""
    printf "${BLUE}Step 5/6${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Installing Helm chart"
    install_helm_chart

    echo ""
    printf "${BLUE}Step 6/6${NC} ─────────────────────────────────────────────────────────\n"
    log_step "Verifying installation"
    verify_installation

    echo ""
    echo ""
    printf "${GREEN}╔═══════════════════════════════════════════════════════════════════╗${NC}\n"
    printf "${GREEN}║${NC}                                                                   ${GREEN}║${NC}\n"
    printf "${GREEN}║${NC}   ${GREEN}✓ Klyro Edge installed successfully!${NC}                          ${GREEN}║${NC}\n"
    printf "${GREEN}║${NC}                                                                   ${GREEN}║${NC}\n"
    printf "${GREEN}╚═══════════════════════════════════════════════════════════════════╝${NC}\n"
    echo ""
    echo "  Agent ID:   ${GREEN}$AGENT_ID${NC}"
    echo "  Namespace:  ${GREEN}$NAMESPACE${NC}"
    echo "  Release:    ${GREEN}$RELEASE_NAME${NC}"
    echo ""
    printf "${BLUE}Useful commands:${NC}\n"
    echo "  kubectl get pods -n $NAMESPACE"
    echo "  kubectl logs -n $NAMESPACE -l app.kubernetes.io/component=agent -f"
    echo "  helm status $RELEASE_NAME -n $NAMESPACE"
    echo ""
}

main() {
    # Parse arguments
    is_local=false
    INSTALL_TOKEN=""
    CUSTOM_API_ENDPOINT=""
    CUSTOM_GRPC_ENDPOINT=""

    while [ $# -gt 0 ]; do
        case "$1" in
            --help|-h|help)
                usage
                ;;
            --local|--local-dev)
                is_local=true
                shift
                ;;
            --api-endpoint|--api_endpoint)
                CUSTOM_API_ENDPOINT="$2"
                shift 2
                ;;
            --grpc-endpoint|--grpc_endpoint)
                CUSTOM_GRPC_ENDPOINT="$2"
                shift 2
                ;;
            *)
                if [ -z "$INSTALL_TOKEN" ]; then
                    INSTALL_TOKEN="$1"
                fi
                shift
                ;;
        esac
    done

    if [ -z "$INSTALL_TOKEN" ]; then
        log_error "Missing install token"
        echo ""
        usage
        exit 1
    fi

    # Export custom endpoints for use in other functions
    export CUSTOM_API_ENDPOINT
    export CUSTOM_GRPC_ENDPOINT

    if [ "$is_local" = "true" ]; then
        main_local "$INSTALL_TOKEN"
    else
        main_production "$INSTALL_TOKEN"
    fi
}

main "$@"
