#!/usr/bin/env python3
"""Run the repository's native XYZAC RTCP case and record LinuxCNC stat ticks.

This is an oracle collector, not a simulator.  It intentionally starts the
same run-5axis-rtcp.sh launcher used by the native case and reads values only
from linuxcnc.stat().  Fields unavailable through the public stat API remain
explicitly null.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import signal
import subprocess
import sys
import time
from pathlib import Path
from typing import Any

# The oracle must bind to the same source build that run-5axis-rtcp.sh starts.
# In particular, importing the distro linuxcnc extension would silently bind to
# a different NML ABI/version and produce plausible but unusable status values.
ROOT = Path(__file__).resolve().parents[1]
SOURCE_PYTHON_DIR = ROOT / "linuxcnc/lib/python"
sys.path.insert(0, str(SOURCE_PYTHON_DIR))
import linuxcnc

from run_axis3_tutorial_native import cleanup_linuxcnc_services


CASE = "linuxcnc/configs/sim/axis/vismach/5axis/table-rotary-tilting/xyzac-trt.ini"
PROGRAM = ROOT / "linuxcnc/configs/sim/axis/vismach/5axis/table-rotary-tilting/demos/xyzac_switchkins.ngc"
SCHEMA = "xyzac-native-linuxcnc-tick-trace.v1"


def canonical(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def sha256_bytes(value: bytes) -> str:
    return hashlib.sha256(value).hexdigest()


def finite_list(value: Any) -> list[float] | None:
    try:
        return [float(item) for item in value]
    except (TypeError, ValueError):
        return None


def scalar(value: Any) -> float | int | bool | str | None:
    if value is None or isinstance(value, (bool, int, str)):
        return value
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return number if number == number and abs(number) != float("inf") else None


def stat_row(stat: Any, sample_index: int, previous_taskbeat: int | None) -> dict[str, Any]:
    stat.poll()
    taskbeat = scalar(getattr(stat, "task_heartbeat", None))
    commanded_pose = finite_list(getattr(stat, "position", None))
    distance_to_go = finite_list(getattr(stat, "dtg", None))
    return {
        "sample_index": sample_index,
        "taskbeat": taskbeat,
        "task_heartbeat": taskbeat,
        "taskbeat_changed": previous_taskbeat is None or taskbeat != previous_taskbeat,
        "task_state": scalar(getattr(stat, "task_state", None)),
        "task_mode": scalar(getattr(stat, "task_mode", None)),
        "interp_state": scalar(getattr(stat, "interp_state", None)),
        "exec_state": scalar(getattr(stat, "exec_state", None)),
        "current_line": scalar(getattr(stat, "current_line", None)),
        "motion_line": scalar(getattr(stat, "motion_line", None)),
        "read_line": scalar(getattr(stat, "read_line", None)),
        "actual_pose": finite_list(getattr(stat, "actual_position", None)),
        "commanded_pose": commanded_pose,
        "joint_actual": finite_list(getattr(stat, "joint_actual_position", None)),
        "joint_commanded": finite_list(getattr(stat, "joint_position", None)),
        "velocity": scalar(getattr(stat, "current_vel", None)),
        "acceleration": scalar(getattr(stat, "acceleration", None)),
        "feed_override": scalar(getattr(stat, "feedrate", None)),
        "rapid_override": scalar(getattr(stat, "rapidrate", None)),
        "feed_override_enabled": scalar(getattr(stat, "feed_override_enabled", None)),
        "queue_depth": scalar(getattr(stat, "queue", None)),
        "active_queue": scalar(getattr(stat, "active_queue", None)),
        "queue_full": scalar(getattr(stat, "queue_full", None)),
        "motion_id": scalar(getattr(stat, "motion_id", None)),
        "active_segment_id": scalar(getattr(stat, "motion_id", None)),
        "active_segment_geometry": None,
        "active_segment_geometry_status": "unavailable_linuxcnc_stat_api",
        "distance_to_go": distance_to_go,
        "position": finite_list(getattr(stat, "position", None)),
        "homed": list(getattr(stat, "homed", [])) if getattr(stat, "homed", None) is not None else None,
        "limits": list(getattr(stat, "limit", [])) if getattr(stat, "limit", None) is not None else None,
        "kinematics_type": scalar(getattr(stat, "kinematics_type", None)),
        "motion_type": scalar(getattr(stat, "motion_type", None)),
        "feed_hold_enabled": scalar(getattr(stat, "feed_hold_enabled", None)),
        "paused": scalar(getattr(stat, "paused", None)),
    }


def wait_stat(stat: Any, predicate: Any, timeout: float) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        stat.poll()
        if predicate(stat):
            return
        time.sleep(0.02)
    raise RuntimeError("native_stat_transition_timeout")


def terminate_process(process: subprocess.Popen[str]) -> None:
    if process.poll() is not None:
        return
    try:
        os.killpg(process.pid, signal.SIGTERM)
        process.wait(timeout=5)
    except (ProcessLookupError, subprocess.TimeoutExpired):
        try:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
        try:
            process.wait(timeout=2)
        except subprocess.TimeoutExpired:
            pass


def native_axis_window(environment: dict[str, str]) -> dict[str, Any] | None:
    result = subprocess.run(["xwininfo", "-root", "-tree"], env=environment, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, check=False)
    for line in result.stdout.splitlines():
        if '("axis" "Axis")' not in line or "xyzac-trt" not in line.lower():
            continue
        marker = line.strip().split('"', 2)
        if len(marker) < 2:
            continue
        size = re.search(r"\)\s+(\d+)x(\d+)\+", line)
        if not size or int(size.group(1)) < 600 or int(size.group(2)) < 400:
            continue
        return {"title": marker[1], "width": int(size.group(1)), "height": int(size.group(2)), "raw": line.strip()}
    return None


def run(duration: float) -> dict[str, Any]:
    if not PROGRAM.is_file():
        return {"schema": SCHEMA, "status": "blocked", "code": "native_program_missing", "program": PROGRAM.as_posix()}
    cleanup = cleanup_linuxcnc_services()
    environment = os.environ.copy()
    environment["DISPLAY"] = environment.get("DISPLAY", ":1")
    command = ["./run-5axis-rtcp.sh", "xyzac-trt.ini"]
    process = subprocess.Popen(command, cwd=ROOT, env=environment, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, start_new_session=True)
    output_lines: list[str] = []
    rows: list[dict[str, Any]] = []
    errors: list[str] = []
    stage = "launch"
    stat: Any = None
    control: Any = None
    connected = False
    window: dict[str, Any] | None = None
    try:
        deadline = time.monotonic() + 20
        while time.monotonic() < deadline and process.poll() is None:
            window = native_axis_window(environment)
            if not window:
                time.sleep(0.1)
                continue
            # The AXIS toplevel is created only after linuxcncsvr/milltask have
            # opened the NML buffers. Creating linuxcnc.stat before that point
            # permanently binds an invalid channel, so wait for the real
            # window first and create fresh channels afterwards.
            time.sleep(0.1)
        if stat is None or control is None or not connected:
            if not window:
                raise RuntimeError("native_linuxcnc_window_unavailable")
            stat = None
            control = None
            attach_deadline = time.monotonic() + 10
            while time.monotonic() < attach_deadline and process.poll() is None:
                try:
                    candidate_stat = linuxcnc.stat()
                    candidate_stat.poll()
                    stat = candidate_stat
                    control = linuxcnc.command()
                    connected = True
                    break
                except BaseException:
                    stat = None
                    control = None
                    time.sleep(0.1)
            if stat is None or control is None or not connected:
                raise RuntimeError("native_linuxcnc_stat_unavailable")

        stage = "state_estop_reset"
        control.state(linuxcnc.STATE_ESTOP_RESET)
        control.wait_complete(5)
        wait_stat(stat, lambda item: item.task_state != linuxcnc.STATE_ESTOP, 5)
        stage = "state_on"
        control.state(linuxcnc.STATE_ON)
        control.wait_complete(5)
        wait_stat(stat, lambda item: item.task_state == linuxcnc.STATE_ON, 5)
        stage = "mode_manual"
        control.mode(linuxcnc.MODE_MANUAL)
        control.wait_complete(5)
        for joint in range(5):
            stage = f"home_joint_{joint}"
            control.home(joint)
            control.wait_complete(5)
        stage = "wait_home_all"
        wait_stat(stat, lambda item: all(bool(value) for value in list(getattr(item, "homed", []))[:5]), 10)
        stage = "program_open"
        control.program_open(str(PROGRAM))
        control.wait_complete(5)
        stage = "mode_auto"
        control.mode(linuxcnc.MODE_AUTO)
        control.wait_complete(5)
        stage = "program_run"
        control.auto(linuxcnc.AUTO_RUN, 0)

        start = time.monotonic()
        sample_index = 0
        previous_taskbeat: int | None = None
        last_change = start
        while time.monotonic() - start < duration:
            row = stat_row(stat, sample_index, previous_taskbeat)
            current_taskbeat = row["taskbeat"]
            if row["taskbeat_changed"] or not rows:
                row["native_elapsed_ms"] = round((time.monotonic() - start) * 1000.0, 3)
                rows.append(row)
                previous_taskbeat = current_taskbeat if isinstance(current_taskbeat, int) else previous_taskbeat
                last_change = time.monotonic()
            if row["task_state"] == linuxcnc.STATE_OFF or row["interp_state"] == linuxcnc.INTERP_IDLE and row["queue_depth"] == 0 and len(rows) > 4:
                break
            sample_index += 1
            time.sleep(0.005)
        final = stat_row(stat, sample_index + 1, previous_taskbeat)
        if final["taskbeat_changed"]:
            final["native_elapsed_ms"] = round((time.monotonic() - start) * 1000.0, 3)
            rows.append(final)
    except Exception as error:
        errors.append(f"{stage}:{type(error).__name__}:{error}")
    finally:
        terminate_process(process)
        try:
            output, _ = process.communicate(timeout=2)
            output_lines.extend(output.splitlines())
        except (subprocess.TimeoutExpired, ValueError):
            pass
        cleanup_linuxcnc_services()

    source_bytes = PROGRAM.read_bytes()
    ini_path = ROOT / CASE
    launcher_path = ROOT / "run-5axis-rtcp.sh"
    heartbeat_values = [row.get("task_heartbeat") for row in rows]
    valid_heartbeats = [value for value in heartbeat_values if isinstance(value, int)]
    heartbeat_monotonic = all(left < right for left, right in zip(valid_heartbeats, valid_heartbeats[1:]))
    trace_hash = sha256_bytes(canonical(rows).encode())
    python_module = getattr(linuxcnc, "__file__", None)
    source_python_module = str(ROOT / "linuxcnc/lib/python/linuxcnc.so")
    if python_module != source_python_module:
        errors.append("NATIVE_PYTHON_MODULE_MISMATCH")
    status = "verified" if rows and not errors and heartbeat_monotonic and len(valid_heartbeats) == len(set(valid_heartbeats)) else "blocked"
    if rows and not heartbeat_monotonic:
        errors.append("NATIVE_TASK_HEARTBEAT_NON_MONOTONIC")
    return {
        "schema": SCHEMA,
        "status": status,
        "case": {"path": CASE, "scope": "xyzac_only"},
        "launcher": command,
        "native_axis_window": window,
        "program": {"path": PROGRAM.relative_to(ROOT).as_posix(), "sha256": sha256_bytes(source_bytes)},
        "native_linuxcnc": {"python_module": python_module, "expected_source_python_module": source_python_module, "display": environment["DISPLAY"], "python_module_sha256": sha256_bytes(Path(python_module).read_bytes()) if python_module and Path(python_module).is_file() else None},
        "inputs": {
            "case_ini": {"path": CASE, "sha256": sha256_bytes(ini_path.read_bytes())},
            "launcher": {"path": "run-5axis-rtcp.sh", "sha256": sha256_bytes(launcher_path.read_bytes())},
            "program": {"path": PROGRAM.relative_to(ROOT).as_posix(), "sha256": sha256_bytes(source_bytes)},
            "source_python_module": {"path": source_python_module, "sha256": sha256_bytes(Path(source_python_module).read_bytes()) if Path(source_python_module).is_file() else None},
        },
        "tick_source": "linuxcnc.stat.task_heartbeat",
        "tick_integrity": {"first": valid_heartbeats[0] if valid_heartbeats else None, "last": valid_heartbeats[-1] if valid_heartbeats else None, "row_count": len(valid_heartbeats), "unique_count": len(set(valid_heartbeats)), "strictly_monotonic": heartbeat_monotonic},
        "rows": rows,
        "row_count": len(rows),
        "trace_sha256": trace_hash,
        "errors": errors,
        "launcher_output_tail": "\n".join(output_lines[-80:]),
        "launcher_output_sha256": sha256_bytes("\n".join(output_lines).encode()),
        "cleanup": cleanup,
        "known_boundaries": [
            "native stat exposes motion_id but not full TP segment geometry; active_segment_geometry is explicit null",
            "taskbeat is the public native task-cycle counter; wall-clock is recorded only as diagnostic elapsed_ms",
            "source run-in-place launcher and installed Python extension must use the same LinuxCNC build; mismatch is blocked",
            "this report is an oracle input and does not claim native/WASM parity by itself",
        ],
        "release_status": "STOP_RELEASE",
        "report_sha256": None,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--duration", type=float, default=8.0)
    parser.add_argument("--output", type=Path, default=Path("coverage/xyzac-native-tick-trace.report.json"))
    args = parser.parse_args()
    report = run(max(0.1, min(args.duration, 300.0)))
    unsigned = {key: value for key, value in report.items() if key != "report_sha256"}
    report["report_sha256"] = sha256_bytes(canonical(unsigned).encode())
    destination = args.output if args.output.is_absolute() else ROOT / args.output
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(json.dumps(report, ensure_ascii=False, sort_keys=True, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({"schema": SCHEMA, "status": report["status"], "row_count": report["row_count"], "trace_sha256": report["trace_sha256"], "errors": report["errors"], "report_sha256": report["report_sha256"]}, ensure_ascii=False))
    return 0 if report["status"] == "verified" else 1


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