#!/usr/bin/env python3
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""Runner for the Comet micro benchmark suites on a single machine (typically EC2).

The script is self-contained and can be downloaded on its own, before the
repository is cloned:

    curl -sSLO https://raw.githubusercontent.com/apache/datafusion-comet/main/benchmarks/micro/run.py
    python3 run.py all

Subcommands:

    setup     install prerequisites, clone/update the repository, build a release
    run       run the benchmark suites, one JVM per suite
    collect   copy the generated results into benchmarks/results/micro
    publish   commit the collected results and optionally open a pull request
    all       setup + run + collect

Examples:

    python3 run.py all --ref main
    python3 run.py run --only Cast
    python3 run.py run --suites my-suites.txt --heap 12g
    python3 run.py collect --publish
    python3 run.py publish --push --open-pr
"""

import argparse
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request
import zipfile
from datetime import datetime, timezone
from pathlib import Path

BENCH_PACKAGE = "org.apache.spark.sql.benchmark"
BENCH_SOURCE_DIR = Path("spark/src/test/scala/org/apache/spark/sql/benchmark")
DEFAULT_REPO = "https://github.com/apache/datafusion-comet.git"
DEFAULT_COMET_HOME = Path.home() / "datafusion-comet"
DEFAULT_RUNS_ROOT = Path.home() / "comet-bench-runs"
DEFAULT_HEAP = "8g"
DEFAULT_JDK = "17"
# Spark 4.x, the default build profile, does not compile on anything older
MINIMUM_JDK = 17
DEFAULT_TIMEOUT_MINUTES = 60
# Only used when the distribution has no protoc package. prost-build needs a
# protoc on the PATH to compile the Comet protobuf definitions.
PROTOC_VERSION = "25.5"
RESULTS_SUBDIR = Path("benchmarks") / "results" / "micro"
IMDS_BASE = "http://169.254.169.254/latest"

# Suites that legitimately take much longer than the default timeout allows.
# Minutes, applied in place of --timeout for that suite only.
SUITE_TIMEOUT_OVERRIDES = {
    "CometShuffleBenchmark": 180,
}

# Suites deliberately left out of the default set, with the reason. They can
# still be run by naming them in a --suites file. Everything else found in
# BENCH_SOURCE_DIR is run, so a new benchmark needs no change here.
EXCLUDED_SUITES = {
    "CometTPCHQueryBenchmark": "needs TPC-H data, pass --data-location",
    "CometTPCDSQueryBenchmark": "needs TPC-DS data, pass --data-location",
    "CometTPCDSMicroBenchmark": "needs TPC-DS data, pass --data-location",
    "CometC2RIsolatedBench": "prints to stdout only, writes no results file",
    "CometReadHdfsBenchmark": "starts a local HDFS mini cluster",
}


# ---------------------------------------------------------------------------
# small helpers
# ---------------------------------------------------------------------------


def log(message):
    stamp = datetime.now().strftime("%H:%M:%S")
    print(f"[{stamp}] {message}", flush=True)


def fail(message):
    print(f"error: {message}", file=sys.stderr)
    sys.exit(1)


def utc_now():
    return datetime.now(timezone.utc)


def format_duration(seconds):
    minutes, seconds = divmod(int(seconds), 60)
    hours, minutes = divmod(minutes, 60)
    if hours:
        return f"{hours}h{minutes:02d}m{seconds:02d}s"
    if minutes:
        return f"{minutes}m{seconds:02d}s"
    return f"{seconds}s"


def run_command(cmd, cwd=None, env=None, log_path=None, timeout=None, check=True, dry_run=False):
    """Run a command, optionally sending its output to a log file."""
    printable = " ".join(str(part) for part in cmd)
    if dry_run:
        log(f"[dry-run] {printable}")
        return 0
    log(printable)
    if log_path is None:
        completed = subprocess.run(cmd, cwd=cwd, env=env, timeout=timeout)
    else:
        with open(log_path, "w") as handle:
            handle.write(f"$ {printable}\n\n")
            handle.flush()
            completed = subprocess.run(
                cmd, cwd=cwd, env=env, stdout=handle, stderr=subprocess.STDOUT, timeout=timeout
            )
    if check and completed.returncode != 0:
        fail(f"command failed with exit code {completed.returncode}: {printable}")
    return completed.returncode


def capture(cmd, cwd=None, env=None):
    """Run a command and return stripped stdout, or None if it fails."""
    try:
        completed = subprocess.run(
            cmd, cwd=cwd, env=env, capture_output=True, text=True, timeout=300
        )
    except (OSError, subprocess.SubprocessError):
        return None
    if completed.returncode != 0:
        return None
    return completed.stdout.strip()


def sudo_prefix():
    if os.geteuid() == 0:
        return []
    if shutil.which("sudo") is None:
        fail("this step needs root privileges but sudo was not found")
    return ["sudo"]


def tail_file(path, lines=20):
    try:
        content = Path(path).read_text(errors="replace").splitlines()
    except OSError:
        return ""
    return "\n".join(content[-lines:])


def read_proc_field(path, prefix):
    """The value after the colon of the first line in `path` starting with `prefix`."""
    try:
        for line in Path(path).read_text().splitlines():
            if line.startswith(prefix):
                return line.split(":", 1)[1].strip()
    except (OSError, IndexError):
        pass
    return None


# The last lines of a failed run are Maven's epilogue, which says nothing about
# the cause. Look for the exception that started it instead.
EXCEPTION_PATTERN = re.compile(r"^\s*[\w.$]+(Exception|Error)(:|\b)")
MAVEN_FAILURE_PATTERN = re.compile(r"^\[ERROR\] Failed to execute goal")


def failure_excerpt(path, context=12, fallback_lines=20):
    """The first exception in a log, falling back to Maven's message or the tail."""
    try:
        content = Path(path).read_text(errors="replace").splitlines()
    except OSError:
        return ""
    for pattern in (EXCEPTION_PATTERN, MAVEN_FAILURE_PATTERN):
        for index, line in enumerate(content):
            if pattern.search(line):
                return "\n".join(content[index : index + context])
    return "\n".join(content[-fallback_lines:])


