#!/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.
"""Snapshot real-Host memory at explicit lifecycle boundaries on macOS.

Use the release native test binary from scripts/agent-scale.py. This rendezvous
pauses the load driver, so its timing results are not performance measurements.
--stacks adds allocator tracing in a separate run; never compare its RSS with
an uninstrumented run. Raw snapshots and command failures are retained.
"""
import argparse
import hashlib
import json
import os
from pathlib import Path
import platform
import re
import subprocess
import time

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--binary", type=Path, required=True)
parser.add_argument("--agents", type=int, nargs="+", default=[10, 100, 1000])
parser.add_argument("--output", type=Path, default=Path("/tmp/maka-agent-memory"))
parser.add_argument("--stacks", action="store_true")
parser.add_argument("--scenario", choices=["barrier", "flowing"], default="barrier")
parser.add_argument("--waves", type=int, default=5)
parser.add_argument("--mixed", action="store_true")
parser.add_argument("--prompt-bytes", type=int)
parser.add_argument("--delta-bytes", type=int)
parser.add_argument("--deltas", type=int, default=32)
args = parser.parse_args()
if platform.system() != "Darwin" or min(args.agents + [args.deltas, args.waves] + [n for n in (args.prompt_bytes, args.delta_bytes) if n is not None]) < 1:
    parser.error("macOS and positive agent counts are required")
if args.scenario == "barrier" and args.mixed:
    parser.error("--mixed requires --scenario flowing")
if args.scenario == "flowing" and (args.prompt_bytes is not None or args.delta_bytes is not None):
    parser.error("flowing uses fixed light sizes or the deterministic --mixed distribution")
test = "owner::tests::scale::real_host_agent_scale" if args.scenario == "barrier" else "owner::tests::scale::flowing::real_host_agent_memory"
args.output.mkdir(parents=True, exist_ok=True)
report = {
    "machine": platform.platform(),
    "binary": str(args.binary.resolve()),
    "binary_sha256": hashlib.file_digest(args.binary.open("rb"), "sha256").hexdigest(),
    "allocator_tracing": args.stacks,
    "workload": {
        "prompt_bytes": ([1024, 8192, 65536, 262144] if args.mixed else 32) if args.scenario == "flowing" else args.prompt_bytes or 5,
        "delta_bytes": ([16, 256, 2048, 8192] if args.mixed else 6) if args.scenario == "flowing" else args.delta_bytes or 6,
        "deltas": args.deltas, "scenario": args.scenario,
        "waves": args.waves if args.scenario == "flowing" else 1, "mixed": args.mixed,
    },
    "scope": "whole process, including the local HTTP fixture; lifecycle pauses invalidate timing metrics",
    "runs": [],
}


def record():
    (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n")


for count in args.agents:
    folder = args.output / str(count)
    folder.mkdir(exist_ok=True)
    env = {**os.environ, "MAKA_SCALE_AGENTS": str(count), "MAKA_SCALE_MEMORY_PROFILE": "1"}
    for name, value in (("MAKA_SCALE_PROMPT_BYTES", args.prompt_bytes), ("MAKA_SCALE_DELTA_BYTES", args.delta_bytes), ("MAKA_SCALE_DELTAS", args.deltas)):
        env.pop(name, None)
        if value is not None:
            env[name] = str(value)
    for name in ("MAKA_SCALE_WAVES", "MAKA_SCALE_MIXED", "MallocStackLoggingNoCompact"):
        env.pop(name, None)
    if args.scenario == "flowing":
        env["MAKA_SCALE_WAVES"] = str(args.waves)
        if args.mixed:
            env["MAKA_SCALE_MIXED"] = "1"
    env.pop("MallocStackLogging", None)
    if args.stacks:
        env["MallocStackLogging"] = "1"
    row = {"agents": count, "phases": []}
    report["runs"].append(row)
    with (folder / "stdout.log").open("w") as stdout, (folder / "stderr.log").open("w") as stderr:
        process = subprocess.Popen(
            [str(args.binary.resolve()), test, "--ignored", "--exact", "--nocapture", "--test-threads=1"],
            env=env, stdin=subprocess.PIPE, stdout=stdout, stderr=subprocess.PIPE,
            text=True, bufsize=1,
        )
        try:
            for line in process.stderr:
                stderr.write(line)
                stderr.flush()
                if not line.startswith("MAKA_AGENT_MEMORY "):
                    continue
                phase = line.strip().split()[-1]
                started = time.monotonic()
                sample = {"phase": phase, "commands": []}
                rss = subprocess.check_output(["ps", "-o", "rss=", "-p", str(process.pid)], text=True)
                sample["rss_bytes"] = int(rss.strip()) * 1024
                commands = [
                    ("vmmap", ["vmmap", "-summary", str(process.pid)]),
                    ("heap", ["heap", "-s", "--noContent", str(process.pid)]),
                ]
                if args.stacks and phase.endswith(("all_active", "all_admitted", "output_settled", "dropped")):
                    commands.append(("stacks", ["malloc_history", str(process.pid), "-callTree", "-invert", "-ignoreThreads", "-noContent", "-consolidateAllBySymbol"]))
                for name, command in commands:
                    destination = folder / f"{phase}-{name}.txt"
                    with destination.open("w") as output:
                        remaining = 40 - (time.monotonic() - started)
                        if remaining <= 0:
                            sample["commands"].append({"tool": name, "skipped": "checkpoint time budget exhausted"})
                            continue
                        timeout = min(20, remaining)
                        try:
                            result = subprocess.run(command, stdout=output, stderr=subprocess.STDOUT, timeout=timeout)
                            status = {"tool": name, "exit_code": result.returncode}
                        except subprocess.TimeoutExpired:
                            status = {"tool": name, "timeout_seconds": timeout}
                    sample["commands"].append(status)
                heap_text = (folder / f"{phase}-heap.txt").read_text()
                allocated = re.search(r"All zones: (\d+) nodes \((\d+) bytes\)", heap_text)
                if allocated:
                    sample["malloc_live_blocks"] = int(allocated[1])
                    sample["malloc_live_bytes"] = int(allocated[2])
                footprint = re.search(r"Physical footprint: *([^\n]+)", heap_text)
                if footprint:
                    sample["physical_footprint"] = footprint[1]
                sample["checkpoint_seconds"] = time.monotonic() - started
                row["phases"].append(sample)
                record()
                print(json.dumps({"agents": count, **sample}), flush=True)
                process.stdin.write("x")
                process.stdin.flush()
            row["exit_code"] = process.wait(timeout=30)
        finally:
            if process.poll() is None:
                process.terminate()
                process.wait(timeout=10)
            process.stdin.close()
    record()
    if not row["phases"]:
        raise SystemExit(f"no checkpoints; rebuild the test binary: {folder}")
    if row["exit_code"]:
        raise SystemExit(f"profiled load failed (inspect workload and checkpoint diagnostics): {folder}")
    if any(command.get("exit_code") != 0 for phase in row["phases"] for command in phase["commands"]):
        raise SystemExit(f"incomplete profiler evidence: {folder}")
print(f"Evidence: {args.output / 'report.json'}")