"""Bin a pytest trace into a coarse failure category.

The goal is a single tag per failure that aggregates well in a report. We do
this with cheap regex on the raw trace string. Categories, checked in order:

  OOM             — torch OutOfMemoryError, "CUDA out of memory"
  load_error      — failure raised before forward (from_pretrained, safetensors,
                    HFValidationError, missing weight, gated/auth)
  cuda_runtime    — runtime CUDA/cuBLAS/cuDNN/nvrtc/Triton compile errors
  output_mismatch — text / tensor comparison failures (assertEqual, Expectations,
                    Tensor-likes are not close, etc.)
  import_or_config — Python-level errors before the test body runs
                    (ImportError, AttributeError on config, TypeError on
                    __init__ signatures).
  other           — fallback
"""

from __future__ import annotations

import re


_OOM_PAT = re.compile(r"OutOfMemoryError|CUDA out of memory|MallocFailure|HIP out of memory", re.I)
_LOAD_PAT = re.compile(
    r"from_pretrained|safetensors\.|HFValidationError|Repository Not Found|gated|"
    r"Cannot read|UnboundLocalError.*loading|FileNotFoundError|access requested|"
    r"401 Client Error|403 Client Error",
    re.I,
)
_CUDA_RUNTIME_PAT = re.compile(
    r"CUDA error|CUBLAS_STATUS|CUDNN_STATUS|cudnn[_ ]frontend|nvrtc|"
    r"triton\.compiler|RuntimeError: Triton|c10::Error|NCCL.*error",
    re.I,
)
_OUTPUT_MISMATCH_PAT = re.compile(
    r"Tensor-likes are not close|"
    r"assertEqual|assertSequenceEqual|self\.assertListEqual|"
    r"assertAlmostEqual|assertGreater|expected_text|"
    r"AssertionError",  # generic fallback — most assertion failures are output mismatches
    re.I | re.DOTALL,
)
_IMPORT_CFG_PAT = re.compile(
    r"^.*ImportError|ModuleNotFoundError|"
    r"AttributeError:.*(config|object has no attribute)|"
    r"TypeError:.*(__init__|got an unexpected keyword argument)|"
    r"ValueError:.*Unrecognized configuration",
    re.I | re.M,
)


def classify(trace: str) -> str:
    if not trace:
        return "other"
    for tag, pat in (
        ("OOM", _OOM_PAT),
        ("load_error", _LOAD_PAT),
        ("cuda_runtime", _CUDA_RUNTIME_PAT),
        ("import_or_config", _IMPORT_CFG_PAT),
        ("output_mismatch", _OUTPUT_MISMATCH_PAT),
    ):
        if pat.search(trace):
            return tag
    return "other"


def short_excerpt(trace: str, max_chars: int = 240) -> str:
    """Take the LAST non-empty line of the trace (the actual exception line).
    Then trim to `max_chars`.
    """
    if not trace:
        return ""
    for line in reversed(trace.splitlines()):
        line = line.strip()
        if line:
            return (line[: max_chars - 1] + "…") if len(line) > max_chars else line
    return ""


if __name__ == "__main__":
    samples = [
        "tests/...py:42: torch.OutOfMemoryError: CUDA out of memory.",
        "RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED",
        "AssertionError: 'Paris' != 'capital of France'",
        "ImportError: cannot import name 'foo'",
        "HFValidationError: Repository Not Found for url: ...",
        "Tensor-likes are not close (...) max_abs_diff=0.5",
    ]
    for s in samples:
        print(f"{classify(s):<18}  {s}")