# ---------------------------------------------------------------------------
# environment
# ---------------------------------------------------------------------------


def find_checkout_from_script():
    """Return the Comet checkout this script lives in, if any."""
    for parent in Path(__file__).resolve().parents:
        if (parent / "Makefile").is_file() and (parent / "native" / "Cargo.toml").is_file():
            return parent
    return None


def resolve_comet_home(args):
    if args.comet_home:
        return Path(args.comet_home).expanduser().resolve()
    if os.environ.get("COMET_HOME"):
        return Path(os.environ["COMET_HOME"]).expanduser().resolve()
    checkout = find_checkout_from_script()
    if checkout is not None:
        return checkout
    return DEFAULT_COMET_HOME


def require_comet_home(args):
    comet_home = resolve_comet_home(args)
    if not (comet_home / "Makefile").is_file():
        fail(f"{comet_home} is not a Comet checkout, run `setup` first or pass --comet-home")
    return comet_home


def detect_java_home():
    if os.environ.get("JAVA_HOME"):
        return os.environ["JAVA_HOME"]
    candidates = []
    for pattern in ("java-*-amazon-corretto*", "java-*-openjdk*", "jdk-*"):
        candidates.extend(Path("/usr/lib/jvm").glob(pattern))

    def version_of(path):
        match = re.search(r"(\d+)", path.name)
        return int(match.group(1)) if match else 0

    # Newest first, so that an older JDK left on the machine is not picked up
    ordered = sorted(candidates, key=version_of, reverse=True)
    for candidate in ordered:
        if (candidate / "bin" / "javac").is_file():
            return str(candidate)
    java = shutil.which("java")
    if java is not None:
        resolved = Path(java).resolve()
        return str(resolved.parent.parent)
    return None


def base_env(comet_home):
    """Environment shared by the build and the benchmark runs."""
    env = dict(os.environ)
    cargo_bin = Path.home() / ".cargo" / "bin"
    if cargo_bin.is_dir() and str(cargo_bin) not in env.get("PATH", ""):
        env["PATH"] = f"{cargo_bin}{os.pathsep}{env.get('PATH', '')}"
    java_home = detect_java_home()
    if java_home:
        env["JAVA_HOME"] = java_home
        java_bin = Path(java_home) / "bin"
        if str(java_bin) not in env.get("PATH", ""):
            env["PATH"] = f"{java_bin}{os.pathsep}{env.get('PATH', '')}"
    env["COMET_CONF_DIR"] = str(comet_home / "conf")
    return env


def java_major_version(env):
    """Major version of the JDK the build will use, or None if unknown."""
    javac = Path(env["JAVA_HOME"]) / "bin" / "javac" if env.get("JAVA_HOME") else "javac"
    try:
        completed = subprocess.run(
            [str(javac), "-version"], capture_output=True, text=True, env=env, timeout=60
        )
    except (OSError, subprocess.SubprocessError):
        return None
    match = re.search(r"(\d+)", (completed.stdout or "") + (completed.stderr or ""))
    return int(match.group(1)) if match else None


