#!/usr/bin/env bash

set -euo pipefail

# init_prole-db.sh
# Purpose:
# - Manage prole-db CloudNative-PG cluster operations (deploy, start, stop, etc.)

# Initialize SCRIPT_DIR
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)

ACTION=${1:-}
VERSION=${2:-latest}

# Load env
if [[ -n "${PROLE_HOME:-}" && -f "$PROLE_HOME/env.sh" ]]; then
  set --
  source "$PROLE_HOME/env.sh"
elif [[ -f "$HOME/.prole/env.sh" ]]; then
  set --
  source "$HOME/.prole/env.sh"
fi

CNPG_CLUSTER_NAME=${CNPG_CLUSTER_NAME:-prole-db}
NAMESPACE=${NAMESPACE:-default}

# Support both PROLE_HOME/k8s and sibling k8s directory
if [[ -d "$SCRIPT_DIR/../k8s/prole" ]]; then
  CNPG_MANIFEST="$SCRIPT_DIR/../k8s/prole/prole-db.yaml"
elif [[ -n "${PROLE_HOME:-}" && -d "$PROLE_HOME/k8s/prole" ]]; then
  CNPG_MANIFEST="$PROLE_HOME/k8s/prole/prole-db.yaml"
else
  CNPG_MANIFEST="$SCRIPT_DIR/../k8s/prole/prole-db.yaml"
fi

ensure_tools() {
  for t in kubectl jq; do
    command -v "$t" >/dev/null || { echo "Missing required tool: $t" >&2; exit 1; }
  done
}

get_latest_image() {
    local version_file
    if [[ -f "$SCRIPT_DIR/../conf/postgresql/.version" ]]; then
        version_file="$SCRIPT_DIR/../conf/postgresql/.version"
    elif [[ -n "${PROLE_HOME:-}" && -f "$PROLE_HOME/conf/postgresql/.version" ]]; then
        version_file="$PROLE_HOME/conf/postgresql/.version"
    else
        version_file="$SCRIPT_DIR/../conf/postgresql/.version"
    fi

    if [[ -f "$version_file" ]]; then
        echo "prole-db:$(cat "$version_file" | tr -d '[:space:]')"
    else
        echo "prole-db:17.7-037"
    fi
}

start() {
    ensure_tools
    echo "Checking dependencies..."
    
    # 1. k3d is running
    if ! command -v k3d >/dev/null || ! k3d cluster list >/dev/null 2>&1; then
        echo "ERROR: k3d is not running or not installed." >&2
        exit 1
    fi
    
    # 2. cnpg operator is loaded
    if ! kubectl get deployment -n cnpg-system cnpg-controller-manager >/dev/null 2>&1; then
        echo "ERROR: CloudNative-PG operator is not loaded." >&2
        exit 1
    fi
    
    # 3. openbao is configured
    if ! kubectl get statefulset openbao -n "$NAMESPACE" >/dev/null 2>&1 && ! kubectl get deployment openbao -n "$NAMESPACE" >/dev/null 2>&1; then
        echo "ERROR: OpenBao is not deployed." >&2
        exit 1
    fi

    # Compare latest image with deployed
    local latest_image
    latest_image=$(get_latest_image)
    
    local current_image
    current_image=$(kubectl get cluster "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" -o jsonpath='{.spec.imageName}' 2>/dev/null || echo "")
    
    if [[ "$current_image" != "$latest_image" ]]; then
        echo "Updating cluster image from '$current_image' to '$latest_image'..."
        kubectl patch cluster "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" --type merge -p "{\"spec\": {\"imageName\": \"$latest_image\"}}"
        # Force a rollout to ensure the new image is pulled even if it was just a tag update (though we use unique tags)
        kubectl cnpg restart "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" || true
    else
        echo "Cluster is already using the latest image: $latest_image"
    fi

    echo "Starting prole-db cluster (ensuring manifest is applied)..."
    kubectl apply -n "$NAMESPACE" -f "$CNPG_MANIFEST"
}

