# Copyright 2023 The OpenXLA 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
#
#     http://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.
# ============================================================================

r"""Produces a `compile_commands.json` from the output of `bazel aquery`.

This tool requires that a build has been completed for all targets in the
query. Generated files such as proto headers and tablegen output are
configuration-specific, so use the same Bazel options for the build and the
action query. If clangd gets out of date, rebuild, regenerate
`compile_commands.json`, or both.

For a debug compilation database with C/C++ assertions enabled, pass `-c dbg`
to both commands. To keep another compilation mode while enabling assertions,
pass `--copt=-UNDEBUG` to both commands so `NDEBUG` is undefined.

Bazel can compile the same source in a target configuration (for the requested
outputs) and an exec configuration (for tools run during the build). Exec
configurations use tool-building flags and typically have `-exec-` in their
`bazel-out` paths. When both exist, this tool keeps the non-exec command so
clangd sees the flags used to compile the target code.

Example usage:
  bazel build -c dbg //xla/... -k
  bazel aquery "mnemonic(CppCompile, //xla/...)" -c dbg --output=jsonproto | \
      python3 build_tools/lint/generate_compile_commands.py
"""

import dataclasses
import json
import logging
import pathlib
import sys
from typing import Any

_JSONDict = dict[Any, Any]  # Approximates parsed JSON

_DISALLOWED_ARGS = frozenset(["-fno-canonical-system-headers"])
_XLA_SRC_ROOT = pathlib.Path(__file__).absolute().parent.parent.parent


def _setup_external_symlink(xla_src_root: pathlib.Path) -> None:
  """Makes Bazel external repositories visible from the source root."""
  external = xla_src_root / "external"
  if external.is_symlink() and not external.exists():
    logging.info("Removing dangling external symlink %s", external)
    external.unlink()

  if not external.exists():
    # bazel-out points to <output_base>/execroot/<workspace>/bazel-out.
    bazel_external = xla_src_root / "bazel-out" / "../../../external"
    logging.info("Symlinking %s to %s", external, bazel_external)
    external.symlink_to(bazel_external)


@dataclasses.dataclass
class CompileCommand:
  """Represents a compilation command with options on a specific file."""

  file: str
  arguments: list[str]

  @classmethod
  def from_args_list(cls, args_list: list[str]) -> "CompileCommand":
    """Alternative constructor which uses the args_list from `bazel aquery`.

    This collects arguments and the file being run on from the output of
    `bazel aquery`. Also filters out arguments which break clang-tidy.

    Arguments:
      args_list: List of arguments generated by `bazel aquery`

    Returns:
      The corresponding ClangTidyCommand.
    """
    cc_file = None
    filtered_args = []

    for arg in args_list:
      if arg in _DISALLOWED_ARGS:
        continue

      if arg.endswith(".cc"):
        cc_file = arg

      # Split generated commands, because otherwise they get wrapped
      # into "command with spaces" when passed to clangd, and clangd
      # can't parse them correctly.
      for s in arg.split(" "):
        filtered_args.append(s)

    return cls(cc_file, filtered_args)  # pyrefly: ignore[bad-argument-type]

  def to_dumpable_json(self, directory: str) -> _JSONDict:
    return {
        "directory": directory,
        "file": self.file,
        "arguments": self.arguments,
    }


def _is_exec_configuration(command: CompileCommand) -> bool:
  return any("-exec-" in arg for arg in command.arguments)


def extract_compile_commands(
    parsed_aquery_output: _JSONDict,
) -> list[CompileCommand]:
  """Gathers compile commands to run from `bazel aquery` JSON output.

  Arguments:
    parsed_aquery_output: Parsed JSON representing the output of `bazel aquery
      --output=jsonproto`.

  Returns:
    The list of CompileCommands that should be executed.
  """
  actions = parsed_aquery_output["actions"]

  commands = []
  command_indices: dict[str, int] = {}
  for action in actions:
    command = CompileCommand.from_args_list(action["arguments"])
    if command.file is None:
      commands.append(command)
      continue

    index = command_indices.get(command.file)
    if index is None:
      command_indices[command.file] = len(commands)
      commands.append(command)
    elif _is_exec_configuration(commands[index]) and not _is_exec_configuration(
        command
    ):
      commands[index] = command

  return commands


def main():
  # Setup logging
  logging.basicConfig()
  logging.getLogger().setLevel(logging.INFO)

  # Setup external symlink if necessary so headers can be found in include paths
  _setup_external_symlink(_XLA_SRC_ROOT)

  logging.info("Reading `bazel aquery` output from stdin...")
  parsed_aquery_output = json.loads(sys.stdin.read())

  commands = extract_compile_commands(parsed_aquery_output)

  with (_XLA_SRC_ROOT / "compile_commands.json").open("w") as f:
    json.dump(
        [
            command.to_dumpable_json(directory=str(_XLA_SRC_ROOT))
            for command in commands
        ],
        f,
    )


if __name__ == "__main__":
  main()