#!/bin/bash

set -euo pipefail

# =============================================================================
# Fetch Cluster CRs Script
# =============================================================================
# This script fetches Kubernetes Cluster CRs with bizType filter,
# adds "migrated: true" label, and stores them in JSON files.
# =============================================================================

# --- Configuration ---
CRD_NAME="clusters.observability.pingcap.com"
LABEL_KEY="observability.pingcap.com/bizType"
MIGRATION_LABEL="observability.pingcap.com/migrated"

# Default values to filter
VALUES_TO_FILTER=(
    "prod-devtier"
    "dev-devtier"
    "staging-devtier"
)

# Directory structure
WORK_DIR="./migration-$(date +%Y%m%d-%H%M%S)"
JSON_DIR="${WORK_DIR}/json"
LOG_FILE="${WORK_DIR}/fetch.log"

# Batch size for splitting
BATCH_SIZE=5000

# Options
FILTER_LABELS=""
DRY_RUN=false
SKIP_LABEL_UPDATE=false

# --- Helper Functions ---

log() {
    local msg="[$(date +'%Y-%m-%d %H:%M:%S')] $*"
    echo "$msg"
    if [[ -n "${LOG_FILE:-}" ]] && [[ -d "$(dirname "${LOG_FILE}")" ]]; then
        echo "$msg" >> "${LOG_FILE}"
    fi
}

error() {
    local msg="[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $*"
    echo "$msg" >&2
    if [[ -n "${LOG_FILE:-}" ]] && [[ -d "$(dirname "${LOG_FILE}")" ]]; then
        echo "$msg" >> "${LOG_FILE}"
    fi
}

usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Fetch Kubernetes Cluster CRs and prepare them for migration.

OPTIONS:
    -b, --biz-types VALUES      Comma-separated bizType values to filter
                                (default: prod-devtier,dev-devtier,staging-devtier)
    -l, --labels LABELS         Additional label selector (e.g., "env=prod,tier=premium")
    -w, --work-dir DIR          Working directory for output files (default: ./migration-TIMESTAMP)
    -d, --dry-run               Dry run mode - don't update CRs
    -s, --skip-label            Skip adding migrated label to CRs
    -h, --help                  Show this help message

EXAMPLES:
    # Basic fetch
    $0

    # Custom bizTypes
    $0 --biz-types "prod-devtier,staging-devtier"

    # Filter by additional labels
    $0 --labels "env=prod,region=us-west-2"

    # Dry run to see what would be fetched
    $0 --dry-run

    # Skip labeling (just fetch)
    $0 --skip-label

OUTPUT:
    The script creates a working directory with:
    - json/all_clusters.jsonl    # All fetched CRs (one per line)
    - json/batch_*               # Split into batches
    - fetch.log                  # Detailed log
    - metadata.json              # Metadata about the fetch

EOF
    exit 0
}

# Parse command line arguments
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -b|--biz-types)
                IFS=',' read -ra VALUES_TO_FILTER <<< "$2"
                shift 2
                ;;
            -l|--labels)
                FILTER_LABELS="$2"
                shift 2
                ;;
            -w|--work-dir)
                WORK_DIR="$2"
                JSON_DIR="${WORK_DIR}/json"
                LOG_FILE="${WORK_DIR}/fetch.log"
                shift 2
                ;;
            -d|--dry-run)
                DRY_RUN=true
                shift
                ;;
            -s|--skip-label)
                SKIP_LABEL_UPDATE=true
                shift
                ;;
            -h|--help)
                usage
                ;;
            *)
                error "Unknown option: $1"
                usage
                ;;
        esac
    done

    mkdir -p "${JSON_DIR}"
    log "Created working directory: $WORK_DIR"
}