stop() {
    ensure_tools
    echo "Stopping prole-db cluster $CNPG_CLUSTER_NAME..."
    
    # Identify instances
    local instances
    instances=$(kubectl get pods -n "$NAMESPACE" -l "cnpg.io/cluster=$CNPG_CLUSTER_NAME" -o jsonpath='{.items[*].metadata.name}')
    
    if [[ -z "$instances" ]]; then
        echo "No instances found for cluster $CNPG_CLUSTER_NAME."
        return
    fi

    # Determine replicas and primary
    local primary
    primary=$(kubectl get cluster "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" -o jsonpath='{.status.currentPrimary}')
    
    echo "Primary: $primary"
    
    # Shutdown replicas first
    for pod in $instances; do
        if [[ "$pod" != "$primary" ]]; then
            echo "Shutting down replica $pod..."
            kubectl exec -n "$NAMESPACE" "$pod" -c postgres -- psql -U postgres -c "CHECKPOINT;" || true
        fi
    done
    
    # Shutdown primary last
    if [[ -n "$primary" ]]; then
        echo "Shutting down primary $primary..."
        kubectl exec -n "$NAMESPACE" "$primary" -c postgres -- psql -U postgres -c "CHECKPOINT;" || true
    fi

    echo "Deleting cluster resource..."
    kubectl delete -f "$CNPG_MANIFEST" --ignore-not-found
}

restart() {
    ensure_tools
    echo "Restarting prole-db cluster..."
    kubectl cnpg restart "$CNPG_CLUSTER_NAME" -n "$NAMESPACE"
}

deploy() {
    ensure_tools
    local image
    if [[ "$VERSION" == "latest" ]]; then
        image=$(get_latest_image)
    else
        image="prole-db:$VERSION"
    fi
    
    echo "Deploying $image to cluster $CNPG_CLUSTER_NAME..."
    
    # If cluster doesn't exist, use init_cloudnative_pg.sh create first
    if ! kubectl get cluster "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" >/dev/null 2>&1; then
        echo "Cluster not found. Running init_cloudnative_pg.sh create..."
        bash "$SCRIPT_DIR/init_cloudnative_pg.sh" create
    fi
    
    # Ensure image is updated if manifest has an older version
    # First, apply the manifest to ensure the cluster exists/is updated
    echo "Applying manifest $CNPG_MANIFEST..."
    kubectl apply -n "$NAMESPACE" -f "$CNPG_MANIFEST"

    # Then force the specific image version via patch if different
    local current_image
    current_image=$(kubectl get cluster "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" -o jsonpath='{.spec.imageName}' 2>/dev/null || echo "")
    if [[ "$current_image" != "$image" ]]; then
        echo "Patching cluster to use image '$image' (was '$current_image')..."
        kubectl patch cluster "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" --type merge -p "{\"spec\": {\"imageName\": \"$image\"}}"
        # Force a rollout
        kubectl cnpg restart "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" || true
    fi
}

rollback() {
    ensure_tools
    echo "Rollback: Updating to previous image (manually specified version or fallback)..."
    if [[ "$VERSION" == "latest" ]]; then
        echo "Please specify a version to rollback to. Usage: $0 rollback <version>"
        exit 1
    fi
    deploy
}

