#!/bin/bash

# Edge Restore Script
# Restores MongoDB from a backup directory: shutdown, override secret, discover primary, mongorestore-local-tls, startup.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
UTIL_DIR="${SCRIPT_DIR}/util"
KUBECONFIG_FILE=""
BACKUP_DIR=""
CLUSTER_UID=""
IS_HELM=false
SECRET_NAMESPACE="hubble-system"
SECRET_NAME="configserversecret"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
log_warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

show_help() {
    cat << EOF
Edge Restore Script

Restores MongoDB from a backup directory for an edge cluster (cluster-<uid>).
Steps: extract password, shutdown edge, override secret and apply, discover primary, extract target TLS, mongorestore-local-tls, startup.

Usage: $0 --kubeconfig FILE --backup-dir DIR [--uid UID] [--is-helm]

OPTIONS:
    --kubeconfig FILE    Path to kubeconfig file (REQUIRED)
    --backup-dir DIR     Backup directory to restore from (REQUIRED)
    --uid UID            Cluster UID; namespace will be cluster-<UID> (REQUIRED unless --is-helm)
    --is-helm            Helm restore mode (default: false). Skips CLUSTER_UID requirement
                         and propagates --is-helm to shutdown/startup (skips UID-scoped steps).
    --help               Show this help message

EXAMPLES:
    $0 --kubeconfig --kubeconfig /path/to/kubeconfig --backup-dir /path/to/backup_dir --uid abc123
    $0 --kubeconfig /path/to/kubeconfig --backup-dir /path/to/backup_dir --is-helm
EOF
}

parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --kubeconfig)
                KUBECONFIG_FILE="$2"
                shift 2
                ;;
            --backup-dir)
                BACKUP_DIR="$2"
                shift 2
                ;;
            --uid)
                CLUSTER_UID="$2"
                shift 2
                ;;
            --is-helm)
                IS_HELM=true
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done
}

