#!/usr/bin/env bash

set -euo pipefail

# Source common functions
SCRIPT_DIR=$(cd $(dirname "$0");pwd)
source "${SCRIPT_DIR}/common.sh"
TICI_SERVER_BINPATH=${TICI_SERVER_BINPATH:-"${PROJECT_ROOT}/target/debug/tici-server"}
declare -a COMPONENT_DOWNLOAD_SUMMARY=()
declare -a PREPARE_STAGE_SUMMARY=()
PREPARE_SUMMARY_PRINTED=0
PREPARE_TOTAL_START_SECONDS=${SECONDS}

print_playground_status() {
    local playground_version=$1
    log "Printing TiDB playground cluster status..."

    # Prefer TiUP status output; if unavailable, fall back to key lines in playground log.
    if [ -n "${PLAYGROUND_TAG:-}" ]; then
        if tiup "playground${playground_version}" display --tag "${PLAYGROUND_TAG}"; then
            return 0
        fi
        log_error "Failed to display playground status by tag (${PLAYGROUND_TAG})"
        grep -E "Connect TiDB:|Connect PD:|Dashboard:|TiFlash .*failed to start" "${TEST_OUTPUT_DIR}/log/playground.log" || true
        return 0
    fi

    if tiup "playground${playground_version}" display; then
        return 0
    fi

    log_error "Failed to run 'tiup playground${playground_version} display', fallback to playground log summary"
    grep -E "Connect TiDB:|Connect PD:|Dashboard:" "${TEST_OUTPUT_DIR}/log/playground.log" || true
    return 0
}

build_tici() {
    log "Building TiCI server..."
    cargo build --features "openssl-vendored failpoints"
}

record_component_download_summary() {
    local component=$1
    local detail=$2

    COMPONENT_DOWNLOAD_SUMMARY+=("${component}|${detail}")
}

run_timed_component_download() {
    local component=$1
    shift
    local start_seconds=${SECONDS}

    if "$@"; then
        record_component_download_summary "${component}" "$((SECONDS - start_seconds))s"
        return 0
    fi

    record_component_download_summary "${component}" "failed after $((SECONDS - start_seconds))s"
    return 1
}

