#!/usr/bin/env python3
"""The probe's subject: what an ordinary ONNX Runtime user does, and nothing more.

Modes:
  default  import onnxruntime, build a one-node Add model in memory, run it once, sleep, exit 0.
  api      the same, but call onnxruntime.disable_telemetry_events() right after the import.
  posctl   instrument positive control: never imports onnxruntime. Python itself looks up the
           collector name and opens TLS to it, to prove the listeners would see such an attempt.

It writes its own timestamps (wall clock, seconds since the epoch) to --marks as JSON, once after the
work and again at exit, so a crash still leaves the first half.
"""
import argparse
import json
import socket
import ssl
import sys
import time

COLLECTOR = "mobile.events.data.microsoft.com"  # the host the recon found in the 1.29/1.30 binaries


def save(path, marks):
    with open(path, "w") as f:
        json.dump(marks, f, indent=1)


def ort_workload(mode, marks):
    import onnxruntime as ort  # the import alone initialises ORT's environment (and its telemetry)
    marks["t_import_done"] = time.time()
    marks["ort_version"] = ort.__version__
    if mode == "api":
        ort.disable_telemetry_events()  # the documented runtime switch; reachable only after import
        marks["t_api_disabled"] = time.time()

    import numpy as np
    from onnx import TensorProto, helper

    # Z = X + Y. Canary strings let us find this model's details in the local queue afterwards.
    vec = [2]
    graph = helper.make_graph(
        [helper.make_node("Add", ["X", "Y"], ["Z"])],
        "probe_graph_canary",
        [helper.make_tensor_value_info("X", TensorProto.FLOAT, vec),
         helper.make_tensor_value_info("Y", TensorProto.FLOAT, vec)],
        [helper.make_tensor_value_info("Z", TensorProto.FLOAT, vec)],
    )
    model = helper.make_model(graph, producer_name="probe-kit",
                              opset_imports=[helper.make_opsetid("", 17)], ir_version=9)
    helper.set_model_props(model, {"probe_meta_key": "probe_meta_value_canary"})

    sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
    marks["t_session_created"] = time.time()
    out = sess.run(None, {"X": np.array([1, 2], np.float32), "Y": np.array([10, 20], np.float32)})
    marks["t_run_done"] = time.time()
    marks["run_output"] = out[0].tolist()


def posctl_workload(marks):
    time.sleep(5)
    marks["t_posctl_connect"] = time.time()
    try:
        ctx = ssl.create_default_context()
        with socket.create_connection((COLLECTOR, 443), timeout=5) as raw:
            with ctx.wrap_socket(raw, server_hostname=COLLECTOR):
                pass
    except Exception as e:  # the listener hangs up mid-handshake, so an error is the expected outcome
        marks["posctl_error"] = repr(e)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=["default", "api", "posctl"], required=True)
    ap.add_argument("--hold", type=float, required=True, help="seconds to stay alive after the work")
    ap.add_argument("--marks", required=True)
    a = ap.parse_args()

    marks = {"t_start": time.time(), "mode": a.mode, "hold_s": a.hold}
    if a.mode == "posctl":
        posctl_workload(marks)
    else:
        ort_workload(a.mode, marks)
    save(a.marks, marks)

    time.sleep(a.hold)
    marks["t_exit"] = time.time()
    save(a.marks, marks)
    return 0


if __name__ == "__main__":
    sys.exit(main())