#!/usr/bin/env python3

"""Rewrite BUILD.gn files to use jumbo_source_set targets.

The script walks a directory tree, finds BUILD.gn files, replaces
source_set/static_library/component target declarations with
jumbo_source_set, and keeps the jumbo import in alphabetical order
with the existing imports near the top of each file.

[AI Disclosure]
This script was generated by GitHub Copilot. The generated code has
been verified and tested for correctness.
"""

from __future__ import annotations

import argparse
import re
from pathlib import Path


JUMBO_IMPORT = 'import("//build/config/jumbo.gni")'
TARGET_RE = re.compile(r'^(\s*)(source_set|static_library|component)(\s*\()')
TARGET_DEF_RE = re.compile(r'^(\s*)(source_set|static_library|component)\s*\(\s*"([^"]+)"\s*\)')
JUMBO_TARGET_MAP = {
    "source_set": "jumbo_source_set",
    "static_library": "jumbo_static_library",
    "component": "jumbo_component",
}
TOP_LEVEL_TARGET_RE = re.compile(
    r'^\s*(jumbo_source_set|jumbo_static_library|jumbo_component|source_set|static_library|component)\s*\('
)
SCOPE_TARGET_RE = re.compile(
    r'^\s*(jumbo_source_set|jumbo_static_library|jumbo_component|source_set|static_library|component)\s*\('
)
LEADING_WS_RE = re.compile(r'^(\s*)')
ASSIGNMENT_START_RE = re.compile(r'^\s*([A-Za-z_][A-Za-z0-9_]*)\s*=')
SKIP_PATCH_ASSIGNMENT_BLACKLIST = frozenset({"testonly", "test_only", "visibility"})

# Temporary hardcoded ignore list. One GN label per line.
# Examples:
#   //components/xx1/xx1:yyyy
#   //components/xx2/xx2
IGNORE_LIST_TEXT = """
//components/dbus/utils
//components/dbus/xdg
//components/language_detection/core
//components/unexportable_keys/mojom
//components/optimization_guide/content/browser:browser
//components/collaboration/internal:core_internal
//components/segmentation_platform/embedder/default_model
//components/saved_tab_groups/internal:tab_group_sync_bridge
//components/ui_devtools/views
//components/safe_browsing/content/renderer/phishing_classifier
//components/spellcheck/renderer
//chrome/browser/ui/color
//chrome/browser/glic:glic
//chrome/browser/ui/views/tabs/projects
"""


def detect_newline(text: str) -> str:
    return "\r\n" if "\r\n" in text else "\n"


def is_header_line(line: str) -> bool:
    stripped = line.strip()
    return not stripped or stripped.startswith("#") or stripped.startswith("import(")


def brace_depth_delta(line: str) -> int:
    return line.count("{") - line.count("}")


def find_leading_assignment_end(lines: list[str]) -> int:
    """Return the first index after leading top-level assignments/comments."""
    index = 0
    while index < len(lines):
        stripped = lines[index].strip()
        if not stripped or stripped.startswith("#"):
            index += 1
            continue

        if not ASSIGNMENT_START_RE.match(lines[index]):
            break

        # Consume the whole assignment statement. Multi-line list/string
        # assignments are handled by balancing [], () and {} counts.
        bracket_depth = 0
        paren_depth = 0
        brace_depth = 0
        while index < len(lines):
            line = lines[index]
            bracket_depth += line.count("[") - line.count("]")
            paren_depth += line.count("(") - line.count(")")
            brace_depth += line.count("{") - line.count("}")
            index += 1
            if bracket_depth <= 0 and paren_depth <= 0 and brace_depth <= 0:
                break

    return index


def has_blacklisted_top_level_assignment(lines: list[str]) -> bool:
    block_depth = 0
    for line in lines:
        stripped = line.strip()
        if block_depth == 0 and stripped and not stripped.startswith("#"):
            match = ASSIGNMENT_START_RE.match(line)
            if match and match.group(1).lower() in SKIP_PATCH_ASSIGNMENT_BLACKLIST:
                return True
        block_depth += brace_depth_delta(line)
    return False