main() {
    parse_arguments "$@"

    if [[ -z "$KUBECONFIG_FILE" ]]; then
        log_error "Kubeconfig path is required (--kubeconfig FILE)"
        show_help
        exit 1
    fi
    if [[ ! -f "$KUBECONFIG_FILE" ]]; then
        log_error "Kubeconfig file not found: $KUBECONFIG_FILE"
        exit 1
    fi
    if [[ -z "$BACKUP_DIR" ]]; then
        log_error "Backup directory is required (--backup-dir DIR)"
        show_help
        exit 1
    fi
    if [[ ! -d "$BACKUP_DIR" ]]; then
        log_error "Backup directory not found: $BACKUP_DIR"
        exit 1
    fi
    if [[ "$IS_HELM" != "true" && -z "$CLUSTER_UID" ]]; then
        log_error "Cluster UID is required (--uid UID)"
        show_help
        exit 1
    fi

    for script in extract-db-password.sh extract-mongo-tls-certs.sh k8s-shutdown-edge.sh get-mongo-primary-pod.sh mongorestore-local-tls.sh override-db-password.sh k8s-startup-edge.sh; do
        if [[ ! -x "${UTIL_DIR}/${script}" ]]; then
            log_error "Required script not found or not executable: ${UTIL_DIR}/${script}"
            exit 1
        fi
    done

    if ! which mongorestore >/dev/null 2>&1; then
        log_error "mongorestore binary not found in PATH. Please install mongo database tools or add mongorestore to PATH."
        exit 1
    fi

    export KUBECONFIG="$KUBECONFIG_FILE"

    # Step 1: Extract DB password
    log_info "Step 1: Extracting DB password from cluster..."
    PASSWORD=$("${UTIL_DIR}/extract-db-password.sh" --kubeconfig "$KUBECONFIG_FILE" 2>/dev/null | tail -1)
    if [[ -z "$PASSWORD" ]]; then
        log_error "Failed to extract DB password"
        exit 1
    fi
    log_success "DB password extracted"

    # Step 2: Kubernetes shutdown (edge)
    log_info "Step 2: Running k8s-shutdown-edge.sh..."
    shutdown_args=(--kubeconfig "$KUBECONFIG_FILE")
    if [[ "$IS_HELM" == "true" ]]; then
        shutdown_args+=(--is-helm)
    else
        shutdown_args+=(--uid "$CLUSTER_UID")
    fi
    if ! "${UTIL_DIR}/k8s-shutdown-edge.sh" "${shutdown_args[@]}"; then
        log_error "k8s-shutdown-edge.sh failed"
        exit 1
    fi
    log_success "Shutdown completed"

    # Step 3: Override DB password in secret from backup and apply
    log_info "Step 3: Overriding dbPassword in configserversecret and applying..."
    SECRET_FILE="${BACKUP_DIR}/cleaned-secrets/configserversecret.yaml"
    if [[ ! -f "$SECRET_FILE" ]]; then
        log_error "Secret file not found: $SECRET_FILE (backup should contain cleaned-secrets from edge-backup)"
        exit 1
    fi
    if ! "${UTIL_DIR}/override-db-password.sh" --password "$PASSWORD" --file "$SECRET_FILE"; then
        log_error "override-db-password.sh failed"
        exit 1
    fi
    if ! kubectl --kubeconfig="$KUBECONFIG_FILE" apply -f "$SECRET_FILE"; then
        log_warning "kubectl apply failed (metadata conflict); trying replace..."
        if ! kubectl --kubeconfig="$KUBECONFIG_FILE" replace -f "$SECRET_FILE" --force; then
            log_error "Failed to apply configserversecret"
            exit 1
        fi
    fi
    log_success "Secret updated and applied"

    log_info "Applying msgbroker-secret..."
    if kubectl --kubeconfig="$KUBECONFIG_FILE" apply -f "${BACKUP_DIR}/cleaned-secrets/msgbroker-secret.yaml"; then
        log_success "msgbroker-secret applied successfully"
    else
        log_error "Failed to apply msgbroker-secret"
        exit 1
    fi

    # Step 4: Extract TLS from target cluster (backup mongo-tls may be from another cluster or rotated)
    TLS_RESTORE_DIR="${BACKUP_DIR}/mongo-tls-restore"
    TLS_CERTS_AVAILABLE=true
    log_info "Step 4: Extracting MongoDB TLS certificates from target cluster to ${TLS_RESTORE_DIR}..."
    if ! "${UTIL_DIR}/extract-mongo-tls-certs.sh" \
        --kubeconfig "$KUBECONFIG_FILE" \
        --output-dir "$TLS_RESTORE_DIR"; then
        log_warning "extract-mongo-tls-certs.sh failed"
        TLS_CERTS_AVAILABLE=false
    else
        log_success "Target cluster TLS certificates extracted"
    fi

    # Step 5: Discover MongoDB primary pod
    log_info "Step 5: Discovering MongoDB primary pod..."
    if [[ "$TLS_CERTS_AVAILABLE" == "true" ]]; then
        POD_NUMBER=$("${UTIL_DIR}/get-mongo-primary-pod.sh" -p "$PASSWORD" -k "$KUBECONFIG_FILE" --tls)
    else
        POD_NUMBER=$("${UTIL_DIR}/get-mongo-primary-pod.sh" -p "$PASSWORD" -k "$KUBECONFIG_FILE")
    fi
    if [[ -z "$POD_NUMBER" ]]; then
        log_error "Failed to get MongoDB primary pod number"
        exit 1
    fi
    log_success "MongoDB primary is mongo-${POD_NUMBER}"

    # Step 6: Run local mongorestore via port-forward
    log_info "Step 6: Running mongorestore-local-tls.sh from $BACKUP_DIR..."
    mongorestore_args=(
        -p "$PASSWORD"
        -l "$BACKUP_DIR"
        -k "$KUBECONFIG_FILE"
        -N "$POD_NUMBER"
    )
    if [[ "$TLS_CERTS_AVAILABLE" == "true" ]]; then
        mongorestore_args+=(--tls-dir "$TLS_RESTORE_DIR")
    fi
    if ! "${UTIL_DIR}/mongorestore-local-tls.sh" "${mongorestore_args[@]}"; then
        log_error "mongorestore-local-tls.sh failed"
        exit 1
    fi
    log_success "Restore completed"

    # Step 7: Kubernetes startup (edge)
    log_info "Step 7: Running k8s-startup-edge.sh..."
    startup_args=(--kubeconfig "$KUBECONFIG_FILE")
    if [[ "$IS_HELM" == "true" ]]; then
        startup_args+=(--is-helm)
    else
        startup_args+=(--uid "$CLUSTER_UID")
    fi
    if ! "${UTIL_DIR}/k8s-startup-edge.sh" "${startup_args[@]}"; then
        log_error "k8s-startup-edge.sh failed"
        exit 1
    fi

    log_success "Edge restore completed successfully."
}

main "$@"
