# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import annotations

import collections
import contextlib
import dataclasses
import enum
import functools
import itertools
import pathlib
import re

# Almost every sysroot just uses #include <>, but fuchsia uses #include ""
# sometimes.
_INCLUDES = re.compile(r'#\s*(?:include|import)(_next)?\s*["<]([^>"]+)[>"]')


class CompileStatus(enum.Enum):
  NotCompiled = 1
  Success = 2
  Failure = 3


class IncludeDir(enum.Enum):
  # Ordered by include order for clang
  LibCxx = 1
  Builtin = 2
  SysrootModule = 3
  Sysroot = 4
  Framework = 5

  def __lt__(self, other):
    return self.value < other.value


@dataclasses.dataclass
class Header:
  include_dir: IncludeDir
  # The string to use in your #include statement to get this file.
  rel: str
  # The absolute path to the file
  # This should be filled by the time the compiler finishes.
  abs: pathlib.Path = None
  # Prev and next come from include_next / include_prev
  prev: None | Header = None
  next: None | Header = None
  root_module: None | str = None
  textual: bool = False
  umbrella: bool = False
  compile_status: CompileStatus = CompileStatus.NotCompiled

  deps: list[Header] = dataclasses.field(default_factory=list)
  direct_deps: set[Header] = dataclasses.field(default_factory=set)
  # Here, None means no exports, and the empty list means 'export *'
  # We default to exporting all to preserve the behaviour of includes.
  exports: None | list[str] = dataclasses.field(default_factory=list)

  # Kwargs that will end up on the BUILD.gn targets.
  kwargs: dict[str, list[str]] = dataclasses.field(
      default_factory=lambda: collections.defaultdict(list))

  def __hash__(self):
    return hash((self.include_dir, self.rel))

  def __eq__(self, other):
    if isinstance(other, Header):
      return (self.include_dir, self.rel) == (other.include_dir, other.rel)
    else:
      # This allows you to write (Sysroot, 'foo.h') in set[Header]
      return (self.include_dir, self.rel) == other

  def __lt__(self, other):
    return (self.include_dir, self.rel) < (other.include_dir, other.rel)

  @property
  def pretty_name(self):
    return f'{self.include_dir.name}/{self.rel}'

  def __repr__(self):
    return self.pretty_name

  @property
  def submodule_name(self):
    return self.rel.replace('.', '_').replace('/', '_').replace('-', '_')

  @functools.cached_property
  def content(self) -> str:
    return self.abs.read_text(errors='ignore')

  def calculate_direct_deps(self, includes: dict[str, Header],
                            sysroot: pathlib.Path) -> set[Header]:

    direct = set()
    found_includes = _INCLUDES.findall(self.content)

    def find_include(is_next, include) -> bool:
      header = None
      first = includes.get(include, None)
      if first is not None:
        # When modules are enabled, #include_next<foo.h> from any file other
        # than foo.h is treated the same as #include <foo.h>.
        if not is_next or (is_next and self.rel != include):
          header = first
        elif self.next is not None:
          header = self.next

        # It might have been conditionally included.
        if header is not None and header in self.deps:
          direct.add(header)
          return True
      return False

    for is_next, include in found_includes:
      if not find_include(is_next, include):
        # This is required, because, for example, libcxx's threading includes
        # pthread.h, but the include scanner sees pthread/pthread.h (the
        # symlink target).
        with contextlib.suppress(OSError, FileNotFoundError):
          find_include(is_next, str((sysroot / include).readlink()))

    return direct

  @functools.cache
  def _required_deps(self) -> tuple[set[Header], set[Header]]:
    nontextual = set()
    textual = set()
    todo = [self]
    while todo:
      hdr = todo.pop()
      for dep in hdr.direct_deps:
        if dep.textual and dep not in textual:
          todo.append(dep)
          textual.add(dep)
        elif not dep.textual:
          nontextual.add(dep)
    return nontextual, textual

  @property
  def required_deps(self) -> set[Header]:
    """The header files required to be built before we can build."""
    return self._required_deps()[0]

  @property
  def required_textual_deps(self) -> set[Header]:
    """The textual header files we directly include.

    This includes textual headers included via other textual headers"""
    return self._required_deps()[1]

  def find_loop(self) -> list[Header] | None:
    """Finds a loop of #includes, if it exists."""
    chain = [self]
    has_chain = True
    while has_chain:
      has_chain = False
      if self in chain[-1].direct_deps:
        return chain + [self]
      for dep in chain[-1].direct_deps:
        if dep not in chain and self in dep.deps:
          chain.append(dep)
          has_chain = True
          break
    # It shouldn't be possible to have a node that has you as a transitive
    # dep without having a dep that has you as a transitive dep.
    assert len(chain) == 1


def calculate_rdeps(headers: list[Header]) -> dict[Header, list[Header]]:
  """Calculates a reverse dependency graph"""
  rdeps = collections.defaultdict(list)
  for header in headers:
    for dep in header.deps:
      rdeps[dep].append(header)
  return rdeps


def all_headers(graph: dict[str, Header]):
  """Iterates through all headers in a graph."""
  for header in graph.values():
    while header is not None:
      yield header
      header = header.next