def starts_with_leading_assignment(lines: list[str]) -> bool:
    for line in lines:
        stripped = line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        return ASSIGNMENT_START_RE.match(line) is not None
    return False


def first_code_index(lines: list[str]) -> int:
    for index, line in enumerate(lines):
        stripped = line.strip()
        if stripped and not stripped.startswith("#"):
            return index
    return 0


def first_top_level_target_index(lines: list[str]) -> int | None:
    block_depth = 0
    for index, line in enumerate(lines):
        stripped = line.strip()
        if block_depth == 0 and stripped and not stripped.startswith("#"):
            if TOP_LEVEL_TARGET_RE.match(line):
                return index
        block_depth += brace_depth_delta(line)
    return None


def first_scope_target_index(lines: list[str], scope_depth: int) -> int | None:
    block_depth = 0
    for index, line in enumerate(lines):
        stripped = line.strip()
        if block_depth == scope_depth and stripped and not stripped.startswith("#"):
            if SCOPE_TARGET_RE.match(line):
                return index
        block_depth += brace_depth_delta(line)
    return None


def compute_line_scopes(lines: list[str]) -> list[tuple[int, ...]]:
    scopes: list[tuple[int, ...]] = []
    stack: list[int] = []
    next_scope_id = 1

    for line in lines:
        scopes.append(tuple(stack))
        for ch in line:
            if ch == "{":
                stack.append(next_scope_id)
                next_scope_id += 1
            elif ch == "}":
                if stack:
                    stack.pop()

    return scopes


def first_scope_target_index_with_scopes(
    lines: list[str],
    scopes: list[tuple[int, ...]],
    scope_key: tuple[int, ...],
) -> int | None:
    for index, line in enumerate(lines):
        stripped = line.strip()
        if scopes[index] == scope_key and stripped and not stripped.startswith("#"):
            if SCOPE_TARGET_RE.match(line):
                return index
    return None


def find_scope_jumbo_import_index_with_scopes(
    lines: list[str],
    scopes: list[tuple[int, ...]],
    scope_key: tuple[int, ...],
) -> int | None:
    first_target = first_scope_target_index_with_scopes(lines, scopes, scope_key)
    if first_target is None:
        return None

    import_indexes: list[int] = []
    for index in range(first_target):
        if scopes[index] != scope_key:
            continue
        if lines[index].strip().startswith("import("):
            import_indexes.append(index)

    if import_indexes:
        for index in import_indexes:
            if lines[index].strip() >= JUMBO_IMPORT:
                return index
        return import_indexes[-1] + 1

    return first_target


def scope_import_indent_with_scopes(
    lines: list[str],
    scopes: list[tuple[int, ...]],
    scope_key: tuple[int, ...],
) -> str:
    first_target = first_scope_target_index_with_scopes(lines, scopes, scope_key)
    if first_target is None:
        return ""
    match = LEADING_WS_RE.match(lines[first_target])
    return "" if match is None else match.group(1)


def find_scope_jumbo_import_index(lines: list[str], scope_depth: int) -> int | None:
    first_target = first_scope_target_index(lines, scope_depth)
    if first_target is None:
        return None

    import_indexes: list[int] = []
    block_depth = 0
    for index, line in enumerate(lines):
        stripped = line.strip()
        if block_depth == scope_depth and index < first_target and stripped.startswith("import("):
            import_indexes.append(index)
        block_depth += brace_depth_delta(line)

    if import_indexes:
        for index in import_indexes:
            if lines[index].strip() >= JUMBO_IMPORT:
                return index
        return import_indexes[-1] + 1

    return first_target


def scope_import_indent(lines: list[str], scope_depth: int) -> str:
    first_target = first_scope_target_index(lines, scope_depth)
    if first_target is None:
        return ""
    match = LEADING_WS_RE.match(lines[first_target])
    return "" if match is None else match.group(1)