def require_supported_jdk(env):
    """Spark 4.x needs JDK 17, and javac silently stubs java.lang.Record below it."""
    major = java_major_version(env)
    if major is None:
        fail(f"could not run javac from JAVA_HOME={env.get('JAVA_HOME')}, install a JDK")
    if major < MINIMUM_JDK:
        fail(
            f"JDK {major} found at JAVA_HOME={env.get('JAVA_HOME')}, but the default Spark 4.x "
            f"profile needs JDK {MINIMUM_JDK} or later. Compiling with an older JDK fails with "
            "'Class java.lang.Record not found'. Install a newer JDK, point JAVA_HOME at it, and "
            "run `./mvnw clean` to discard classes built against the old one."
        )


# ---------------------------------------------------------------------------
# setup
# ---------------------------------------------------------------------------


def install_packages(args):
    dnf_packages = [
        "git",
        "make",
        "cmake",
        "gcc",
        "gcc-c++",
        "protobuf-compiler",
        "python3",
        f"java-{args.jdk}-amazon-corretto-devel",
    ]
    apt_packages = [
        "git",
        "make",
        "cmake",
        "build-essential",
        "protobuf-compiler",
        "python3",
        "curl",
        f"openjdk-{args.jdk}-jdk",
    ]
    if shutil.which("dnf"):
        run_command(sudo_prefix() + ["dnf", "install", "-y"] + dnf_packages, dry_run=args.dry_run)
    elif shutil.which("yum"):
        run_command(sudo_prefix() + ["yum", "install", "-y"] + dnf_packages, dry_run=args.dry_run)
    elif shutil.which("apt-get"):
        run_command(sudo_prefix() + ["apt-get", "update"], dry_run=args.dry_run)
        run_command(
            sudo_prefix() + ["apt-get", "install", "-y"] + apt_packages, dry_run=args.dry_run
        )
    else:
        log("no supported package manager found, skipping package installation")
    install_protoc(args)


def protoc_asset_name():
    machine = platform.machine().lower()
    if platform.system() == "Darwin":
        return f"protoc-{PROTOC_VERSION}-osx-universal_binary.zip"
    if machine in ("aarch64", "arm64"):
        return f"protoc-{PROTOC_VERSION}-linux-aarch_64.zip"
    return f"protoc-{PROTOC_VERSION}-linux-x86_64.zip"


def install_protoc(args):
    """Install protoc from the protobuf releases when the distribution has none.

    Only the Rust build needs a protoc on the PATH. The JVM side downloads its
    own through protoc-jar-maven-plugin.
    """
    if shutil.which("protoc"):
        log(f"protoc already installed: {capture(['protoc', '--version'])}")
        return
    asset = protoc_asset_name()
    url = f"https://github.com/protocolbuffers/protobuf/releases/download/v{PROTOC_VERSION}/{asset}"
    if args.dry_run:
        log(f"[dry-run] download {url} and install bin/protoc into /usr/local/bin")
        return
    log(f"protoc not found, installing {PROTOC_VERSION} from {url}")
    with tempfile.TemporaryDirectory() as work_dir:
        archive = Path(work_dir) / asset
        try:
            with urllib.request.urlopen(url, timeout=120) as response:
                archive.write_bytes(response.read())
            with zipfile.ZipFile(archive) as bundle:
                bundle.extract("bin/protoc", work_dir)
        except (OSError, urllib.error.URLError, zipfile.BadZipFile) as error:
            fail(f"could not install protoc: {error}")
        run_command(
            sudo_prefix() + ["install", "-m", "755", f"{work_dir}/bin/protoc", "/usr/local/bin/"]
        )
    log(f"installed {capture(['/usr/local/bin/protoc', '--version'])}")


def install_rust(args):
    if shutil.which("cargo") or (Path.home() / ".cargo" / "bin" / "cargo").is_file():
        log("rust toolchain already installed")
        return
    log("installing rust toolchain")
    if args.dry_run:
        log("[dry-run] curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y")
        return
    installer = subprocess.run(
        ["curl", "--proto", "=https", "--tlsv1.2", "-sSf", "https://sh.rustup.rs"],
        capture_output=True,
    )
    if installer.returncode != 0:
        fail("failed to download rustup")
    shell = subprocess.run(["sh", "-s", "--", "-y"], input=installer.stdout)
    if shell.returncode != 0:
        fail("rustup installation failed")