@dataclasses.dataclass
class Target:
  """Represents a single clang module / gn target."""
  include_dir: IncludeDir
  name: str
  headers: list[Header] = dataclasses.field(default_factory=list)

  def __lt__(self, other):
    return self.name < other.name

  def __eq__(self, other):
    return self.name == other.name

  def __hash__(self):
    return hash(self.name)

  @property
  def kwargs(self) -> dict[str, set[str]]:
    """The kwargs associated with a build target.

    eg. if you have kwargs = {"defines": ["FOO"]}, then it outputs:

    target_type(target.name) {
      defines = ["FOO"]
    }
    """
    kwargs = collections.defaultdict(set)
    for header in self.headers:
      for single in header.group:
        for dep in {single} | single.required_textual_deps:
          for k, v in dep.kwargs.items():
            kwargs[k].update(v)
    return kwargs

  @property
  def header_deps(self) -> set[Header]:
    direct_deps = set()
    for hdr in self.headers:
      direct_deps.update(hdr.required_deps)
    return direct_deps

  @property
  def public_deps(self) -> list[str]:
    return sorted(
        set([
            hdr.root_module for hdr in self.header_deps
            if hdr.root_module is not None and hdr.root_module != self.name
        ]))


def run_build(graph: dict[str, Header]) -> list[Target]:
  """Calculates the correct way to run a build."""
  unbuilt_modules: dict[str, list[Header]] = collections.defaultdict(list)
  unbuilt_headers: set[Header] = set()

  for header in all_headers(graph):
    if not header.textual:
      if header.root_module is None:
        unbuilt_headers.add(header)
      else:
        unbuilt_modules[header.root_module].append(header)
    header.rdeps = set()
    # A group is a set of headers that must be compiled together, because they
    # form a dependency loop.
    header.group = [header]
    header.mod_deps = set()
    header.unbuilt_deps = set(
        dep for dep in header.required_deps
        # You don't need to wait for a dependency within the same module.
        if (header.root_module is None or header.root_module != dep.root_module)
        and dep != header)

  # Perform a union find to find all dependency loops.
  # Since we can easily tell if a given edge represents a dependency loop, we
  # simply union together all pairs of nodes on loop edges.
  # We perform a simple form of union find where we don't bother with rank or
  # size, and only do path compression (always on the lexicographically first
  # header). It's slower but still plenty fast and simplifies things.
  parents = {}

  def find(header):
    if header in parents:
      # Optimization: path compression
      parents[header] = find(parents[header])
      return parents[header]
    else:
      return header

  for header in sorted(unbuilt_headers):
    for dep in header.required_deps:
      if dep > header and header in dep.deps:
        assert header.include_dir == IncludeDir.Sysroot and dep.include_dir == IncludeDir.Sysroot, (
            header, dep)
        # Perform the 'union' operation.
        x, y = sorted([find(header), find(dep)])
        if x != y:
          parents[y] = x

  loops = collections.defaultdict(list)
  for header in unbuilt_headers:
    loops[find(header)].append(header)

  for headers in loops.values():
    # Not a loop
    if len(headers) == 1:
      continue
    headers.sort()
    headers[0].group = headers
    headers[0].unbuilt_deps = set.union(
        *[header.unbuilt_deps for header in headers]) - set(headers)
    for header in headers[1:]:
      unbuilt_headers.remove(header)

  for header in all_headers(graph):
    for dep in header.unbuilt_deps:
      dep.rdeps.add(header)

  build_gn = []

  for i in itertools.count():
    # Try and build any buildable modules from modulemaps.
    while True:
      n_remaining = len(unbuilt_modules)
      for mod, headers in list(unbuilt_modules.items()):
        if not any(header.unbuilt_deps for header in headers):
          build_gn.append(
              Target(
                  include_dir=headers[0].include_dir,
                  name=mod,
                  headers=sorted(headers),
              ))
          del unbuilt_modules[mod]
          for header in headers:
            for rdep in header.rdeps:
              rdep.mod_deps.add(header.root_module)
              rdep.unbuilt_deps.remove(header)

      if n_remaining == len(unbuilt_modules):
        break

    # Try and build any builtable sysroot modules
    sysroot_mod = f'sys_stage{i + 1}'
    build_gn.append(Target(
        include_dir=IncludeDir.Sysroot,
        name=sysroot_mod,
    ))
    while True:
      n_remaining = len(unbuilt_headers)
      for header in list(unbuilt_headers):
        if not header.unbuilt_deps:
          build_gn[-1].headers.append(header)
          unbuilt_headers.remove(header)
          for header in header.group:
            header.root_module = sysroot_mod
            for rdep in header.rdeps:
              rdep.mod_deps.add(header.root_module)
              rdep.unbuilt_deps.remove(header)

      if n_remaining == len(unbuilt_headers):
        break

    build_gn[-1].headers.sort()
    if not build_gn[-1].headers:
      break

  build_gn.pop()

  if not unbuilt_modules and not unbuilt_headers:
    # Success. Everything is built
    return build_gn
  else:
    print(
        "Dependency loop in sysroot. You probably want to make one of them textual."
    )
    print("The following headers are in a dependency loop:")
    seen = set()
    for header in unbuilt_headers:
      if header not in seen:
        chain = header.find_loop()
        if chain is not None:
          print(' -> '.join([header.pretty_name for header in chain]))
          seen.update(chain)

    # If you get to this point, you probably want a debugger to help understand what the problem is.
    breakpoint()
    exit(1)