def find_jumbo_import_index(lines: list[str]) -> int:
    first_target = first_top_level_target_index(lines)

    # Keep imports before leading top-level assignments, but after comments.
    if starts_with_leading_assignment(lines):
        code_index = first_code_index(lines)
        if first_target is not None:
            return min(code_index, first_target)
        return code_index

    min_insert_index = find_leading_assignment_end(lines)
    import_indexes: list[int] = []
    block_depth = 0

    for index, line in enumerate(lines):
        stripped = line.strip()
        if block_depth > 0:
            block_depth += brace_depth_delta(line)
            continue

        if not stripped or stripped.startswith("#"):
            continue

        if stripped.startswith("import(") and (first_target is None or index < first_target):
            import_indexes.append(index)

        block_depth += brace_depth_delta(line)

    candidate_indexes = [i for i in import_indexes if i >= min_insert_index]

    if candidate_indexes:
        for index in candidate_indexes:
            if lines[index].strip() >= JUMBO_IMPORT:
                return index
        return candidate_indexes[-1] + 1

    if first_target is not None:
        return first_target

    return min_insert_index


def parse_ignore_list(ignore_items: list[str]) -> tuple[dict[str, set[str]], set[str], dict[str, list[str]]]:
    """Parse ignore items.

    Input format examples:
    - //components/xx1/xx1:yyyy  -> skip source_set("yyyy") in that BUILD.gn
    - //components/xx2/xx2       -> skip whole BUILD.gn
    """
    skip_targets_by_module: dict[str, set[str]] = {}
    skip_modules: set[str] = set()
    grouped: dict[str, list[str]] = {}

    for raw in ignore_items:
        item = raw.strip()
        if not item:
            continue

        if item.startswith("//"):
            item = item[2:]
        item = item.strip("/")
        if not item:
            continue

        module, sep, target = item.partition(":")
        module = module.strip("/")
        if not module:
            continue

        group_key, _, group_rest = module.partition("/")
        value_base = group_rest if group_rest else module
        value = value_base if not sep else f"{value_base}:{target}"
        grouped.setdefault(group_key, []).append(value)

        if not sep:
            skip_modules.add(module)
            continue

        target_name = target.strip()
        if not target_name:
            skip_modules.add(module)
            continue

        skip_targets_by_module.setdefault(module, set()).add(target_name)

    return skip_targets_by_module, skip_modules, grouped


def parse_ignore_list_text(ignore_text: str) -> list[str]:
    items: list[str] = []
    for raw in ignore_text.splitlines():
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        items.append(line)
    return items


def module_path_from_build_file(path: Path, root: Path) -> str:
    rel = path.relative_to(root).as_posix()
    if not rel.endswith("/BUILD.gn"):
        return ""
    return rel[: -len("/BUILD.gn")]


def transform_build_gn(text: str, skip_target_names: set[str]) -> tuple[str, bool]:
    lines = text.splitlines(keepends=True)
    newline = detect_newline(text)
    line_scopes = compute_line_scopes(lines)

    # Skip entire file when blacklisted top-level assignments are present.
    if has_blacklisted_top_level_assignment(lines):
        return text, False

    changed = False
    had_target = False
    kept_lines: list[str] = []
    kept_scopes: list[tuple[int, ...]] = []
    rewritten_scopes: set[tuple[int, ...]] = set()

    for line, line_scope in zip(lines, line_scopes):

        if line.strip() == JUMBO_IMPORT:
            changed = True
            kept_lines.append(line)
            kept_scopes.append(line_scope)
            continue

        target_def_match = TARGET_DEF_RE.match(line)
        if target_def_match and target_def_match.group(3) in skip_target_names:
            kept_lines.append(line)
            kept_scopes.append(line_scope)
            continue

        match = TARGET_RE.match(line)
        if match:
            jumbo_target = JUMBO_TARGET_MAP[match.group(2)]
            kept_lines.append(f"{match.group(1)}{jumbo_target}{line[match.end(2):]}")
            kept_scopes.append(line_scope)
            changed = True
            had_target = True
            rewritten_scopes.add(line_scope)
            continue

        kept_lines.append(line)
        kept_scopes.append(line_scope)

    if not had_target:
        return text, False

    no_jumbo_lines: list[str] = []
    no_jumbo_scopes: list[tuple[int, ...]] = []
    for line, line_scope in zip(kept_lines, kept_scopes):
        if line.strip() == JUMBO_IMPORT and line_scope in rewritten_scopes:
            continue
        no_jumbo_lines.append(line)
        no_jumbo_scopes.append(line_scope)

    insert_points: list[tuple[int, str]] = []
    for scope_key in sorted(rewritten_scopes):
        insert_at = find_scope_jumbo_import_index_with_scopes(
            no_jumbo_lines,
            no_jumbo_scopes,
            scope_key,
        )
        if insert_at is not None:
            indent = scope_import_indent_with_scopes(
                no_jumbo_lines,
                no_jumbo_scopes,
                scope_key,
            )
            insert_points.append((insert_at, f"{indent}{JUMBO_IMPORT}{newline}"))

    for insert_at, import_line in sorted(insert_points, key=lambda item: item[0], reverse=True):
        no_jumbo_lines.insert(insert_at, import_line)
        no_jumbo_scopes.insert(insert_at, ())

    return "".join(no_jumbo_lines), True


