#!/usr/bin/env python3
"""Run one simulator pass with a whole-host CPU guard."""
import argparse
import json
import os
import queue
import re
import select
import subprocess
import sys
import threading
import time
from pathlib import Path


def cpu_ticks():
    with open("/proc/stat", encoding="ascii") as source:
        fields = list(map(int, source.readline().split()[1:]))
    idle = fields[3] + (fields[4] if len(fields) > 4 else 0)
    return sum(fields), idle


def read_output(stream, lines):
    for line in stream:
        lines.put(line)
    lines.put(None)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("devices", type=int)
    parser.add_argument("--duration", type=int, default=300, help="run time in seconds (default: 300)")
    parser.add_argument("--cpu-limit", type=float, help="optional stop threshold in aggregate CPU percent")
    args = parser.parse_args()

    core_count = os.cpu_count() or 1
    capacity = core_count * 100.0
    environment = os.environ.copy()
    environment["LT_MQTT_CLIENT_PREFIX"] = f"loadtest-guarded-{args.devices}-{int(time.time())}"
    command = [sys.executable, "simulator.py", "--devices", str(args.devices), "--duration", str(args.duration)]

    previous = cpu_ticks()
    previous_at = time.monotonic()
    process = subprocess.Popen(command, env=environment, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)
    output_lines = queue.Queue()
    threading.Thread(target=read_output, args=(process.stdout, output_lines), daemon=True).start()
    report_path = None
    peak_cpu = 0.0
    sample_count = 0
    next_sample = previous_at + 1.0
    stopped_for_cpu = False

    limit_text = f"stop threshold={args.cpu_limit:.0f}%" if args.cpu_limit is not None else "no CPU stop threshold"
    print(f"Starting {args.devices} devices for {args.duration}s; {limit_text}; host capacity={capacity:.0f}%")
    while process.poll() is None:
        now = time.monotonic()
        if now >= next_sample:
            current = cpu_ticks()
            total_delta = current[0] - previous[0]
            idle_delta = current[1] - previous[1]
            elapsed = now - previous_at
            cpu_percent = (1.0 - idle_delta / total_delta) * capacity if total_delta else 0.0
            peak_cpu = max(peak_cpu, cpu_percent)
            sample_count += 1
            if sample_count == 1 or sample_count % 5 == 0:
                print(f"CPU host sample {sample_count}: {cpu_percent:.1f}% / {capacity:.0f}% capacity ({elapsed:.2f}s)", flush=True)
            previous, previous_at = current, now
            next_sample = now + 1.0
            if args.cpu_limit is not None and cpu_percent > args.cpu_limit:
                stopped_for_cpu = True
                print(f"CPU exceeded {args.cpu_limit:.0f}%; terminating simulator immediately", flush=True)
                process.terminate()
                break

        try:
            line = output_lines.get(timeout=max(0.01, min(0.25, next_sample - time.monotonic())))
        except queue.Empty:
            continue
        if line is None:
            break
        if "Starting load test:" in line or "Connected " in line or "heartbeats=" in line:
            print(line.rstrip(), flush=True)
        match = re.search(r"REPORT: (.+\.json)", line)
        if match:
            report_path = Path(match.group(1))
            print(f"Report: {report_path}", flush=True)

    try:
        return_code = process.wait(timeout=10)
    except subprocess.TimeoutExpired:
        process.kill()
        return_code = process.wait()

    while not output_lines.empty():
        line = output_lines.get_nowait()
        if line and "heartbeats=" in line:
            print(line.rstrip(), flush=True)
        if line:
            match = re.search(r"REPORT: (.+\.json)", line)
            if match:
                report_path = Path(match.group(1))

    print(f"CPU peak: {peak_cpu:.1f}% / {capacity:.0f}% capacity; samples={sample_count}; stopped_for_cpu={stopped_for_cpu}; exit={return_code}")
    if report_path and report_path.is_file():
        report = json.loads(report_path.read_text(encoding="utf-8"))
        summary = {key: report.get(key) for key in ("device_count", "heartbeats_sent", "connect_failures", "publish_errors", "acks_sent", "ack_errors", "acks_unconfirmed", "pushes_received", "duplicate_pushes", "gap_devices_never_pushed", "latency_seconds_by_type")}
        print(json.dumps(summary, indent=2))
    elif not stopped_for_cpu:
        print("No report was produced.")

    return 2 if stopped_for_cpu or return_code != 0 else 0


if __name__ == "__main__":
    raise SystemExit(main())