# Copyright 2026 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""XML to JSON Converter for JAX GHA Test Results.

This utility parses JUnit XML test reports generated by either Bazel or Pytest,
converts them into a flat, newline-delimited JSON format, and streams them to
stdout. The generated JSON records are fully compatible with the schema of the
BigQuery table `jax-dev:jax_ci.tests`.

Key features include:
1. Automatic detection of test execution engine (Bazel vs. Pytest).
2. Heuristic-based inference of source file paths from XML names/classnames.
3. Mapping of execution-engine specific statuses to a unified format (FAILED,
ERROR, SKIPPED, PASSED).
4. Fallback timestamp resolution using GHA jobs.json to satisfy BQ required
field constraints.
"""

import datetime
import json
import os
import sys
import xml.etree.ElementTree as ET


def extract_file_name(xml_file_name):
  """Heuristically recovers the Python source file name from a Bazel XML name.

  Bazel XML report filenames are typically generated using double underscores to
  encode the directory structure and target names, e.g.,
  'tests__scaled_dot_test_gpu__shard_2_of_5__test.xml'.
  This function splits the name, strips standard suffixes, and reconstructs the
  original file path (e.g., 'tests/scaled_dot_test_gpu.py').
  """
  name = xml_file_name.replace(".xml", "")
  parts = name.split("__")

  target_parts = []
  for p in parts:
    # Skip standard placeholder tokens.
    if p == "test":
      continue
    # Stop at shard parameters to avoid duplicate target naming.
    if p.startswith("shard_"):
      break
    target_parts.append(p)

  target_path = "/".join(target_parts)

  # Truncate the target path immediately after '_test' to clean up inferred paths.
  idx = target_path.rfind("_test")
  if idx != -1:
    target_path = target_path[: idx + len("_test")]

  return target_path + ".py"


def process_xml(file_path, job_id, fallback_timestamp=None):
  """Parses a JUnit XML file and streams flat JSON records to stdout.

  Overall Logic:
  1. Parse the XML tree and detect whether it is a Bazel or Pytest report.
  2. Extract the test suite and loop over each '<testcase>' element.
  3. Map the status, infer the source file path, and capture failure/error
  tracebacks.
  4. Resolve missing timestamps using GHA jobs.json or the current execution
  time.
  5. Serialize the finalized flat record to stdout as a single-line JSON string.
  """
  base_name = os.path.basename(file_path)

  try:
    tree = ET.parse(file_path)
    root = tree.getroot()
  except Exception as e:
    sys.stderr.write(f"Error parsing {file_path}: {e}\n")
    return

  # Detect the execution engine format.
  # Pytest XML reports usually contain 'pytest' in the filename or are explicitly named 'pytest tests'.
  # Bazel reports encode targets using double underscores '__' in their filenames.
  test_type = "pytest"
  if "pytest" in base_name or root.get("name") == "pytest tests":
    test_type = "pytest"
  elif "__" in base_name and file_path.endswith(".xml"):
    test_type = "bazel"

  # Handle XML structures with either a '<testsuites>' root (multiple suites)
  # or a single '<testsuite>' root.
  if root.tag == "testsuites":
    suites = root.findall("testsuite")
  else:
    suites = [root]

  # For Bazel reports, heuristically infer the source file path from the filename.
  file_path_inferred = (
      extract_file_name(base_name) if test_type == "bazel" else None
  )

  for suite in suites:
    suite_timestamp = suite.get("timestamp")

    for tc in suite.findall("testcase"):
      # BigQuery marks 'timestamp' as REQUIRED. We fallback to:
      #   1. Testcase-specific timestamp (if present).
      #   2. Suite-level timestamp (if present).
      #   3. GHA jobs.json workflow run start timestamp (passed via process_test_results.sh).
      #   4. Current UTC execution time (as a final recovery default).
      timestamp = (
          tc.get("timestamp")
          or suite_timestamp
          or fallback_timestamp
          or datetime.datetime.now(datetime.UTC).isoformat()
      )

      record = {
          "name": tc.get("name"),
          "classname": tc.get("classname"),
          "time": tc.get("time"),
          "job_id": job_id,
          "file_name": base_name,
          "timestamp": timestamp,
      }

      # Extract failure/error details and detailed tracebacks.
      failure = tc.find("failure")
      error = tc.find("error")
      skipped = tc.find("skipped")

      message = None
      detail = None

      if failure is not None:
        message = failure.get("message")
        detail = failure.text.strip() if failure.text else ""
      elif error is not None:
        message = error.get("message")
        detail = error.text.strip() if error.text else ""
      elif skipped is not None:
        message = skipped.get("message")

      if test_type == "bazel":
        record["file_path"] = file_path_inferred
        parts = base_name.split("__")
        if len(parts) > 1:
          record["target"] = parts[1]

        # Map Bazel-specific result/status fields to unified BigQuery statuses.
        result = tc.get("result")
        status = tc.get("status")
        if failure is not None:
          record["status"] = "FAILED"
        elif status == "notrun":
          record["status"] = "SKIPPED"
        elif (
            result == "completed" or result is None
        ):  # Default to passed if completed or not specified.
          record["status"] = "PASSED"
        else:
          record["status"] = "PASSED"  # Final safety fallback.

      else:  # Pytest format.
        cname = tc.get("classname")
        if cname and "." in cname:
          # Infer the source file path from the Python classname module structure.
          idx = cname.rfind(".")
          record["file_path"] = cname[:idx] + ".py"

        # Map standard JUnit XML elements directly to BigQuery status.
        if failure is not None:
          record["status"] = "FAILED"
        elif error is not None:
          record["status"] = "ERROR"
        elif skipped is not None:
          record["status"] = "SKIPPED"
        else:
          record["status"] = "PASSED"

      if message:
        record["message"] = message
      if detail:
        record["detail"] = detail

      # Extract system-out and system-err console logs for debugging records.
      stdout = tc.find("system-out")
      stderr = tc.find("system-err")
      if stdout is not None:
        record["system_out"] = stdout.text.strip() if stdout.text else ""
      if stderr is not None:
        record["system_err"] = stderr.text.strip() if stderr.text else ""

      print(json.dumps(record))


if __name__ == "__main__":
  if len(sys.argv) < 3:
    print("Usage: xml2json.py <xml_file> <job_id> [fallback_timestamp]")
    sys.exit(1)
  fallback = sys.argv[3] if len(sys.argv) > 3 else None
  process_xml(sys.argv[1], int(sys.argv[2]), fallback)