rollout() {
    ensure_tools
    echo "Starting manual recreate rollout for $CNPG_CLUSTER_NAME in $NAMESPACE..."
    
    # 1. Monitor status
    echo "Current status:"
    kubectl cnpg status "$CNPG_CLUSTER_NAME" -n "$NAMESPACE"
    
    # 2. Identify the primary instance
    local primary
    primary=$(kubectl get cluster "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" -o jsonpath='{.status.currentPrimary}')
    if [[ -z "$primary" ]]; then
        echo "ERROR: Could not identify primary instance." >&2
        exit 1
    fi
    echo "Primary instance: $primary"
    
    # 3. Identify all pod instances
    local instances
    instances=$(kubectl get pods -n "$NAMESPACE" -l "cnpg.io/cluster=$CNPG_CLUSTER_NAME" -o jsonpath='{.items[*].metadata.name}')
    
    # 4. Recreate non-primary instances one at a time
    for pod in $instances; do
        if [[ "$pod" != "$primary" ]]; then
            echo "Recreating non-primary pod: $pod..."
            kubectl delete pod -n "$NAMESPACE" "$pod"
            
            echo "Waiting for $pod to be recreated and healthy..."
            while true; do
                if kubectl get pod -n "$NAMESPACE" "$pod" >/dev/null 2>&1; then
                    local status
                    status=$(kubectl get pod -n "$NAMESPACE" "$pod" -o jsonpath='{.status.phase}')
                    local ready
                    ready=$(kubectl get pod -n "$NAMESPACE" "$pod" -o jsonpath='{.status.containerStatuses[0].ready}')
                    if [[ "$status" == "Running" && "$ready" == "true" ]]; then
                        echo "Pod $pod is active and healthy."
                        break
                    fi
                fi
                echo -n "."
                sleep 5
            done
            echo ""
            
            # Additional wait for CNPG to recognize it as joined and healthy
            echo "Waiting for CNPG cluster to stabilize..."
            sleep 10
            kubectl cnpg status "$CNPG_CLUSTER_NAME" -n "$NAMESPACE"
        fi
    done
    
    # 5. When the last primary instance needs to be recreated, promote a new primary
    echo "Promoting a new primary to recreate the old primary $primary..."
    # We promote one of the newly recreated instances. 
    # kubectl cnpg promote <cluster> <instance>
    
    local new_primary=""
    for pod in $instances; do
        if [[ "$pod" != "$primary" ]]; then
            new_primary="$pod"
            break
        fi
    done
    
    if [[ -z "$new_primary" ]]; then
         echo "ERROR: No candidate for new primary found." >&2
         exit 1
    fi
    
    echo "Promoting $new_primary..."
    kubectl cnpg promote "$CNPG_CLUSTER_NAME" "$new_primary" -n "$NAMESPACE" || {
        echo "Promotion failed, but continuing rollout (CNPG might have already promoted it)..."
    }
    
    echo "Waiting for $new_primary to become primary..."
    while true; do
        local current_primary
        current_primary=$(kubectl get cluster "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" -o jsonpath='{.status.currentPrimary}')
        if [[ "$current_primary" == "$new_primary" ]]; then
            echo "$new_primary is now the primary."
            break
        fi
        echo -n "."
        sleep 5
    done
    echo ""
    
    # 6. Rebuild the last pod (the old primary) to make the cluster consistent
    echo "Recreating the old primary pod: $primary..."
    kubectl delete pod -n "$NAMESPACE" "$primary"
    
    echo "Waiting for $primary to be recreated and healthy..."
    while true; do
        if kubectl get pod -n "$NAMESPACE" "$primary" >/dev/null 2>&1; then
            local status
            status=$(kubectl get pod -n "$NAMESPACE" "$primary" -o jsonpath='{.status.phase}')
            local ready
            ready=$(kubectl get pod -n "$NAMESPACE" "$primary" -o jsonpath='{.status.containerStatuses[0].ready}')
            if [[ "$status" == "Running" && "$ready" == "true" ]]; then
                echo "Pod $primary is active and healthy."
                break
            fi
        fi
        echo -n "."
        sleep 5
    done
    echo ""
    
    echo "Rollout complete. Final cluster status:"
    kubectl cnpg status "$CNPG_CLUSTER_NAME" -n "$NAMESPACE"
}

backup() {
    echo "Backup - tbd, when we configure s3 or other block store"
}

reset() {
    ensure_tools
    echo "Resetting prole-db cluster..."
    echo "Removing cnpg instances and pods..."
    kubectl delete cluster "$CNPG_CLUSTER_NAME" -n "$NAMESPACE" --ignore-not-found
    
    # Wait for deletion
    kubectl wait --for=delete cluster/"$CNPG_CLUSTER_NAME" -n "$NAMESPACE" --timeout=60s || true
    
    echo "Reinstantiating..."
    deploy
}

case "$ACTION" in
    start)
        start
        ;;
    stop)
        stop
        ;;
    restart)
        restart
        ;;
    deploy)
        deploy
        ;;
    rollback)
        rollback
        ;;
    rollout)
        rollout
        ;;
    backup)
        backup
        ;;
    reset)
        reset
        ;;
    *)
        echo "Usage: $0 {start|stop|restart|deploy|rollback|backup|reset} [version]" >&2
        exit 2
        ;;
esac