#!/usr/bin/env python3
# Copyright 2026 Google LLC
#
# 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.

"""Set up clangd-based C++ code navigation for TorchTPU repository.

Run it once to bootstrap, and re-run it regularly on your local machine after
BUILD files, compiler flags, or generated sources have significant changes. Each
run performs the following steps:
  1. Symlink repo-root external/ to Bazel's external repositories so that
     include paths in the compilation database resolve. Skipped once the
     symlink is in place.
  2. Regenerate compile_commands.json in the repository root.
  3. Build the C++ targets so generated files (e.g. proto headers) exist.
In cases where users have run `bazel build` on targets of interest recently,
generated files may already exist on disk. Users can use `--no-build` option to
skip the build step.

WARNING: Skipping actual builds may lead to missing generated files, so clangd
might report errors on files that depend on them.

Example usage:
  # Full setup / refresh with build.
  $ tools/setup_clangd.py

  # Fast database-only refresh without build.
  $ tools/setup_clangd.py --no-build

  # Refresh for specific combined target patterns.
  $ tools/setup_clangd.py //csrc/common/... //src/torch_tpu/ops/...
"""

from __future__ import annotations

import argparse
from collections.abc import Iterable, Mapping, Sequence
import dataclasses
import json
import logging  # PYTHON_LOGGING_OK=tool setup
import pathlib
import subprocess
import sys
from typing import Any, Final

# GCC-only flags that clang-based tools cannot parse.
_UNSUPPORTED_ARGS: Final[frozenset[str]] = frozenset([
    "-fno-canonical-system-headers",
])

_DEFAULT_TARGET_PATTERN: Final[str] = "//..."
_REPO_ROOT: Final[pathlib.Path] = pathlib.Path(__file__).resolve().parents[1]