def clone_or_update(args, comet_home):
    # .git is a directory in a normal clone and a file in a worktree
    if not (comet_home / ".git").exists():
        comet_home.parent.mkdir(parents=True, exist_ok=True)
        run_command(
            ["git", "clone", args.repo, str(comet_home)],
            dry_run=args.dry_run,
        )
        ref = args.ref or "main"
    else:
        log(f"using existing checkout at {comet_home}")
        ref = args.ref
    if ref:
        run_command(["git", "fetch", "origin", "--tags"], cwd=comet_home, dry_run=args.dry_run)
        run_command(["git", "checkout", ref], cwd=comet_home, dry_run=args.dry_run)
        # Fast-forward when the ref is a branch. Detached heads and tags fail
        # here, which is fine.
        run_command(
            ["git", "merge", "--ff-only", f"origin/{ref}"],
            cwd=comet_home,
            check=False,
            dry_run=args.dry_run,
        )


def build_release(args, comet_home):
    env = base_env(comet_home)
    if env.get("JAVA_HOME") is None:
        fail("JAVA_HOME could not be determined, install a JDK or export JAVA_HOME")
    require_supported_jdk(env)
    if args.profile:
        env["PROFILES"] = args.profile
    log("building Comet in release mode, this takes a while")
    run_command(["make", "release"], cwd=comet_home, env=env, dry_run=args.dry_run)


def cmd_setup(args):
    comet_home = resolve_comet_home(args)
    if not args.skip_packages:
        install_packages(args)
        install_rust(args)
    clone_or_update(args, comet_home)
    if not args.skip_build:
        build_release(args, comet_home)
    log(f"setup complete: {comet_home}")
    return comet_home


# ---------------------------------------------------------------------------
# run
# ---------------------------------------------------------------------------


# Top-level `object Name`, with whatever it extends up to the start of its body
OBJECT_PATTERN = re.compile(r"^object\s+(\w+)\s*(extends[^{]*)?\{", re.MULTILINE)


def discover_suites(comet_home):
    """The benchmark entry points, read from the sources rather than listed here.

    An entry point is a top-level `object` in the benchmark package that either
    mixes in a `*Benchmark*` base, which supplies `main`, or declares its own.
    Helpers such as `TPCDSSchemaHelper` do neither and are skipped. Discovering
    them this way means a newly added benchmark runs without anyone having to
    remember to also add it to this file.
    """
    source_dir = comet_home / BENCH_SOURCE_DIR
    if not source_dir.is_dir():
        fail(f"no benchmark sources at {source_dir}, is {comet_home} a Comet checkout?")
    suites = set()
    for path in sorted(source_dir.glob("*.scala")):
        text = path.read_text(errors="replace")
        declares_main = "def main(" in text
        for name, bases in OBJECT_PATTERN.findall(text):
            if "Benchmark" in bases or declares_main:
                suites.add(name)
    if not suites:
        fail(f"no benchmark objects found in {source_dir}")
    return sorted(suites)


def stale_suite_entries(comet_home):
    """Names in this file that no longer match the benchmark sources.

    EXCLUDED_SUITES and SUITE_TIMEOUT_OVERRIDES are keyed by suite name, so a
    rename upstream silently turns an exclusion into a suite that runs, or a
    timeout override into one that does not apply. Report them instead.
    """
    discovered = set(discover_suites(comet_home))
    return [
        f"{suite} is listed in {label} but is not a benchmark in {BENCH_SOURCE_DIR}"
        for label, names in (
            ("EXCLUDED_SUITES", EXCLUDED_SUITES),
            ("SUITE_TIMEOUT_OVERRIDES", SUITE_TIMEOUT_OVERRIDES),
        )
        for suite in sorted(set(names) - discovered)
    ]


def load_suites(args, comet_home):
    if args.suites:
        path = Path(args.suites).expanduser()
        if not path.is_file():
            fail(f"suite file not found: {path}")
        suites = []
        for line in path.read_text().splitlines():
            line = line.split("#", 1)[0].strip()
            if line:
                suites.append(line.rsplit(".", 1)[-1])
    else:
        suites = [s for s in discover_suites(comet_home) if s not in EXCLUDED_SUITES]
    if args.only:
        suites = [s for s in suites if any(re.search(pattern, s) for pattern in args.only)]
    if args.skip:
        suites = [s for s in suites if not any(re.search(pattern, s) for pattern in args.skip)]
    if not suites:
        fail("no suites selected")
    return suites


