#!/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.

"""Consolidated TPC benchmark runner for Spark-based engines.

Usage:
    python3 run.py --engine comet --benchmark tpch
    python3 run.py --engine comet --benchmark tpcds --iterations 3
    python3 run.py --engine comet-iceberg --benchmark tpch
    python3 run.py --engine comet --benchmark tpch --dry-run
    python3 run.py --engine spark --benchmark tpch --no-restart
"""

import argparse
import os
import re
import subprocess
import sys

# ---------------------------------------------------------------------------
# TOML loading – prefer stdlib tomllib (3.11+), else minimal fallback
# ---------------------------------------------------------------------------

try:
    import tomllib  # Python 3.11+

    def load_toml(path):
        with open(path, "rb") as f:
            return tomllib.load(f)

except ModuleNotFoundError:

    def _parse_toml(text):
        """Minimal TOML parser supporting tables, quoted-key strings, plain
        strings, arrays of strings, booleans, and comments.  Sufficient for
        the engine config files used by this runner."""
        root = {}
        current = root
        for line in text.splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            # Table header: [env.defaults] or [spark_conf]
            m = re.match(r"^\[([^\]]+)\]$", line)
            if m:
                keys = m.group(1).split(".")
                current = root
                for k in keys:
                    current = current.setdefault(k, {})
                continue
            # Key = value
            m = re.match(r'^("(?:[^"\\]|\\.)*"|[A-Za-z0-9_.]+)\s*=\s*(.+)$', line)
            if not m:
                continue
            raw_key, raw_val = m.group(1), m.group(2).strip()
            key = raw_key.strip('"')
            val = _parse_value(raw_val)
            current[key] = val
        return root

    def _parse_value(raw):
        if raw == "true":
            return True
        if raw == "false":
            return False
        if raw.startswith('"') and raw.endswith('"'):
            return raw[1:-1]
        if raw.startswith("["):
            # Simple array of strings
            items = []
            for m in re.finditer(r'"((?:[^"\\]|\\.)*)"', raw):
                items.append(m.group(1))
            return items
        if raw.startswith("{"):
            # Inline table: { KEY = "VAL", ... }
            tbl = {}
            for m in re.finditer(r'([A-Za-z0-9_]+)\s*=\s*"((?:[^"\\]|\\.)*)"', raw):
                tbl[m.group(1)] = m.group(2)
            return tbl
        return raw

    def load_toml(path):
        with open(path, "r") as f:
            return _parse_toml(f.read())


# ---------------------------------------------------------------------------
# Common Spark configuration (shared across all engines)
# ---------------------------------------------------------------------------

COMMON_SPARK_CONF = {
    "spark.driver.memory": "8G",
    "spark.executor.memory": "16g",
    "spark.memory.offHeap.enabled": "true",
    "spark.memory.offHeap.size": "16g",
    "spark.eventLog.enabled": "true",
    "spark.eventLog.dir": os.environ.get("SPARK_EVENT_LOG_DIR", "/tmp/spark-events"),
    "spark.hadoop.fs.s3a.impl": "org.apache.hadoop.fs.s3a.S3AFileSystem",
    "spark.hadoop.fs.s3a.aws.credentials.provider": "com.amazonaws.auth.DefaultAWSCredentialsProviderChain",
}

# ---------------------------------------------------------------------------
# Benchmark profiles
# ---------------------------------------------------------------------------

BENCHMARK_PROFILES = {
    "tpch": {
        "executor_instances": "2",
        "executor_cores": "8",
        "max_cores": "16",
        "data_env": "TPCH_DATA",
        "format": "parquet",
    },
    "tpcds": {
        "executor_instances": "2",
        "executor_cores": "8",
        "max_cores": "16",
        "data_env": "TPCDS_DATA",
        "format": None,  # omit --format for TPC-DS
    },
}

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def resolve_env(value):
    """Expand $VAR and ${VAR} references using os.environ."""
    if not isinstance(value, str):
        return value
    return re.sub(
        r"\$\{([^}]+)\}|\$([A-Za-z_][A-Za-z0-9_]*)",
        lambda m: os.environ.get(m.group(1) or m.group(2), ""),
        value,
    )


def resolve_env_in_list(lst):
    return [resolve_env(v) for v in lst]


