#!/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 BOOTSTRAP_STAGE_SUMMARY=()
BOOTSTRAP_SUMMARY_PRINTED=0
BOOTSTRAP_TOTAL_START_SECONDS=${SECONDS}

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_bootstrap_stage_summary() {
    local stage=$1
    local detail=$2

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

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

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

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

print_bootstrap_stage_summary() {
    local item
    local stage
    local detail

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

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

    BOOTSTRAP_SUMMARY_PRINTED=1
    log "Bootstrap summary:"
    print_bootstrap_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

    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
}

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() {
    log "Starting e2e bootstrap..."

    trap 'print_bootstrap_summary' EXIT

    if ! run_timed_bootstrap_stage "check prerequisites" check_prerequisites; then
        log_error "Prerequisites check failed"
        exit 1
    fi

    log "Build or verify TiCI server binary..."
    if [ ! -f "${TICI_SERVER_BINPATH}" ]; then
        if ! run_timed_bootstrap_stage "build TiCI server" build_tici; then
            log_error "Failed to build TiCI server"
            exit 1
        fi
    else
        record_bootstrap_stage_summary "build TiCI server" "skipped (existing binary)"
    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_bootstrap_stage "install TiUP nightly components" install_nightly_components; then
        log_error "Failed to install TiUP components (nightly)"
        exit 1
    fi

    if ! run_timed_bootstrap_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

    log "e2e bootstrap completed successfully!"
}

main "$@"