def benchmark_args(comet_home, env, profile, heap):
    """The benchmark invocation, read from the Makefile's `print-benchmark-args`.

    The Makefile is the one definition of how a suite is invoked. The
    `benchmark-%` target itself cannot be used here because its `release`
    prerequisite would rebuild Comet before every suite.
    """
    command = ["make", "print-benchmark-args", f"BENCH_HEAP={heap}"]
    if profile:
        command.append(f"PROFILES={profile}")
    output = capture(command, cwd=comet_home, env=env)
    if output is None:
        fail(f"`{' '.join(command)}` failed in {comet_home}, is make installed?")
    values = {}
    for line in output.splitlines():
        key, separator, value = line.partition("=")
        if separator:
            values[key.strip()] = value.strip()
    missing = [key for key in ("MAVEN_OPTS", "MVN_ARGS") if not values.get(key)]
    if missing:
        fail(f"print-benchmark-args did not report {', '.join(missing)}")
    return values


def suite_command(suite, mvn_args, profile):
    """The `mvnw` invocation for one suite, run from the `spark` directory."""
    command = ["../mvnw"] + mvn_args.split() + [f"-Dexec.mainClass={BENCH_PACKAGE}.{suite}"]
    if profile:
        command.append(profile)
    return command


def warn_unless_release_build(comet_home):
    release_dir = comet_home / "native" / "target" / "release"
    if any((release_dir / name).is_file() for name in ("libcomet.so", "libcomet.dylib")):
        return
    log(
        "warning: no native release build found in native/target/release. "
        "Run `setup` or `make release` first, otherwise the numbers will not reflect "
        "a release build."
    )