@dataclasses.dataclass(frozen=True)
class CompileCommand:
  """Represents a C++ compilation command with options on a specific file.

  Attributes:
    file: The C++ source file path being compiled, or None if not found.
    arguments: Sequence of command-line arguments for the compilation.
  """

  file: str | None
  arguments: Sequence[str]

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

    It collects arguments and the file being compiled from the output of
    `bazel aquery`, and filters out arguments incompatible with clangd.

    Args:
      args_list: Arguments generated by `bazel aquery`.

    Returns:
      The corresponding `CompileCommand` object.
    """
    source_file = None
    filtered_args = []

    for i, arg in enumerate(args_list):
      if arg in _UNSUPPORTED_ARGS:
        continue

      # The argument after `-c` is the file being compiled.
      if i > 0 and args_list[i - 1] == "-c":
        source_file = arg
      elif source_file is None and arg.endswith(".cc"):
        source_file = arg

      # Split multi-token arguments (e.g. produced by --per_file_copt),
      # otherwise clangd cannot parse them correctly.
      filtered_args.extend(arg.split())

    return cls(source_file, filtered_args)

  def to_dict(self, directory: str) -> Mapping[str, Any]:
    return {
        "directory": directory,
        "file": self.file,
        "arguments": self.arguments,
    }


def _extract_compile_commands(
    parsed_aquery_output: Mapping[str, Any],
) -> Sequence[CompileCommand]:
  """Extracts compile commands from `bazel aquery` JSON output.

  Deduplicates by source file (keeping the first occurrence), since the same
  file may be compiled several times under different configurations and clangd
  only uses one command per file.

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

  Returns:
    The sequence of `CompileCommand` objects for the compilation database.
  """
  commands = []
  seen = set()

  for action in parsed_aquery_output.get("actions", []):
    command = CompileCommand.from_args_list(action["arguments"])
    if command.file is None or command.file in seen:
      continue
    seen.add(command.file)
    commands.append(command)

  return commands


def _setup_external_dir_symlink() -> None:
  """Symlinks from repo-root external/ directory to Bazel's external repositories.

  Compile arguments reference external repositories via paths relative to
  Bazel's execution root (e.g. -isystem external/abseil-cpp+). Mirroring
  the output base's external/ directory into the repository root makes
  those paths resolve from `directory` in compile_commands.json.

  Per https://bazel.build/remote/output-directories#layout-diagram, the external
  directory under the output base always contains every fetched repository.
  """
  repo_external_dir = _REPO_ROOT / "external"

  # A dangling symlink (e.g. the output base was cleaned) will fall through and
  # be recreated.
  if repo_external_dir.is_symlink() and repo_external_dir.resolve().is_dir():
    logging.info("external directory symlink is already set up.")
    return

  if repo_external_dir.exists() and not repo_external_dir.is_symlink():
    logging.error(
        "%s exists but is not a symlink. Exiting without overwriting it.",
        repo_external_dir,
    )
    sys.exit(1)

  if repo_external_dir.is_symlink():
    repo_external_dir.unlink()

  output_base = subprocess.run(
      ["bazel", "info", "output_base"],
      stdout=subprocess.PIPE,
      cwd=_REPO_ROOT,
      check=True,
      text=True,
  ).stdout.strip()

  external_dir = pathlib.Path(output_base) / "external"
  repo_external_dir.symlink_to(external_dir)


def _bazel_aquery(target_patterns: Iterable[str]) -> Mapping[str, Any]:
  """Runs Bazel aquery to obtain C++ compile actions for the given targets."""
  query_expr = " union ".join(target_patterns)

  result = subprocess.run(
      [
          "bazel",
          "aquery",
          f'mnemonic("CppCompile", {query_expr})',
          "--output=jsonproto",
          # For each action, we only care about its `arguments`.
          # Skip listing artifacts shrinks the output considerably.
          "--include_artifacts=false",
      ],
      stdout=subprocess.PIPE,
      cwd=_REPO_ROOT,
      check=True,
  )

  return json.loads(result.stdout)


def _generate_compile_commands(target_patterns: Iterable[str]) -> None:
  """Generates compile_commands.json in the repository root."""
  parsed_aquery_output = _bazel_aquery(target_patterns)
  commands = _extract_compile_commands(parsed_aquery_output)
  if not commands:
    logging.warning(
        "No C++ compile actions found for %s; compile_commands.json will be"
        " empty.",
        " ".join(target_patterns),
    )

  output_path = _REPO_ROOT / "compile_commands.json"
  with output_path.open("w") as f:
    json.dump(
        [command.to_dict(directory=str(_REPO_ROOT)) for command in commands],
        f,
        indent=2,
    )

  logging.info(
      "Successfully wrote %d compile commands to %s", len(commands), output_path
  )


def _bazel_build(target_patterns: Iterable[str]) -> None:
  """Builds specified target patterns so generated sources and headers exist on disk.

  Certain compilation inputs (proto headers, genrule outputs) only exist under
  bazel-out/ after actual builds. Without them, clangd cannot index the files
  that depend on them. Bazel keeps this incremental, so refreshes are cheap.

  Args:
    target_patterns: Bazel target patterns to build.
  """
  # Use --keep_going to build as much as possible by downgrading build failures
  # to warnings. A stale but mostly complete index is more useful than none.
  result = subprocess.run(
      ["bazel", "build", "--keep_going", *target_patterns],
      cwd=_REPO_ROOT,
      check=False,
  )

  if result.returncode != 0:
    logging.warning(
        "`bazel build` failed for some targets. The generated files belonging"
        " to the failed targets may be missing from the index."
    )


def _setup_clangd(args: argparse.Namespace) -> None:
  """Sets up clangd configurations following the given arguments."""
  logging.info("[1/3] Setting up the external directory symlink...")
  _setup_external_dir_symlink()

  logging.info("[2/3] Generating compile_commands.json...")
  _generate_compile_commands(args.target_patterns)

  if args.no_build:
    logging.info("[3/3] Skipping the Bazel build (--no-build is specified)...")
  else:
    logging.info(
        "[3/3] Building %s so generated files exist (pass --no-build to"
        " skip)...",
        " ".join(args.target_patterns),
    )
    _bazel_build(args.target_patterns)


def main() -> None:
  logging.basicConfig(format="%(levelname)s: %(message)s")
  logging.getLogger().setLevel(logging.INFO)

  parser = argparse.ArgumentParser(
      description=__doc__,
      formatter_class=argparse.RawDescriptionHelpFormatter,
  )
  parser.add_argument(
      "target_patterns",
      nargs="*",
      default=[_DEFAULT_TARGET_PATTERN],
      metavar="target_pattern",
      help=(
          "Bazel target patterns (or query expressions) whose C++ compile"
          " commands are included in the database. Multiple patterns are"
          f" combined (default: {_DEFAULT_TARGET_PATTERN})."
      ),
  )
  parser.add_argument(
      "--no-build",
      action="store_true",
      help=(
          "Skip actual Bazel build. The compilation database still"
          " refreshes, but generated files (e.g. proto headers) that were"
          " never built will be missing, so clangd may report errors in"
          " files that depend on them."
      ),
  )
  args = parser.parse_args()

  _setup_clangd(args)

  logging.info(
      "clangd will pick up compile_commands.json soon and start indexing"
      " automatically."
  )


if __name__ == "__main__":
  main()