print_component_download_summary() {
    local item
    local component
    local detail

    log "Component download summary:"
    for item in "${COMPONENT_DOWNLOAD_SUMMARY[@]}"; do
        component=${item%%|*}
        detail=${item#*|}
        log "  - ${component}: ${detail}"
    done
}

record_prepare_stage_summary() {
    local stage=$1
    local detail=$2

    PREPARE_STAGE_SUMMARY+=("${stage}|${detail}")
}

run_timed_prepare_stage() {
    local stage=$1
    shift
    local start_seconds=${SECONDS}

    if "$@"; then
        record_prepare_stage_summary "${stage}" "$((SECONDS - start_seconds))s"
        return 0
    fi

    record_prepare_stage_summary "${stage}" "failed after $((SECONDS - start_seconds))s"
    return 1
}

print_prepare_stage_summary() {
    local item
    local stage
    local detail

    log "Preparation stage summary:"
    for item in "${PREPARE_STAGE_SUMMARY[@]}"; do
        stage=${item%%|*}
        detail=${item#*|}
        log "  - ${stage}: ${detail}"
    done
    log "  - total: $((SECONDS - PREPARE_TOTAL_START_SECONDS))s"
}

print_preparation_summary() {
    if [ "${PREPARE_SUMMARY_PRINTED}" -eq 1 ]; then
        return 0
    fi

    PREPARE_SUMMARY_PRINTED=1
    log "Preparation summary:"
    print_prepare_stage_summary

    if [ ${#COMPONENT_DOWNLOAD_SUMMARY[@]} -gt 0 ]; then
        print_component_download_summary
    fi
}

install_tiup_component_path_from_oci_image() {
    local image_repo=$1
    local image_tag=$2
    local archive_path=$3
    local install_path=$4
    local platform_os
    platform_os=$(compute_platform_os)
    local platform_arch
    platform_arch=$(compute_platform_arch) || return 1
    local image_ref="${image_repo}:${image_tag}"
    local install_dir
    install_dir=$(dirname "${install_path}")
    local tmp_dir
    tmp_dir=$(mktemp -d)
    local platform_digest
    local resolved_ref
    local layer_digest
    local layer_tar="${tmp_dir}/layer.tar.gz"
    local archive_member
    local normalized_archive_member
    local extract_root="${tmp_dir}/extract"

    mkdir -p "${install_dir}"
    log "Installing ${archive_path} from OCI image ${image_ref} to ${install_path}"

    if ! platform_digest=$(oras manifest fetch "${image_ref}" | python3 -c '
import json
import sys

data = json.load(sys.stdin)
manifests = data.get("manifests")

if not manifests:
    print("")
    raise SystemExit(0)

platform_os = sys.argv[1]
platform_arch = sys.argv[2]
for manifest in manifests:
    platform = manifest.get("platform", {})
    if platform.get("os") == platform_os and platform.get("architecture") == platform_arch:
        print(manifest["digest"])
        raise SystemExit(0)

raise SystemExit(1)
' "${platform_os}" "${platform_arch}"
    ); then
        log_error "Failed to resolve ${image_ref} for ${platform_os}/${platform_arch}"
        rm -rf "${tmp_dir}"
        return 1
    fi

    resolved_ref="${image_ref}"
    if [ -n "${platform_digest}" ]; then
        resolved_ref="${image_repo}@${platform_digest}"
    fi

    # TiDB/TiKV release images copy the server binary in the top layer.
    if ! layer_digest=$(oras manifest fetch "${resolved_ref}" | python3 -c '
import json
import sys

data = json.load(sys.stdin)
layers = data.get("layers", [])
if not layers:
    raise SystemExit(1)

print(layers[-1]["digest"])
'
    ); then
        log_error "Failed to locate binary layer from ${resolved_ref}"
        rm -rf "${tmp_dir}"
        return 1
    fi

    if ! oras blob fetch --output "${layer_tar}" "${image_repo}@${layer_digest}" >/dev/null; then
        log_error "Failed to download layer ${layer_digest} from ${image_repo}"
        rm -rf "${tmp_dir}"
        return 1
    fi

    if ! archive_member=$(tar -tzf "${layer_tar}" | python3 -c '
import sys

expected_path = sys.argv[1].rstrip("/")
for line in sys.stdin:
    path = line.strip()
    normalized = path.lstrip("./").lstrip("/")
    if normalized.rstrip("/") == expected_path:
        print(path)
        raise SystemExit(0)

raise SystemExit(1)
' "${archive_path}"
    ); then
        log_error "Failed to find ${archive_path} in ${image_ref}"
        rm -rf "${tmp_dir}"
        return 1
    fi

    normalized_archive_member=$(python3 -c '
import sys

print(sys.argv[1].lstrip("./").lstrip("/").rstrip("/"))
' "${archive_member}")

    mkdir -p "${extract_root}"
    rm -rf "${install_path}"
    if ! tar -xzf "${layer_tar}" -C "${extract_root}" "${archive_member}"; then
        log_error "Failed to extract ${archive_path} from ${image_ref}"
        rm -rf "${tmp_dir}"
        return 1
    fi

    if ! cp -a "${extract_root}/${normalized_archive_member}" "${install_path}"; then
        log_error "Failed to install ${archive_path} to ${install_path}"
        rm -rf "${tmp_dir}"
        return 1
    fi

    rm -rf "${tmp_dir}"
    log "${archive_path} installed successfully: ${install_path}"
    return 0
}

needs_tici_build() {
    if [ "${TICI_FORCE_BUILD:-0}" = "1" ]; then
        log "TICI_FORCE_BUILD=1, rebuilding TiCI server"
        return 0
    fi

    if [ ! -f "${TICI_SERVER_BINPATH}" ]; then
        log "TiCI server binary not found at ${TICI_SERVER_BINPATH}"
        return 0
    fi

    local rebuild_inputs=(
        "${PROJECT_ROOT}/Cargo.toml"
        "${PROJECT_ROOT}/Cargo.lock"
        "${PROJECT_ROOT}/rust-toolchain.toml"
        "${PROJECT_ROOT}/rust-toolchain"
    )

    local input
    for input in "${rebuild_inputs[@]}"; do
        if [ -f "${input}" ] && [ "${input}" -nt "${TICI_SERVER_BINPATH}" ]; then
            log "Detected newer build input: ${input}"
            return 0
        fi
    done

    local newer_source
    newer_source=$(
        find "${PROJECT_ROOT}/src" "${PROJECT_ROOT}/components" \
            -type f \
            \( -name '*.rs' -o -name 'Cargo.toml' \) \
            -newer "${TICI_SERVER_BINPATH}" \
            -print \
            -quit
    )
    if [ -n "${newer_source}" ]; then
        log "Detected newer Rust source: ${newer_source}"
        return 0
    fi

    return 1
}

install_tiup_component_binary_from_oci_image() {
    local image_repo=$1
    local image_tag=$2
    local binary_name=$3
    local install_path=$4

    if ! install_tiup_component_path_from_oci_image "${image_repo}" "${image_tag}" "${binary_name}" "${install_path}"; then
        return 1
    fi

    chmod +x "${install_path}"
    return 0
}

install_e2e_tikv_binary_from_oci() {
    local tikv_binary_version=$1
    local tikv_component_version=$2
    local tiup_home="${TIUP_HOME:-$HOME/.tiup}"
    local oci_artifact_host="${OCI_ARTIFACT_HOST:-us-docker.pkg.dev/pingcap-testing-account/hub}"

    install_tiup_component_binary_from_oci_image \
        "${oci_artifact_host}/tikv/tikv/image" \
        "${tikv_binary_version}" \
        "tikv-server" \
        "${tiup_home}/components/tikv/${tikv_component_version}/tikv-server"
}

install_e2e_tikv_binary_from_tiup() {
    local tikv_binary_version=$1
    local tikv_component_version=$2

    tiup install "tikv:${tikv_binary_version}"
    prepare_tikv_component_binary "${tikv_binary_version}" "${tikv_component_version}"
}

install_e2e_tidb_binary_from_oci() {
    local tidb_image_tag=$1
    local tidb_install_version=$2
    local tiup_home="${TIUP_HOME:-$HOME/.tiup}"
    local oci_artifact_host="${OCI_ARTIFACT_HOST:-us-docker.pkg.dev/pingcap-testing-account/hub}"

    install_tiup_component_binary_from_oci_image \
        "${oci_artifact_host}/pingcap/tidb/images/tidb-server" \
        "${tidb_image_tag}" \
        "tidb-server" \
        "${tiup_home}/components/tidb/${tidb_install_version}/tidb-server"
}

install_nightly_components() {
    if [ -n "${TIUP_MIRROR_NIGHTLY:-}" ]; then
        log "TiUP mirror (nightly): ${TIUP_MIRROR_NIGHTLY}"
        tiup mirror set "${TIUP_MIRROR_NIGHTLY}"
    else
        log "TiUP mirror (nightly): not set (using default)"
    fi

    log "Installing TiUP components (nightly)..."
    run_timed_component_download "pd:nightly (tiup)" tiup install pd:nightly || return 1
    run_timed_component_download "cdc:nightly (tiup)" tiup install cdc:nightly || return 1
    run_timed_component_download "prometheus:nightly (tiup)" tiup install prometheus:nightly || return 1
    run_timed_component_download "grafana:nightly (tiup)" tiup install grafana:nightly || return 1
    # I don't know why tiup playground:v1.16.2-feature.fts will download prometheus and grafana latest stable version, but it does.
    run_timed_component_download "prometheus:latest stable (tiup)" tiup install prometheus || return 1
    run_timed_component_download "grafana:latest stable (tiup)" tiup install grafana || return 1
}

install_fts_components() {
    local playground_component_version=$1
    local tidb_component_version=$2
    local tiflash_component_version=$3
    local tikv_binary_version=$4
    local tikv_component_version=$5
    local fts_install_version

    fts_install_version=$(fts_component_install_version)

    if [ -z "${TIUP_MIRROR_FTS:-}" ]; then
        TIUP_MIRROR_FTS="http://staging.tiup-server.pingcap.net"
        export TIUP_MIRROR_FTS
    fi

    log "TiUP mirror (fts): ${TIUP_MIRROR_FTS}"
    tiup mirror set "${TIUP_MIRROR_FTS}"

    log "Installing FTS playground from TiUP mirror and TiDB/TiKV/TiFlash from OCI images..."
    run_timed_component_download "playground:${playground_component_version} (tiup)" tiup install "playground:${playground_component_version}" || return 1
    run_timed_component_download "tidb:${tidb_component_version} -> ${fts_install_version} (oci)" install_e2e_tidb_binary_from_oci "${tidb_component_version}" "${fts_install_version}" || return 1
    run_timed_component_download "tikv:${tikv_binary_version} -> ${fts_install_version} (oci)" install_e2e_tikv_binary_from_oci "${tikv_binary_version}" "${fts_install_version}" || return 1
    run_timed_component_download "tiflash:${tiflash_component_version} -> ${fts_install_version} (oci)" install_e2e_tiflash_binary_from_oci "${tiflash_component_version}" "${fts_install_version}" || return 1
}

# Main preparation process
main() {
    log "Starting e2e test preparation..."

    trap 'print_preparation_summary' EXIT

    # Check prerequisites
    if ! run_timed_prepare_stage "check prerequisites" check_prerequisites; then
        log_error "Prerequisites check failed"
        exit 1
    fi

    # Setup cleanup trap
    trap 'cleanup; exit 130' INT TERM

    # Rebuild if the binary is missing or stale relative to Rust sources/manifests.
    log "Build or verify TiCI server binary..."
    if needs_tici_build; then
        if ! run_timed_prepare_stage "build TiCI server" build_tici; then
            log_error "Failed to build TiCI server"
            exit 1
        fi
    else
        record_prepare_stage_summary "build TiCI server" "skipped (existing binary)"
        log "Using existing TiCI server binary at ${TICI_SERVER_BINPATH}"
    fi

    # Prepare dataset
    log "Prepare dataset..."
    if ! run_timed_prepare_stage "prepare dataset" prepare_dataset; then
        log_error "Failed to prepare dataset"
        exit 1
    fi

    # Start MinIO
    log "Starting MinIO server..."
    if ! run_timed_prepare_stage "start MinIO" start_minio; then
        log_error "Failed to start MinIO"
        exit 1
    fi

    # Prepare playground arguments
    local PLAYGROUND_ARGS="--tici.binpath ${TICI_SERVER_BINPATH} --s3.bucket ticidefaultbucket"
    PLAYGROUND_ARGS=$(append_fts_playground_binpath_args "${PLAYGROUND_ARGS}")
    if [ -n "${TICI_META_CONFIG:-}" ]; then
        PLAYGROUND_ARGS+=" --tici.meta.config $TICI_META_CONFIG"
    fi
    local meta_count=${TICI_META_COUNT:-1}
    PLAYGROUND_ARGS+=" --tici.meta $meta_count"
    local worker_count=${TICI_WORKER_COUNT:-1}
    PLAYGROUND_ARGS+=" --tici.worker $worker_count"
    local reader_count=${TICI_READER_COUNT:-1}
    local reader_config=${TICI_READER_CONFIG:-}
    if [ -z "${reader_config}" ]; then
        reader_config=${TIFLASH_CONFIG:-"${PROJECT_ROOT}/tests/config/tiflash_reader.toml"}
    fi
    PLAYGROUND_ARGS+=" --tiflash $reader_count"
    PLAYGROUND_ARGS+=" --tiflash.config $reader_config"
    if [ -n "${PLAYGROUND_TAG:-}" ]; then
        PLAYGROUND_ARGS+=" --tag $PLAYGROUND_TAG"
    fi
    if [ -n "${PLAYGROUND_PORT_OFFSET:-}" ]; then
        PLAYGROUND_ARGS+=" --port-offset $PLAYGROUND_PORT_OFFSET"
    fi

    local playground_component_version=${PLAYGROUND_COMPONENT_VERSION:-v1.16.2-feature.fts}
    local tidb_component_version=${TIDB_COMPONENT_VERSION:-feature-fts}
    local tiflash_component_version=${TIFLASH_VERSION:-feature-fts}
    local tikv_binary_version=${TIKV_BINARY_VERSION:-feature-fts}
    local tikv_component_version=${TIKV_COMPONENT_VERSION:-feature-fts}
    if ! run_timed_prepare_stage "install TiUP nightly components" install_nightly_components; then
        log_error "Failed to install TiUP components (nightly)"
        exit 1
    fi

    if ! run_timed_prepare_stage "install FTS components" install_fts_components \
        "${playground_component_version}" \
        "${tidb_component_version}" \
        "${tiflash_component_version}" \
        "${tikv_binary_version}" \
        "${tikv_component_version}"; then
        log_error "Failed to install FTS components"
        exit 1
    fi

    if ! run_timed_prepare_stage "validate FTS component binaries" validate_fts_playground_binpaths; then
        log_error "FTS component binaries are not ready"
        exit 1
    fi

    # Start TiUP playground cluster
    log "Starting TiUP playground cluster..."
    local playground_version=":${playground_component_version}"
    if ! run_timed_prepare_stage "start TiUP playground cluster" start_playground "${playground_version}" "--mode tidb-fts" "$PLAYGROUND_ARGS"; then
        log_error "Failed to start TiUP playground cluster"
        exit 1
    fi

    if ! run_timed_prepare_stage "configure TiDB global sort URI" configure_tidb_global_sort_uri "ticidefaultbucket" "${PLAYGROUND_TAG:-default}/global-sort"; then
        log_error "Failed to configure TiDB global sort URI"
        exit 1
    fi
    if ! run_timed_prepare_stage "print playground status" print_playground_status "${playground_version}"; then
        log_error "Failed to print playground status"
        exit 1
    fi

    log "e2e test environment prepared successfully!"
    log "Services running:"
    log "  - TiDB: 127.0.0.1:$(get_tidb_port)"
    log "  - PD: 127.0.0.1:$(get_pd_port)"
    log "  - TiCDC: 127.0.0.1:$(get_cdc_port)"
    log "  - MinIO: 127.0.0.1:${DEFAULT_MINIO_PORT}"
    log "  - MinIO Console: 127.0.0.1:${DEFAULT_MINIO_ADMIN_PORT}"
    log ""
    log "Logs are available in: ${TEST_OUTPUT_DIR}/log/"
    log "Ready to run e2e tests!"

    local playground_log="${TEST_OUTPUT_DIR}/log/playground.log"
    if [ -f "${playground_log}" ]; then
        log "playground.log:"
        cat "${playground_log}"
    else
        log "playground.log not exist in: ${TEST_OUTPUT_DIR}/log/"
    fi
}

print_usage() {
    echo "Usage: $0"
    echo "Prepare the environment for e2e tests."
}

case "${1:-}" in
    "")
        ;;
    --help|-h)
        print_usage
        exit 0
        ;;
    *)
        echo "Unknown option: $1"
        echo "Use --help or -h for usage information."
        exit 1
        ;;
esac

# Run main function
main "$@"