def parse_subdirs(subdirs_items: list[str], root: Path) -> list[Path]:
    if not subdirs_items:
        return [root]

    subdirs: list[Path] = []
    seen: set[str] = set()
    for raw in subdirs_items:
        item = raw.strip().strip("/")
        if not item:
            continue
        candidate = Path(item)
        if candidate.is_absolute() or ".." in candidate.parts:
            raise SystemExit(f"invalid subdir path: {raw}")
        full = root / candidate
        if not full.is_dir():
            raise SystemExit(f"subdir is not a directory: {full}")
        key = full.resolve().as_posix()
        if key in seen:
            continue
        seen.add(key)
        subdirs.append(full)

    if not subdirs:
        return [root]
    return subdirs


def iter_build_files(scan_roots: list[Path]):
    files: set[Path] = set()
    for scan_root in scan_roots:
        for path in scan_root.rglob("BUILD.gn"):
            if path.is_file():
                files.add(path)
    for path in sorted(files, key=lambda p: p.as_posix()):
        yield path


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Rewrite BUILD.gn files for jumbo builds."
    )
    parser.add_argument("root", help="Directory to scan for BUILD.gn files.")
    parser.add_argument(
        "--check",
        action="store_true",
        help="Do not write files; exit 1 if any file would change.",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print rewritten file paths.",
    )
    parser.add_argument(
        "--subdirs",
        action="append",
        default=[],
        help=(
            "Subdirectory under root to scan. Repeat this option to pass "
            "multiple subdirectories, e.g. --subdirs components --subdirs base/test. "
            "If omitted, scan all subdirectories."
        ),
    )
    args = parser.parse_args()

    root = Path(args.root)
    if not root.is_dir():
        raise SystemExit(f"root is not a directory: {root}")

    ignore_items = parse_ignore_list_text(IGNORE_LIST_TEXT)

    subdirs_items = args.subdirs

    scan_roots = parse_subdirs(subdirs_items, root)
    if args.verbose and subdirs_items:
        print(f"scan subdirs: {[p.relative_to(root).as_posix() for p in scan_roots]}")

    skip_targets_by_module, skip_modules, grouped = parse_ignore_list(ignore_items)
    if args.verbose and grouped:
        print(f"parsed ignore groups: {grouped}")

    changed_files = []
    for path in iter_build_files(scan_roots):
        module_path = module_path_from_build_file(path, root)
        module_candidates = {module_path}
        if module_path:
            module_candidates.add(f"{root.name}/{module_path}")

        if any(candidate in skip_modules for candidate in module_candidates):
            continue

        skip_targets: set[str] = set()
        for candidate in module_candidates:
            skip_targets.update(skip_targets_by_module.get(candidate, set()))
        original = path.read_text(encoding="utf-8")
        rewritten, changed = transform_build_gn(original, skip_targets)
        if not changed:
            continue
        changed_files.append(path)
        if args.verbose:
            print(path)
        if not args.check:
            path.write_text(rewritten, encoding="utf-8")

    if args.check and changed_files:
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
