#!/usr/bin/env python3
"""Classify evidence patterns in a bounded NCCL/PyTorch log read from stdin.

Does not run CUDA, Docker, PyTorch, host commands, or network requests.
It outputs only categories and signal counts, never the original log lines.
Categories are hints to investigate; they are not definitive root causes.
"""
from __future__ import annotations
import argparse
import json
import re
import sys

MAX_BYTES = 2_000_000
PATTERNS = (
    ("SHM_ALLOCATION", re.compile(r"failed to (?:extend|create|open).{0,80}/dev/shm|no space left on device|shm.{0,45}fail", re.I)),
    ("NETWORK_CONNECTIVITY", re.compile(r"socket.{0,70}(?:failed|error|refused)|connection refused|network is unreachable|connect\(\).{0,60}(?:fail|error)|no route to host", re.I)),
    ("GPU_PEER_ACCESS", re.compile(r"p2p.{0,60}(?:failed|unsupported|unavailable)|peer.access.{0,50}(?:failed|unsupported)|cudaErrorPeerAccessUnsupported", re.I)),
    ("WATCHDOG_TIMEOUT", re.compile(r"watchdog.{0,80}timeout|collective.{0,50}timed out|operation timeout|timeout during collective", re.I)),
    ("CUDA_OR_DRIVER_ERROR", re.compile(r"cuda error|cudaErrorUnknown|driver/library version mismatch|NVRM: Xid|GPU is lost", re.I)),
)
GENERIC = re.compile(r"nccl.{0,100}unhandled system error", re.I)


def summarize(log: str) -> dict[str, object]:
    if len(log.encode("utf-8")) > MAX_BYTES:
        raise ValueError("Input exceeds 2 MB")
    counts = {name: len(pattern.findall(log)) for name, pattern in PATTERNS}
    evidence = [name for name, count in counts.items() if count]
    if not evidence:
        verdict = "GENERIC_ERROR_NO_DIAGNOSTIC_SIGNAL" if GENERIC.search(log) else "NO_KNOWN_SIGNAL"
    elif len(evidence) == 1:
        verdict = f"INVESTIGATE_{evidence[0]}"
    else:
        verdict = "MULTIPLE_SIGNALS_REQUIRES_TIMELINE"
    return {
        "verdict": verdict,
        "signal_counts": counts,
        "note": "Matched text only: correlate timestamps, ranks, platform, device topology and exact versions before changing settings.",
        "recommendations": {
            "SHM_ALLOCATION": "Check /dev/shm usage and container IPC limits without assuming host IPC access is needed.",
            "NETWORK_CONNECTIVITY": "Inspect reachable interfaces, firewall and which address NCCL actually selected.",
            "GPU_PEER_ACCESS": "Inspect GPU topology; test P2P transport changes only in an isolated, approved reproduction.",
            "WATCHDOG_TIMEOUT": "Do not hide hangs by increasing timeout; check which collective/rank stopped progressing.",
            "CUDA_OR_DRIVER_ERROR": "Inspect CUDA driver and GPU diagnostics before treating this as an NCCL environment flag issue.",
        },
    }


def self_test() -> dict[str, object]:
    cases = [
        ("NCCL WARN failed to extend /dev/shm/nccl-test to 4194660 bytes", "INVESTIGATE_SHM_ALLOCATION"),
        ("NCCL WARN socket connect() failed: Connection refused", "INVESTIGATE_NETWORK_CONNECTIVITY"),
        ("NCCL WARN P2P failed between GPU devices", "INVESTIGATE_GPU_PEER_ACCESS"),
        ("Watchdog caught collective operation timeout", "INVESTIGATE_WATCHDOG_TIMEOUT"),
        ("NVRM: Xid 79", "INVESTIGATE_CUDA_OR_DRIVER_ERROR"),
        ("NCCL unhandled system error", "GENERIC_ERROR_NO_DIAGNOSTIC_SIGNAL"),
        ("startup successful", "NO_KNOWN_SIGNAL"),
        ("NCCL WARN P2P failed; socket connection refused", "MULTIPLE_SIGNALS_REQUIRES_TIMELINE"),
    ]
    for value, expected in cases:
        out = summarize(value)
        assert out["verdict"] == expected, (value, expected, out)
        assert value not in json.dumps(out), "Tool leaked the log line"
    try:
        summarize("X" * (MAX_BYTES + 1))
        raise AssertionError("Oversize input was accepted")
    except ValueError:
        pass
    return {"status": "PASS", "synthetic_cases": len(cases), "oversize_rejection_checks": 1,
            "gpu_workloads": 0, "network_requests": 0}


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--self-test", action="store_true")
    args = p.parse_args()
    if args.self_test:
        print(json.dumps(self_test(), indent=2))
        return 0
    raw = sys.stdin.buffer.read(MAX_BYTES + 1)
    try:
        if len(raw) > MAX_BYTES:
            raise ValueError("too many bytes")
        value = raw.decode("utf-8", errors="replace")
        print(json.dumps(summarize(value), indent=2))
        return 0
    except ValueError as e:
        print(json.dumps({"status": "INVALID_INPUT", "error_type": type(e).__name__}), file=sys.stderr)
        return 2


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