#!/usr/bin/env bash
set -euo pipefail

CANN_ENV=${CANN_ENV:-/usr/local/Ascend/cann-8.5.1/cann-8.5.1/set_env.sh}
VENV_ACTIVATE=${VENV_ACTIVATE:-/data/vllm-ascend-v0.18.0/.venv/bin/activate}
source "${CANN_ENV}"
source "${VENV_ACTIVATE}"

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"

VLLM_ASCEND_CUSTOM_OP_ROOT=${VLLM_ASCEND_CUSTOM_OP_ROOT-/data/vllm-ascend-v0.18.0/vllm_ascend/_cann_ops_custom/vendors/vllm-ascend}
if [[ -n "${VLLM_ASCEND_CUSTOM_OP_ROOT}" ]]; then
    export ASCEND_CUSTOM_OPP_PATH="${VLLM_ASCEND_CUSTOM_OP_ROOT}${ASCEND_CUSTOM_OPP_PATH:+:${ASCEND_CUSTOM_OPP_PATH}}"
    export LD_LIBRARY_PATH="${VLLM_ASCEND_CUSTOM_OP_ROOT}/op_api/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
fi
export CANN_ENV VENV_ACTIVATE VLLM_ASCEND_CUSTOM_OP_ROOT

export VLLM_TARGET_DEVICE=ascend
export VLLM_USE_MODELSCOPE=True
export ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-0}
export NPU_INDEX=${NPU_INDEX:-$(python "${SCRIPT_DIR}/detect_npu_index.py")}
export NPU_CHIP_INDEX=${NPU_CHIP_INDEX:-0}
export XDG_CACHE_HOME=${XDG_CACHE_HOME:-/data/cache}
export HF_HOME=${HF_HOME:-/data/huggingface-cache}
export MODELSCOPE_CACHE=${MODELSCOPE_CACHE:-/data/modelscope-cache}
export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-/data/triton-cache-cann851}

MODEL_PATH=${MODEL_PATH:?MODEL_PATH must point to the model}
REPRO_RESULTS_DIR=${REPRO_RESULTS_DIR:-/data/reproduction-results}
PERFORMANCE_DIR="${REPRO_RESULTS_DIR}/performance"
ACCURACY_DIR="${REPRO_RESULTS_DIR}/accuracy"
STARTED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
mkdir -p "${PERFORMANCE_DIR}" "${ACCURACY_DIR}"

RESULTS_DIR="${PERFORMANCE_DIR}" bash "${SCRIPT_DIR}/run_benchmark_all.sh"
RESULTS_DIR="${ACCURACY_DIR}" bash "${SCRIPT_DIR}/run_accuracy_gsm8k_gpqa.sh"

FINISHED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
MANIFEST_PATH="${REPRO_RESULTS_DIR}/reproduction-manifest.json"
export REPO_ROOT MODEL_PATH REPRO_RESULTS_DIR PERFORMANCE_DIR ACCURACY_DIR
export STARTED_AT FINISHED_AT MANIFEST_PATH

python - <<'PY'
import importlib.metadata
import json
import os
import shlex
import subprocess
from pathlib import Path


def package_version(distribution):
    try:
        return importlib.metadata.version(distribution)
    except importlib.metadata.PackageNotFoundError:
        return None


source_commit = subprocess.run(
    ["git", "-C", "/data/vllm-ascend-v0.18.0", "rev-parse", "HEAD"],
    check=True,
    capture_output=True,
    text=True,
).stdout.strip()
performance_dir = Path(os.environ["PERFORMANCE_DIR"])
accuracy_dir = Path(os.environ["ACCURACY_DIR"])
repo_root = Path(os.environ["REPO_ROOT"])
performance_command = shlex.join([
    "env",
    f"RESULTS_DIR={performance_dir}",
    "bash",
    str(repo_root / "scripts" / "run_benchmark_all.sh"),
])
accuracy_command = shlex.join([
    "env",
    f"RESULTS_DIR={accuracy_dir}",
    "bash",
    str(repo_root / "scripts" / "run_accuracy_gsm8k_gpqa.sh"),
])
payload = {
    "started_at": os.environ["STARTED_AT"],
    "finished_at": os.environ["FINISHED_AT"],
    "commands": [
        performance_command,
        accuracy_command,
    ],
    "package_versions": {
        name: package_version(name)
        for name in (
            "vllm", "vllm-ascend", "torch", "torch-npu", "triton-ascend"
        )
    },
    "source_commit": source_commit,
    "model_path": os.environ["MODEL_PATH"],
    "artifacts": {
        "performance_jsonl": str(performance_dir / "performance-runs.jsonl"),
        "performance_csv": str(performance_dir / "performance-summary.csv"),
        "performance_markdown": str(performance_dir / "performance-summary.md"),
        "performance_json": str(performance_dir / "performance-summary.json"),
        "accuracy_directory": str(accuracy_dir),
        "accuracy_summary_json": str(accuracy_dir / "accuracy-summary.json"),
        "accuracy_summary_csv": str(accuracy_dir / "accuracy-summary.csv"),
        "accuracy_summary_markdown": str(accuracy_dir / "accuracy-summary.md"),
    },
}
manifest = Path(os.environ["MANIFEST_PATH"])
manifest.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
print(manifest)
PY