def load_engine_config(engine_name):
    """Load and return the TOML config for the given engine."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = os.path.join(script_dir, "engines", f"{engine_name}.toml")
    if not os.path.exists(config_path):
        available = sorted(
            f.removesuffix(".toml")
            for f in os.listdir(os.path.join(script_dir, "engines"))
            if f.endswith(".toml")
        )
        print(f"Error: Unknown engine '{engine_name}'", file=sys.stderr)
        print(f"Available engines: {', '.join(available)}", file=sys.stderr)
        sys.exit(1)
    return load_toml(config_path)


def apply_env_defaults(config):
    """Set environment variable defaults from [env.defaults]."""
    defaults = config.get("env", {}).get("defaults", {})
    for key, val in defaults.items():
        if key not in os.environ:
            os.environ[key] = val


def apply_env_exports(config):
    """Export environment variables from [env.exports]."""
    exports = config.get("env", {}).get("exports", {})
    for key, val in exports.items():
        os.environ[key] = val


def check_required_env(config):
    """Validate that required environment variables are set."""
    required = config.get("env", {}).get("required", [])
    missing = [v for v in required if not os.environ.get(v)]
    if missing:
        print(
            f"Error: Required environment variable(s) not set: {', '.join(missing)}",
            file=sys.stderr,
        )
        sys.exit(1)


def check_common_env():
    """Validate SPARK_HOME and SPARK_MASTER are set."""
    for var in ("SPARK_HOME", "SPARK_MASTER"):
        if not os.environ.get(var):
            print(f"Error: {var} is not set", file=sys.stderr)
            sys.exit(1)


def check_benchmark_env(config, benchmark):
    """Validate benchmark-specific environment variables."""
    profile = BENCHMARK_PROFILES[benchmark]
    use_iceberg = config.get("tpcbench_args", {}).get("use_iceberg", False)

    required = []
    if not use_iceberg:
        required.append(profile["data_env"])

    missing = [v for v in required if not os.environ.get(v)]
    if missing:
        print(
            f"Error: Required environment variable(s) not set for "
            f"{benchmark}: {', '.join(missing)}",
            file=sys.stderr,
        )
        sys.exit(1)

    # Default ICEBERG_DATABASE to the benchmark name if not already set
    if use_iceberg and not os.environ.get("ICEBERG_DATABASE"):
        os.environ["ICEBERG_DATABASE"] = benchmark


def build_spark_submit_cmd(config, benchmark, args):
    """Build the spark-submit command list."""
    spark_home = os.environ["SPARK_HOME"]
    spark_master = os.environ["SPARK_MASTER"]
    profile = BENCHMARK_PROFILES[benchmark]

    cmd = [os.path.join(spark_home, "bin", "spark-submit")]
    cmd += ["--master", spark_master]

    # --jars
    jars = config.get("spark_submit", {}).get("jars", [])
    if jars:
        cmd += ["--jars", ",".join(resolve_env_in_list(jars))]

    # --driver-class-path
    driver_cp = config.get("spark_submit", {}).get("driver_class_path", [])
    if driver_cp:
        cmd += ["--driver-class-path", ":".join(resolve_env_in_list(driver_cp))]

    # Merge spark confs: common + benchmark profile + engine overrides
    conf = dict(COMMON_SPARK_CONF)
    conf["spark.executor.instances"] = profile["executor_instances"]
    conf["spark.executor.cores"] = profile["executor_cores"]
    conf["spark.cores.max"] = profile["max_cores"]

    engine_conf = config.get("spark_conf", {})
    for key, val in engine_conf.items():
        if isinstance(val, bool):
            val = "true" if val else "false"
        conf[resolve_env(key)] = resolve_env(str(val))

    # JFR profiling: append to extraJavaOptions (preserving any existing values)
    if args.jfr:
        jfr_dir = args.jfr_dir
        driver_jfr = (
            f"-XX:StartFlightRecording=disk=true,dumponexit=true,"
            f"filename={jfr_dir}/driver.jfr,settings=profile"
        )
        executor_jfr = (
            f"-XX:StartFlightRecording=disk=true,dumponexit=true,"
            f"filename={jfr_dir}/executor.jfr,settings=profile"
        )
        for spark_key, jfr_opts in [
            ("spark.driver.extraJavaOptions", driver_jfr),
            ("spark.executor.extraJavaOptions", executor_jfr),
        ]:
            existing = conf.get(spark_key, "")
            conf[spark_key] = f"{existing} {jfr_opts}".strip()

    # async-profiler: attach as a Java agent via -agentpath
    if args.async_profiler:
        ap_home = os.environ.get("ASYNC_PROFILER_HOME", "")
        if not ap_home:
            print(
                "Error: ASYNC_PROFILER_HOME is not set. "
                "Set it to the async-profiler installation directory.",
                file=sys.stderr,
            )
            sys.exit(1)
        lib_ext = "dylib" if sys.platform == "darwin" else "so"
        ap_lib = os.path.join(ap_home, "lib", f"libasyncProfiler.{lib_ext}")
        ap_dir = args.async_profiler_dir
        ap_event = args.async_profiler_event
        ap_fmt = args.async_profiler_format
        ext = {"flamegraph": "html", "jfr": "jfr", "collapsed": "txt", "text": "txt"}[ap_fmt]

        driver_ap = (
            f"-agentpath:{ap_lib}=start,event={ap_event},"
            f"{ap_fmt},file={ap_dir}/driver.{ext}"
        )
        executor_ap = (
            f"-agentpath:{ap_lib}=start,event={ap_event},"
            f"{ap_fmt},file={ap_dir}/executor.{ext}"
        )
        for spark_key, ap_opts in [
            ("spark.driver.extraJavaOptions", driver_ap),
            ("spark.executor.extraJavaOptions", executor_ap),
        ]:
            existing = conf.get(spark_key, "")
            conf[spark_key] = f"{existing} {ap_opts}".strip()

    for key, val in sorted(conf.items()):
        cmd += ["--conf", f"{key}={val}"]

    # tpcbench.py path
    cmd.append("tpcbench.py")

    # tpcbench args
    engine_name = config.get("engine", {}).get("name", args.engine)
    cmd += ["--name", engine_name]
    cmd += ["--benchmark", benchmark]

    use_iceberg = config.get("tpcbench_args", {}).get("use_iceberg", False)
    if use_iceberg:
        cmd += ["--catalog", resolve_env("${ICEBERG_CATALOG}")]
        cmd += ["--database", resolve_env("${ICEBERG_DATABASE}")]
    else:
        data_var = profile["data_env"]
        data_val = os.environ.get(data_var, "")
        cmd += ["--data", data_val]

    cmd += ["--output", args.output]
    cmd += ["--iterations", str(args.iterations)]

    if args.query is not None:
        cmd += ["--query", str(args.query)]

    if profile["format"] and not use_iceberg:
        cmd += ["--format", profile["format"]]

    return cmd


def restart_spark():
    """Stop and start Spark master and worker."""
    spark_home = os.environ["SPARK_HOME"]
    sbin = os.path.join(spark_home, "sbin")
    spark_master = os.environ["SPARK_MASTER"]

    # Stop (ignore errors)
    subprocess.run(
        [os.path.join(sbin, "stop-master.sh")],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    subprocess.run(
        [os.path.join(sbin, "stop-worker.sh")],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )

    # Start (check errors)
    r = subprocess.run([os.path.join(sbin, "start-master.sh")])
    if r.returncode != 0:
        print("Error: Failed to start Spark master", file=sys.stderr)
        sys.exit(1)

    r = subprocess.run([os.path.join(sbin, "start-worker.sh"), spark_master])
    if r.returncode != 0:
        print("Error: Failed to start Spark worker", file=sys.stderr)
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description="Consolidated TPC benchmark runner for Spark-based engines."
    )
    parser.add_argument(
        "--engine",
        required=True,
        help="Engine name (matches a TOML file in engines/)",
    )
    parser.add_argument(
        "--benchmark",
        required=True,
        choices=["tpch", "tpcds"],
        help="Benchmark to run",
    )
    parser.add_argument(
        "--iterations", type=int, default=1, help="Number of iterations (default: 1)"
    )
    parser.add_argument(
        "--output", default=".", help="Output directory (default: .)"
    )
    parser.add_argument(
        "--query", type=int, default=None, help="Run a single query number"
    )
    parser.add_argument(
        "--no-restart",
        action="store_true",
        help="Skip Spark master/worker restart",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the spark-submit command without executing",
    )
    parser.add_argument(
        "--jfr",
        action="store_true",
        help="Enable Java Flight Recorder profiling for driver and executors",
    )
    parser.add_argument(
        "--jfr-dir",
        default="/results/jfr",
        help="Directory for JFR output files (default: /results/jfr)",
    )
    parser.add_argument(
        "--async-profiler",
        action="store_true",
        help="Enable async-profiler for driver and executors (profiles Java + native code)",
    )
    parser.add_argument(
        "--async-profiler-dir",
        default="/results/async-profiler",
        help="Directory for async-profiler output files (default: /results/async-profiler)",
    )
    parser.add_argument(
        "--async-profiler-event",
        default="cpu",
        help="async-profiler event type: cpu, wall, alloc, lock, etc. (default: cpu)",
    )
    parser.add_argument(
        "--async-profiler-format",
        default="flamegraph",
        choices=["flamegraph", "jfr", "collapsed", "text"],
        help="async-profiler output format (default: flamegraph)",
    )
    args = parser.parse_args()

    config = load_engine_config(args.engine)

    # Apply env defaults and exports before validation
    apply_env_defaults(config)
    apply_env_exports(config)

    check_common_env()
    check_required_env(config)
    check_benchmark_env(config, args.benchmark)

    # Restart Spark unless --no-restart or --dry-run
    if not args.no_restart and not args.dry_run:
        restart_spark()

    # Create profiling output directories (skip for dry-run)
    if not args.dry_run:
        if args.jfr:
            os.makedirs(args.jfr_dir, exist_ok=True)
        if args.async_profiler:
            os.makedirs(args.async_profiler_dir, exist_ok=True)

    cmd = build_spark_submit_cmd(config, args.benchmark, args)

    if args.dry_run:
        # Group paired arguments (e.g. --conf key=value) on one line
        parts = []
        i = 0
        while i < len(cmd):
            token = cmd[i]
            if token.startswith("--") and i + 1 < len(cmd) and not cmd[i + 1].startswith("--"):
                parts.append(f"{token} {cmd[i + 1]}")
                i += 2
            else:
                parts.append(token)
                i += 1
        print(" \\\n    ".join(parts))
    else:
        r = subprocess.run(cmd)
        sys.exit(r.returncode)


if __name__ == "__main__":
    main()