# Build kubectl label selector
build_kubectl_selector() {
    local selectors=()

    # Add additional label filters
    if [[ -n "$FILTER_LABELS" ]]; then
        IFS=',' read -ra label_pairs <<< "$FILTER_LABELS"
        for pair in "${label_pairs[@]}"; do
            selectors+=("$pair")
        done
    fi

    # Join with commas
    if [[ ${#selectors[@]} -gt 0 ]]; then
        local selector=$(IFS=','; echo "${selectors[*]}")
        echo "$selector"
    else
        echo ""
    fi
}

# Build jq filter for CR selection (bizType filtering)
build_jq_filter() {
    local jq_values=$(printf "\"%s\", " "${VALUES_TO_FILTER[@]}" | sed 's/, $//')
    local filter=".items[] | select(.metadata.labels[\"${LABEL_KEY}\"] | IN(${jq_values}))"

    echo "$filter"
}

# Add migrated label to JSON (not to K8s object)
add_migrated_label_to_json() {
    local input_file=$1
    local output_file=$2

    if [[ "$DRY_RUN" == true ]]; then
        log "  [DRY-RUN] Would add label to JSON"
        cp "$input_file" "$output_file"
        return 0
    fi

    # Add the migrated label to each JSON object (line by line for JSONL)
    while IFS= read -r line; do
        echo "$line" | jq -c ".metadata.labels[\"${MIGRATION_LABEL}\"] = \"true\""
    done < "$input_file" > "$output_file"
}

# Fetch all matching CRs
fetch_clusters() {
    log "=== Fetching Cluster CRs ==="
    log "Filtering by ${LABEL_KEY} IN: ${VALUES_TO_FILTER[*]}"
    [[ -n "$FILTER_LABELS" ]] && log "Additional labels: $FILTER_LABELS"

    local kubectl_selector=$(build_kubectl_selector)
    local jq_filter=$(build_jq_filter)

    # Build kubectl command with label selector
    local kubectl_cmd="kubectl get ${CRD_NAME} -A -o json --chunk-size=500"
    if [[ -n "$kubectl_selector" ]]; then
        kubectl_cmd="${kubectl_cmd} -l ${kubectl_selector}"
        log "Using kubectl label selector: ${kubectl_selector}"
    fi


    # Fetch and save CRs
    log "Fetching cluster CRs..."
    eval "$kubectl_cmd" | \
        jq -c "${jq_filter}" > "${JSON_DIR}/all_clusters.jsonl"

    local fetched_count=$(wc -l < "${JSON_DIR}/all_clusters.jsonl" | tr -d ' ')
    log "Fetched ${fetched_count} clusters"

		label_clusters

    # Split into batches
    log "Splitting into batches of ${BATCH_SIZE}..."
    split -l "${BATCH_SIZE}" -d -a 4 "${JSON_DIR}/all_clusters.jsonl" "${JSON_DIR}/batch_"


    local batch_count=$(ls -1 "${JSON_DIR}"/batch_* 2>/dev/null | wc -l | tr -d ' ')
    log "Created ${batch_count} batches"
}

# Add migrated labels to JSON files
label_clusters() {
    if [[ "$SKIP_LABEL_UPDATE" == true ]]; then
        log "Skipping label updates (--skip-label)"
        return 0
    fi

    log "=== Adding Migrated Labels to JSON ==="

    # Add label to all_clusters.jsonl
    local temp_file="${JSON_DIR}/all_clusters.jsonl.tmp"
    add_migrated_label_to_json "${JSON_DIR}/all_clusters.jsonl" "$temp_file"
    mv "$temp_file" "${JSON_DIR}/all_clusters.jsonl"

    local total=$(wc -l < "${JSON_DIR}/all_clusters.jsonl" | tr -d ' ')
    log "Added migrated label to ${total} clusters in JSON"
}

# Generate metadata file
generate_metadata() {
    log "=== Generating Metadata ==="

    local metadata_file="${WORK_DIR}/metadata.json"
    local total_clusters=$(wc -l < "${JSON_DIR}/all_clusters.jsonl" | tr -d ' ')
    local batch_count=$(ls -1 "${JSON_DIR}"/batch_* 2>/dev/null | wc -l | tr -d ' ')

    cat > "$metadata_file" <<EOF
{
  "fetch_timestamp": "$(date -u +%Y-%m-%dT%H:%M:%SZ)",
  "work_dir": "${WORK_DIR}",
  "filters": {
    "biz_types": $(printf '%s\n' "${VALUES_TO_FILTER[@]}" | jq -R . | jq -s .),
    "additional_labels": "${FILTER_LABELS}",
    "exclude_migrated": $([ "$SKIP_LABEL_UPDATE" == false ] && echo "true" || echo "false")
  },
  "statistics": {
    "total_clusters": ${total_clusters},
    "batch_size": ${BATCH_SIZE},
    "batch_count": ${batch_count}
  },
  "options": {
    "dry_run": ${DRY_RUN},
    "skip_label": ${SKIP_LABEL_UPDATE}
  }
}
EOF

    log "Metadata written to: ${metadata_file}"
    cat "$metadata_file" | jq . | tee -a "${LOG_FILE}"
}

# Main execution
main() {
    log "=== Cluster CR Fetch Script ==="
    log "Start time: $(date)"

    parse_args "$@"

    if [[ "$DRY_RUN" == true ]]; then
        log "=== DRY RUN MODE ==="
    fi

    # Check prerequisites
    if ! command -v kubectl &> /dev/null; then
        error "kubectl not found. Please install kubectl."
        exit 1
    fi

    if ! command -v jq &> /dev/null; then
        error "jq not found. Please install jq."
        exit 1
    fi

    # Execute fetch
    fetch_clusters

    generate_metadata

    log "=== Fetch Complete ==="
    log "End time: $(date)"
    log "Working directory: ${WORK_DIR}"
    log "Log file: ${LOG_FILE}"
    log ""
    log "Next steps:"
    log "1. Review the fetched data: ${JSON_DIR}/all_clusters.jsonl"
    log "2. Generate SQL: ./generate-sql.sh --input-dir ${WORK_DIR}"
}

# Run main function
main "$@"