def git_info(comet_home):
    return {
        "commit": capture(["git", "rev-parse", "HEAD"], cwd=comet_home),
        "commit_short": capture(["git", "rev-parse", "--short", "HEAD"], cwd=comet_home),
        "branch": capture(["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=comet_home),
        "describe": capture(["git", "log", "-1", "--pretty=%s"], cwd=comet_home),
    }


def cmd_run(args, comet_home=None):
    comet_home = comet_home or require_comet_home(args)
    suites = load_suites(args, comet_home)
    for problem in stale_suite_entries(comet_home):
        log(f"warning: {problem}")
    if args.list:
        for suite in suites:
            print(f"{BENCH_PACKAGE}.{suite}")
        for suite, reason in sorted(EXCLUDED_SUITES.items()):
            print(f"# {BENCH_PACKAGE}.{suite} ({reason})")
        return None

    warn_unless_release_build(comet_home)
    runs_root = Path(args.runs_root).expanduser()
    run_dir = runs_root / utc_now().strftime("%Y%m%dT%H%M%SZ")
    log_dir = run_dir / "logs"
    log_dir.mkdir(parents=True, exist_ok=True)

    env = base_env(comet_home)
    if env.get("JAVA_HOME") is None:
        fail("JAVA_HOME could not be determined, install a JDK or export JAVA_HOME")
    require_supported_jdk(env)
    bench = benchmark_args(comet_home, env, args.profile, args.heap)
    env["MAVEN_OPTS"] = bench["MAVEN_OPTS"]
    env["SPARK_GENERATE_BENCHMARK_FILES"] = "1"

    spark_dir = comet_home / "spark"
    started = time.time()
    started_utc = utc_now()
    log(f"running {len(suites)} suite(s), logs in {log_dir}")
    results = []

    for index, suite in enumerate(suites, start=1):
        log_path = log_dir / f"{suite}.log"
        timeout_minutes = SUITE_TIMEOUT_OVERRIDES.get(suite, args.timeout)
        timeout = timeout_minutes * 60 if timeout_minutes else None
        if args.timeout and timeout_minutes != args.timeout:
            log(f"[{index}/{len(suites)}] {suite} (timeout raised to {timeout_minutes}m)")
        else:
            log(f"[{index}/{len(suites)}] {suite}")
        suite_started = time.time()
        status = "ok"
        try:
            code = run_command(
                suite_command(suite, bench["MVN_ARGS"], args.profile),
                cwd=spark_dir,
                env=env,
                log_path=log_path,
                timeout=timeout,
                check=False,
                dry_run=args.dry_run,
            )
            if code != 0:
                status = "failed"
        except subprocess.TimeoutExpired:
            status = "timeout"
        duration = time.time() - suite_started
        results.append(
            {
                "suite": suite,
                "status": status,
                "duration_seconds": round(duration, 1),
                "log": str(log_path),
            }
        )
        if status == "ok":
            log(f"    ok in {format_duration(duration)}")
        else:
            excerpt = tail_file(log_path) if status == "timeout" else failure_excerpt(log_path)
            log(f"    {status} after {format_duration(duration)}, from {log_path}:")
            print(excerpt, flush=True)

    summary = {
        "started_utc": started_utc.isoformat(),
        "finished_utc": utc_now().isoformat(),
        "duration_seconds": round(time.time() - started, 1),
        "comet_home": str(comet_home),
        "profile": args.profile or "default",
        "heap": args.heap,
        "git": git_info(comet_home),
        "suites": results,
    }
    summary_path = run_dir / "summary.json"
    summary_path.write_text(json.dumps(summary, indent=2) + "\n")

    failed = [entry for entry in results if entry["status"] != "ok"]
    log(f"finished in {format_duration(summary['duration_seconds'])}, summary: {summary_path}")
    for entry in results:
        marker = " " if entry["status"] == "ok" else "!"
        log(
            f"  {marker} {entry['suite']:<45} {entry['status']:<8} "
            f"{format_duration(entry['duration_seconds'])}"
        )
    if failed:
        log(f"{len(failed)} suite(s) did not complete successfully")
    return run_dir


# ---------------------------------------------------------------------------
# collect
# ---------------------------------------------------------------------------


def imds(path):
    """Read an EC2 instance metadata value, returning None when not on EC2."""
    try:
        token_request = urllib.request.Request(
            f"{IMDS_BASE}/api/token",
            method="PUT",
            headers={"X-aws-ec2-metadata-token-ttl-seconds": "60"},
        )
        token = urllib.request.urlopen(token_request, timeout=2).read().decode()
        value_request = urllib.request.Request(
            f"{IMDS_BASE}/meta-data/{path}",
            headers={"X-aws-ec2-metadata-token": token},
        )
        return urllib.request.urlopen(value_request, timeout=2).read().decode()
    except Exception:
        return None


def machine_info(env):
    info = {
        "instance_type": imds("instance-type") or "unknown (not an EC2 instance?)",
        "availability_zone": imds("placement/availability-zone") or "unknown",
        "cpus": str(os.cpu_count()),
        "cpu_model": "unknown",
        "memory": "unknown",
        "os": platform.platform(),
        "java": "unknown",
        "rustc": capture(["rustc", "--version"], env=env) or "unknown",
    }
    cpu_model = read_proc_field("/proc/cpuinfo", "model name")
    if cpu_model:
        info["cpu_model"] = cpu_model
    mem_total = read_proc_field("/proc/meminfo", "MemTotal")
    if mem_total:
        # "16311236 kB"
        info["memory"] = f"{int(mem_total.split()[0]) / 1024 / 1024:.1f} GiB"
    if info["cpu_model"] == "unknown" and platform.system() == "Darwin":
        # /proc does not exist on macOS, which is where a local collect runs
        info["cpu_model"] = capture(["sysctl", "-n", "machdep.cpu.brand_string"]) or "unknown"
        memory_bytes = capture(["sysctl", "-n", "hw.memsize"])
        if memory_bytes:
            info["memory"] = f"{int(memory_bytes) / 1024**3:.1f} GiB"
    # `java -version` writes to stderr rather than stdout
    java_version = None
    try:
        completed = subprocess.run(
            ["java", "-version"], capture_output=True, text=True, env=env, timeout=60
        )
        java_version = (completed.stderr or completed.stdout).strip().splitlines()[0]
    except (OSError, subprocess.SubprocessError, IndexError):
        pass
    info["java"] = java_version or "unknown"
    return info


def latest_run_dir(runs_root):
    runs_root = Path(runs_root).expanduser()
    if not runs_root.is_dir():
        return None
    candidates = sorted(path for path in runs_root.iterdir() if (path / "summary.json").is_file())
    return candidates[-1] if candidates else None


LICENSE_HEADER = """<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied.  See the License for the
specific language governing permissions and limitations
under the License.
-->
"""


def write_run_info(destination, summary, info, copied):
    git = summary.get("git", {})
    lines = [
        LICENSE_HEADER,
        "# Comet Micro Benchmark Run",
        "",
        "Generated by `benchmarks/micro/run.py collect`.",
        "",
        "## Environment",
        "",
        "| Field | Value |",
        "| --- | --- |",
        f"| Date (UTC) | {summary.get('started_utc', 'unknown')} |",
        f"| Instance type | {info['instance_type']} |",
        f"| Availability zone | {info['availability_zone']} |",
        f"| CPU | {info['cpu_model']} ({info['cpus']} vCPU) |",
        f"| Memory | {info['memory']} |",
        f"| OS | {info['os']} |",
        f"| Java | {info['java']} |",
        f"| Rust | {info['rustc']} |",
        f"| Maven profile | {summary.get('profile', 'default')} |",
        f"| Benchmark heap | {summary.get('heap', 'unknown')} |",
        f"| Comet branch | {git.get('branch', 'unknown')} |",
        f"| Comet commit | {git.get('commit', 'unknown')} |",
        f"| Total run time | {format_duration(summary.get('duration_seconds', 0))} |",
        "",
        "## Suites",
        "",
        "| Suite | Status | Duration | Results file |",
        "| --- | --- | --- | --- |",
    ]
    for entry in summary.get("suites", []):
        results_file = copied.get(entry["suite"], "-")
        lines.append(
            f"| {entry['suite']} | {entry['status']} | "
            f"{format_duration(entry['duration_seconds'])} | {results_file} |"
        )
    lines.append("")
    (destination / "RUN-INFO.md").write_text("\n".join(lines))


def cmd_collect(args, comet_home=None, run_dir=None):
    comet_home = comet_home or require_comet_home(args)
    if run_dir is None:
        run_dir = (
            Path(args.run_dir).expanduser() if args.run_dir else latest_run_dir(args.runs_root)
        )
    if run_dir is None or not (run_dir / "summary.json").is_file():
        fail("no benchmark run found, pass --run-dir or run `run` first")
    summary = json.loads((run_dir / "summary.json").read_text())

    source = comet_home / "spark" / "benchmarks"
    if not source.is_dir():
        fail(f"no results directory at {source}, did the suites run?")
    destination = comet_home / RESULTS_SUBDIR
    destination.mkdir(parents=True, exist_ok=True)

    started = datetime.fromisoformat(summary["started_utc"]).timestamp()
    copied = {}
    for path in sorted(source.glob("*.txt")):
        if not args.all_results and path.stat().st_mtime < started:
            continue
        shutil.copy2(path, destination / path.name)
        suite = path.name.split("-")[0]
        copied[suite] = path.name
        log(f"collected {path.name}")
    if not copied:
        log("no results files were produced by this run")

    info = machine_info(base_env(comet_home))
    write_run_info(destination, summary, info, copied)
    log(f"wrote {destination / 'RUN-INFO.md'}")
    log(f"results are in {destination}")
    if args.publish:
        cmd_publish(args, comet_home=comet_home)
    return destination


# ---------------------------------------------------------------------------
# publish
# ---------------------------------------------------------------------------


PR_BODY_TEMPLATE = """## Which issue does this PR close?

N/A

## Rationale for this change

Refresh the checked-in micro benchmark results so that changes in expression and
operator performance are visible over time.

## What changes are included in this PR?

Micro benchmark results produced by `benchmarks/micro/run.py` on a {instance_type}
instance, from Comet commit {commit}.

## How are these changes tested?

These are benchmark result files only, produced by running the micro benchmark
suites. See `benchmarks/results/micro/RUN-INFO.md` for the environment details.
"""


def cmd_publish(args, comet_home=None):
    comet_home = comet_home or require_comet_home(args)
    destination = comet_home / RESULTS_SUBDIR
    if not destination.is_dir():
        fail(f"no collected results at {destination}, run `collect` first")

    run_info = destination / "RUN-INFO.md"
    instance_type = "unknown"
    if run_info.is_file():
        for line in run_info.read_text().splitlines():
            if line.startswith("| Instance type |"):
                instance_type = line.split("|")[2].strip()
                break
    commit = git_info(comet_home).get("commit_short") or "unknown"
    branch = args.branch or f"benchmark-results-{utc_now().strftime('%Y%m%d')}"
    title = "chore: add micro benchmark results"
    if not instance_type.startswith("unknown"):
        title = f"{title} from {instance_type}"

    run_command(["git", "checkout", "-b", branch], cwd=comet_home, dry_run=args.dry_run)
    run_command(
        ["git", "add", str(RESULTS_SUBDIR)],
        cwd=comet_home,
        dry_run=args.dry_run,
    )
    run_command(
        ["git", "commit", "-m", title],
        cwd=comet_home,
        dry_run=args.dry_run,
    )

    if not args.push:
        log("commit created, push it and open a pull request with:")
        print(f"  git -C {comet_home} push -u {args.remote} {branch}")
        print(f'  gh pr create --base {args.base} --title "{title}"')
        return

    run_command(
        ["git", "push", "-u", args.remote, branch],
        cwd=comet_home,
        dry_run=args.dry_run,
    )
    if not args.open_pr:
        log("branch pushed, open the pull request when ready")
        return
    if shutil.which("gh") is None:
        log("gh is not installed, open the pull request manually")
        return
    body_path = comet_home / ".git" / "comet-benchmark-pr-body.md"
    body_path.write_text(PR_BODY_TEMPLATE.format(instance_type=instance_type, commit=commit))
    run_command(
        [
            "gh",
            "pr",
            "create",
            "--base",
            args.base,
            "--title",
            title,
            "--body-file",
            str(body_path),
        ],
        cwd=comet_home,
        dry_run=args.dry_run,
    )


# ---------------------------------------------------------------------------
# all
# ---------------------------------------------------------------------------


def cmd_all(args):
    comet_home = cmd_setup(args)
    run_dir = cmd_run(args, comet_home=comet_home)
    if run_dir is None:
        return
    cmd_collect(args, comet_home=comet_home, run_dir=run_dir)


# ---------------------------------------------------------------------------
# argument parsing
# ---------------------------------------------------------------------------


def add_common_arguments(parser):
    parser.add_argument(
        "--comet-home",
        help="Comet checkout to use (default: the checkout this script lives in, "
        f"otherwise {DEFAULT_COMET_HOME})",
    )
    parser.add_argument("--profile", help="Maven profile to build and run with, e.g. -Pspark-3.5")
    parser.add_argument(
        "--dry-run", action="store_true", help="print the commands without running them"
    )


def add_setup_arguments(parser):
    parser.add_argument("--repo", default=DEFAULT_REPO, help="repository to clone")
    parser.add_argument("--ref", help="branch, tag or commit to check out")
    parser.add_argument("--jdk", default=DEFAULT_JDK, help="JDK major version to install")
    parser.add_argument(
        "--skip-packages", action="store_true", help="do not install system packages or Rust"
    )
    parser.add_argument("--skip-build", action="store_true", help="do not run `make release`")


def add_run_arguments(parser):
    parser.add_argument("--suites", help="file listing suites to run, one per line")
    parser.add_argument(
        "--only", action="append", metavar="REGEX", help="only run suites matching this pattern"
    )
    parser.add_argument(
        "--skip", action="append", metavar="REGEX", help="skip suites matching this pattern"
    )
    parser.add_argument("--heap", default=DEFAULT_HEAP, help="JVM max heap for each suite")
    parser.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_TIMEOUT_MINUTES,
        help="per-suite timeout in minutes, 0 to disable",
    )
    parser.add_argument("--list", action="store_true", help="list the selected suites and exit")


def add_runs_root_argument(parser):
    parser.add_argument(
        "--runs-root", default=str(DEFAULT_RUNS_ROOT), help="directory holding per-run logs"
    )


def add_collect_arguments(parser):
    parser.add_argument("--run-dir", help="run directory to collect (default: most recent)")
    parser.add_argument(
        "--all-results",
        action="store_true",
        help="collect every results file, not only those written by this run",
    )
    parser.add_argument(
        "--publish", action="store_true", help="commit the results after collecting them"
    )


def add_publish_arguments(parser):
    parser.add_argument("--branch", help="branch name to create (default: benchmark-results-DATE)")
    parser.add_argument("--remote", default="origin", help="git remote to push to")
    parser.add_argument("--base", default="main", help="base branch for the pull request")
    parser.add_argument("--push", action="store_true", help="push the branch to the remote")
    parser.add_argument(
        "--open-pr", action="store_true", help="open a pull request with gh after pushing"
    )


def parse_arguments(argv):
    parser = argparse.ArgumentParser(
        description="Run the Comet micro benchmark suites and publish the results.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    setup = subparsers.add_parser("setup", help="install prerequisites and build Comet")
    add_common_arguments(setup)
    add_setup_arguments(setup)
    setup.set_defaults(func=cmd_setup)

    run = subparsers.add_parser("run", help="run the benchmark suites")
    add_common_arguments(run)
    add_run_arguments(run)
    add_runs_root_argument(run)
    run.set_defaults(func=cmd_run)

    collect = subparsers.add_parser("collect", help="copy results into benchmarks/results/micro")
    add_common_arguments(collect)
    add_collect_arguments(collect)
    add_runs_root_argument(collect)
    add_publish_arguments(collect)
    collect.set_defaults(func=cmd_collect)

    publish = subparsers.add_parser("publish", help="commit and optionally open a pull request")
    add_common_arguments(publish)
    add_publish_arguments(publish)
    publish.set_defaults(func=cmd_publish)

    every = subparsers.add_parser("all", help="setup, run and collect in one go")
    add_common_arguments(every)
    add_setup_arguments(every)
    add_run_arguments(every)
    add_collect_arguments(every)
    add_runs_root_argument(every)
    add_publish_arguments(every)
    every.set_defaults(func=cmd_all)

    return parser.parse_args(argv)


def main(argv=None):
    args = parse_arguments(argv if argv is not None else sys.argv[1:])
    if getattr(args, "timeout", None) == 0:
        args.timeout = None
    args.func(args)


if __name__ == "__main__":
    main()