generate_compose.py

#!/usr/bin/env python3

# Copyright Institute for Automotive Engineering (ika), RWTH Aachen University
# SPDX-License-Identifier: Apache-2.0

"""Generate template Docker Compose files from ROS 2 launch files."""

from __future__ import annotations

import argparse
import ast
import difflib
import os
import re
import subprocess
import sys
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from pathlib import Path

try:
    from jinja2 import Environment, FileSystemLoader, TemplateNotFound
except ModuleNotFoundError as exc:
    if exc.name == "jinja2":
        Environment = None
        FileSystemLoader = None
        TemplateNotFound = Exception
    else:
        raise


DEFAULT_NAMESPACE = "/"
COMPOSE_DIR = Path("deployment/compose")
COMPOSE_PATH = Path("deployment/compose/docker-compose.yml")
GITLAB_REGISTRY_ENV_NAME = "OPENADS_GITLAB_REGISTRY"
STANDARD_LAUNCH_ARGUMENT_NAMES = ("namespace", "name", "log_level", "use_sim_time", "params")
HOST_EXPOSED_ENV_NAMES = frozenset({"LOG_LEVEL", "USE_SIM_TIME", "ROS_TRACING"})


@dataclass(frozen=True)
class LaunchArgument:
    name: str
    default_value: str
    description: str
    installed_default_path: str | None = None


@dataclass(frozen=True)
class LaunchData:
    node_packages: frozenset[str]
    launch_file_name: str
    arguments: list[LaunchArgument]
    remappable_topic_names: list[str]


@dataclass(frozen=True)
class PackageMetadata:
    name: str
    version: str


@dataclass(frozen=True)
class EnvironmentVariable:
    name: str
    value: str
    host_exposed: bool = False


@dataclass(frozen=True)
class RepoRemote:
    host: str
    port: str | None
    path: str


@dataclass(frozen=True)
class LaunchCommandArgument:
    name: str
    env_name: str
    only_if_set: bool = False


def constant_string(node: ast.AST | None) -> str | None:
    if isinstance(node, ast.Constant) and isinstance(node.value, str):
        return node.value
    return None


def call_name(node: ast.AST) -> str | None:
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        return node.attr
    return None


def call_keyword(call: ast.Call, keyword_name: str) -> ast.AST | None:
    for keyword in call.keywords:
        if keyword.arg == keyword_name:
            return keyword.value
    return None


def installed_share_path(node: ast.AST | None) -> str | None:
    """Resolve os.path.join(get_package_share_directory(...), ...) defaults."""
    if not isinstance(node, ast.Call) or call_name(node.func) != "join" or len(node.args) < 2:
        return None
    if not isinstance(node.func, ast.Attribute) or not isinstance(node.func.value, ast.Attribute):
        return None
    if (
        node.func.value.attr != "path"
        or not isinstance(node.func.value.value, ast.Name)
        or node.func.value.value.id != "os"
    ):
        return None

    package_share_call = node.args[0]
    if (
        not isinstance(package_share_call, ast.Call)
        or call_name(package_share_call.func) != "get_package_share_directory"
        or len(package_share_call.args) != 1
    ):
        return None

    package_name = constant_string(package_share_call.args[0])
    relative_parts = [constant_string(part) for part in node.args[1:]]
    if package_name is None or any(part is None for part in relative_parts):
        return None

    installed_path = Path("/docker-ros/ws/install") / package_name / "share" / package_name
    for part in relative_parts:
        if part is not None:
            installed_path /= part
    return str(installed_path)


def launch_argument_from_call(call: ast.Call) -> LaunchArgument | None:
    if call_name(call.func) != "DeclareLaunchArgument":
        return None
    if not call.args:
        return None

    name = constant_string(call.args[0])
    if name is None:
        return None

    default_node = call_keyword(call, "default_value")
    default_value = constant_string(default_node)
    description = constant_string(call_keyword(call, "description")) or ""
    if default_value is None:
        default_value = ""

    return LaunchArgument(
        name=name,
        default_value=default_value,
        description=description,
        installed_default_path=installed_share_path(default_node),
    )


def extract_remappable_topic_names(tree: ast.AST) -> list[str]:
    for node in ast.walk(tree):
        if not isinstance(node, ast.Assign):
            continue
        if not any(isinstance(target, ast.Name) and target.id == "remappable_topics" for target in node.targets):
            continue
        if not isinstance(node.value, ast.List):
            raise ValueError("remappable_topics must be a list of DeclareLaunchArgument calls")

        topic_names: list[str] = []
        for element in node.value.elts:
            if not isinstance(element, ast.Call):
                raise ValueError("remappable_topics entries must be DeclareLaunchArgument calls")
            launch_argument = launch_argument_from_call(element)
            if launch_argument is None:
                raise ValueError("remappable_topics entries must be parseable DeclareLaunchArgument calls")
            topic_names.append(launch_argument.name)
        return topic_names

    return []


def extract_launch_arguments(tree: ast.AST) -> list[LaunchArgument]:
    arguments: list[LaunchArgument] = []
    seen: set[str] = set()
    for node in ast.walk(tree):
        if not isinstance(node, ast.Call):
            continue
        launch_argument = launch_argument_from_call(node)
        if launch_argument is None or launch_argument.name in seen:
            continue
        arguments.append(launch_argument)
        seen.add(launch_argument.name)

    if not arguments:
        raise ValueError("default launch file does not define any DeclareLaunchArgument calls")
    return arguments


def extract_node_packages(tree: ast.AST) -> frozenset[str]:
    nodes = [node for node in ast.walk(tree) if isinstance(node, ast.Call) and call_name(node.func) == "Node"]
    if not nodes:
        raise ValueError("default launch file must define at least one Node call")

    packages: set[str] = set()
    for node in nodes:
        package = constant_string(call_keyword(node, "package"))
        if package is None:
            raise ValueError("Node package must be a string literal")
        packages.add(package)

    return frozenset(packages)


def parse_launch_file(launch_file: Path) -> LaunchData:
    tree = ast.parse(launch_file.read_text(encoding="utf-8"), filename=str(launch_file))
    return LaunchData(
        node_packages=extract_node_packages(tree),
        launch_file_name=launch_file.name,
        arguments=extract_launch_arguments(tree),
        remappable_topic_names=extract_remappable_topic_names(tree),
    )


def candidate_launch_files(launch_dir: Path, package_name: str) -> list[Path]:
    preferred = launch_dir / f"{package_name}_launch.py"
    candidates = sorted({*launch_dir.glob("*_launch.py"), *launch_dir.glob("*.launch.py")})
    if preferred in candidates:
        return [preferred, *[candidate for candidate in candidates if candidate != preferred]]
    return candidates


def find_default_launch_file(repo_root: Path, package_name: str) -> Path:
    launch_dir = repo_root / package_name / "launch"
    if not launch_dir.is_dir():
        raise FileNotFoundError(f"default launch directory not found: {launch_dir}")

    preferred_launch_file = launch_dir / f"{package_name}_launch.py"
    matching_launch_files: list[Path] = []
    parsed_launch_files: list[tuple[Path, frozenset[str]]] = []
    errors: list[str] = []
    for launch_file in candidate_launch_files(launch_dir, package_name):
        try:
            launch_data = parse_launch_file(launch_file)
        except Exception as exc:
            errors.append(f"{launch_file}: {exc}")
            continue
        parsed_launch_files.append((launch_file, launch_data.node_packages))
        if package_name in launch_data.node_packages:
            matching_launch_files.append(launch_file)

    if len(matching_launch_files) == 1:
        return matching_launch_files[0]
    if len(matching_launch_files) > 1:
        if preferred_launch_file in matching_launch_files:
            return preferred_launch_file
        launch_files = ", ".join(str(path) for path in matching_launch_files)
        raise ValueError(f"multiple default launch files found for package {package_name!r}: {launch_files}")
    if len(parsed_launch_files) == 1:
        return parsed_launch_files[0][0]

    details = "\n".join(errors)
    message = f"default launch file for package {package_name!r} not found in {launch_dir}"
    if parsed_launch_files:
        launch_files = ", ".join(
            f"{path} (launches packages: {', '.join(sorted(node_packages))})"
            for path, node_packages in parsed_launch_files
        )
        message = f"{message}\nparseable launch files: {launch_files}"
    if details:
        message = f"{message}\n{details}"
    raise FileNotFoundError(message)


def find_compose_launch_files(repo_root: Path, package_name: str) -> list[Path]:
    """Find launch files to deploy while preserving the package-default convention."""
    launch_dir = repo_root / package_name / "launch"
    if not launch_dir.is_dir():
        raise FileNotFoundError(f"default launch directory not found: {launch_dir}")

    preferred_launch_file = launch_dir / f"{package_name}_launch.py"
    matching_launch_files: list[Path] = []
    parsed_launch_files: list[tuple[Path, frozenset[str]]] = []
    errors: list[str] = []
    for launch_file in candidate_launch_files(launch_dir, package_name):
        try:
            launch_data = parse_launch_file(launch_file)
        except Exception as exc:
            errors.append(f"{launch_file}: {exc}")
            continue
        parsed_launch_files.append((launch_file, launch_data.node_packages))
        if package_name in launch_data.node_packages:
            matching_launch_files.append(launch_file)

    if preferred_launch_file in matching_launch_files:
        return [preferred_launch_file]
    if matching_launch_files:
        return matching_launch_files
    if len(parsed_launch_files) == 1:
        return [parsed_launch_files[0][0]]

    details = "\n".join(errors)
    message = f"default launch file for package {package_name!r} not found in {launch_dir}"
    if parsed_launch_files:
        launch_files = ", ".join(
            f"{path} (launches packages: {', '.join(sorted(node_packages))})"
            for path, node_packages in parsed_launch_files
        )
        message = f"{message}\nparseable launch files: {launch_files}"
    if details:
        message = f"{message}\n{details}"
    raise FileNotFoundError(message)


def launch_file_stem(launch_file: Path) -> str:
    for suffix in ("_launch.py", ".launch.py"):
        if launch_file.name.endswith(suffix):
            stem = launch_file.name[: -len(suffix)]
            if stem:
                return stem
    raise ValueError(f"unsupported launch file name: {launch_file.name}")


def compose_paths_for_launch_files(launch_files: list[Path]) -> dict[Path, Path]:
    if len(launch_files) == 1:
        return {launch_files[0]: COMPOSE_PATH}

    compose_paths: dict[Path, Path] = {}
    used_stems: dict[str, Path] = {}
    for launch_file in launch_files:
        stem = launch_file_stem(launch_file)
        previous = used_stems.get(stem)
        if previous is not None:
            raise ValueError(
                f"launch file name collision after removing the launch suffix: {previous.name}, {launch_file.name}"
            )
        used_stems[stem] = launch_file
        compose_paths[launch_file] = COMPOSE_DIR / f"docker-compose.{stem}.yml"
    return compose_paths


def parse_package_metadata(package_xml: Path) -> PackageMetadata:
    root = ET.parse(package_xml).getroot()
    name = " ".join(root.findtext("name", default="").split())
    version = " ".join(root.findtext("version", default="").split())
    if not name or not version:
        raise ValueError(f"package.xml must define name and version: {package_xml}")
    return PackageMetadata(name=name, version=version)


def discover_package_metadata(repo_root: Path) -> list[PackageMetadata]:
    metadata: list[PackageMetadata] = []
    for package_xml in sorted(repo_root.rglob("package.xml")):
        if package_xml.parent == repo_root:
            continue
        metadata.append(parse_package_metadata(package_xml))
    return metadata


def package_metadata_for_subdirectory(repo_root: Path, package_name: str) -> PackageMetadata | None:
    package_xml = repo_root / package_name / "package.xml"
    if not package_xml.is_file():
        return None

    package_metadata = parse_package_metadata(package_xml)
    if package_metadata.name != package_name:
        raise ValueError(
            f"default package.xml name {package_metadata.name!r} does not match package name {package_name!r}"
        )
    return package_metadata


def repository_name_from_origin(repo_root: Path) -> str | None:
    result = subprocess.run(
        ["git", "remote", "get-url", "origin"],
        cwd=repo_root,
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    if result.returncode != 0:
        return None

    remote_url = result.stdout.strip().removesuffix(".git").rstrip("/")
    if not remote_url:
        return None

    return remote_url[max(remote_url.rfind("/"), remote_url.rfind(":")) + 1 :]


def find_default_package_metadata(repo_root: Path) -> PackageMetadata:
    package_metadata = package_metadata_for_subdirectory(repo_root, repo_root.name)
    if package_metadata is not None:
        return package_metadata

    origin_repository_name = repository_name_from_origin(repo_root)
    if origin_repository_name and origin_repository_name != repo_root.name:
        package_metadata = package_metadata_for_subdirectory(repo_root, origin_repository_name)
        if package_metadata is not None:
            return package_metadata

    packages = discover_package_metadata(repo_root)
    if len(packages) == 1:
        return packages[0]
    if not packages:
        raise FileNotFoundError(f"no ROS package.xml found below repository root: {repo_root}")
    package_names = ", ".join(package.name for package in packages)
    raise ValueError(
        f"repository root name {repo_root.name!r} does not identify a package and multiple packages were found: "
        f"{package_names}"
    )


def run_git(args: list[str], repo_root: Path) -> str:
    result = subprocess.run(
        ["git", *args],
        cwd=repo_root,
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    if result.returncode != 0:
        raise RuntimeError(result.stderr.strip() or f"git {' '.join(args)} failed")
    return result.stdout.strip()


def parse_repo_remote(remote_url: str) -> RepoRemote:
    remote_url = remote_url.strip()
    url_match = re.match(
        r"^[a-z][a-z0-9+.-]*://(?:(?:[^@/]+)@)?([^/]+)/(.+?)(?:\.git)?/?$",
        remote_url,
        re.IGNORECASE,
    )
    ssh_match = re.match(r"^(?:[^@]+@)?([^:]+):(.+?)(?:\.git)?$", remote_url)

    if url_match:
        host, path = url_match.group(1), url_match.group(2)
    elif ssh_match:
        host, path = ssh_match.group(1), ssh_match.group(2)
    else:
        raise ValueError(f"origin remote is not a supported GitHub or GitLab URL: {remote_url}")

    port: str | None = None
    if ":" in host and not host.startswith("["):
        host, port = host.rsplit(":", 1)

    path_parts = [part for part in path.split("/") if part]
    if len(path_parts) < 2:
        raise ValueError(f"origin remote is not a supported GitHub or GitLab URL: {remote_url}")

    return RepoRemote(host=host, port=port, path="/".join(path_parts))


def gitlab_registry_host(remote: RepoRemote, registry: str | None) -> str:
    if registry:
        return registry.rstrip("/")
    return f"{remote.host}:5050"


def image_repository_from_reference(image_reference: str) -> str:
    image_reference = image_reference.split("@", 1)[0]
    last_slash_index = image_reference.rfind("/")
    last_colon_index = image_reference.rfind(":")
    if last_colon_index > last_slash_index:
        return image_reference[:last_colon_index]
    return image_reference


def container_image_repository_from_existing_compose(repo_root: Path, package_name: str) -> str | None:
    compose_dir = repo_root / COMPOSE_DIR
    compose_paths = [repo_root / COMPOSE_PATH]
    if compose_dir.is_dir():
        compose_paths.extend(sorted(compose_dir.glob("docker-compose.*.yml")))

    for compose_path in dict.fromkeys(compose_paths):
        if not compose_path.is_file():
            continue
        for line in compose_path.read_text(encoding="utf-8").splitlines():
            match = re.match(r"\s*image:\s*(\S+)", line)
            if not match:
                continue
            image_repository = image_repository_from_reference(match.group(1))
            if image_repository.split("/")[-1] == package_name:
                return image_repository
    return None


def container_image_repository(repo_root: Path, package_name: str, gitlab_registry: str | None) -> str:
    remote_url = ""
    try:
        remote_url = run_git(["remote", "get-url", "origin"], repo_root)
        remote = parse_repo_remote(remote_url)
    except ValueError:
        image_repository = container_image_repository_from_existing_compose(repo_root, package_name)
        if image_repository:
            return image_repository
        raise

    host_lower = remote.host.lower()
    namespace = "/".join(remote.path.split("/")[:-1])

    if host_lower == "github.com":
        return f"ghcr.io/{namespace}/{package_name}"
    if host_lower.startswith("gitlab."):
        if gitlab_registry is None:
            image_repository = container_image_repository_from_existing_compose(repo_root, package_name)
            if image_repository:
                return image_repository
        return f"{gitlab_registry_host(remote, gitlab_registry)}/{remote.path}"
    image_repository = container_image_repository_from_existing_compose(repo_root, package_name)
    if image_repository:
        return image_repository
    raise ValueError(f"origin remote is not a supported GitHub or GitLab URL: {remote_url}")


def env_name(argument_name: str) -> str:
    return argument_name.upper()


def compose_service_name(package_name: str) -> str:
    return package_name.replace("_", "-")


def sorted_launch_arguments(launch_data: LaunchData) -> dict[str, LaunchArgument]:
    return {argument.name: argument for argument in launch_data.arguments}


def command_argument_names(launch_data: LaunchData) -> list[str]:
    handled_names = {*STANDARD_LAUNCH_ARGUMENT_NAMES, *launch_data.remappable_topic_names}
    extra_names = [argument.name for argument in launch_data.arguments if argument.name not in handled_names]
    names = [*STANDARD_LAUNCH_ARGUMENT_NAMES, *extra_names, *launch_data.remappable_topic_names]
    return [name for name in names if name in sorted_launch_arguments(launch_data)]


def _environment_variable(name: str, value: str) -> EnvironmentVariable:
    return EnvironmentVariable(
        name=name,
        value=value,
        host_exposed=name in HOST_EXPOSED_ENV_NAMES,
    )


def extra_launch_environment_variables(launch_data: LaunchData) -> list[EnvironmentVariable]:
    arguments = sorted_launch_arguments(launch_data)
    handled_names = {*STANDARD_LAUNCH_ARGUMENT_NAMES, *launch_data.remappable_topic_names}
    return [
        _environment_variable(
            name=env_name(argument.name),
            value=argument.default_value,
        )
        for argument in arguments.values()
        if argument.name not in handled_names
    ]


def _compose_environment_variables(launch_data: LaunchData, params_default_path: str | None) -> list[EnvironmentVariable]:
    arguments = sorted_launch_arguments(launch_data)
    variables = [
        _environment_variable(
            name="LOG_LEVEL",
            value=arguments.get("log_level", LaunchArgument("log_level", "info", "")).default_value or "info",
        ),
        _environment_variable(
            name="USE_SIM_TIME",
            value=arguments.get("use_sim_time", LaunchArgument("use_sim_time", "false", "")).default_value or "false",
        ),
    ]
    if params_default_path is not None:
        variables.append(_environment_variable(name="PARAMS", value=params_default_path))
    variables.extend(extra_launch_environment_variables(launch_data))
    return variables


def strip_markdown_code(value: str) -> str:
    value = value.strip()
    if value.startswith("`") and value.endswith("`") and len(value) >= 2:
        return value[1:-1]
    return value


def extract_readme_topic_directions(package_readme: Path) -> dict[str, str]:
    if not package_readme.is_file():
        return {}

    directions: dict[str, str] = {}
    current_direction: str | None = None
    for line in package_readme.read_text(encoding="utf-8").splitlines():
        stripped = line.strip()
        if stripped == "#### Subscribed Topics":
            current_direction = "input"
            continue
        if stripped == "#### Published Topics":
            current_direction = "output"
            continue
        if stripped.startswith("#### "):
            current_direction = None
            continue
        if current_direction is None or not stripped.startswith("|"):
            continue

        cells = [cell.strip() for cell in stripped.strip("|").split("|")]
        if len(cells) < 3 or cells[0] in {"Topic", "---"}:
            continue
        topic = strip_markdown_code(cells[0])
        existing_direction = directions.get(topic)
        if existing_direction is not None and existing_direction != current_direction:
            directions[topic] = "other"
            continue
        directions[topic] = current_direction

    return directions


def topic_environment_variables(
    launch_data: LaunchData, package_readme: Path
) -> tuple[list[EnvironmentVariable], list[EnvironmentVariable], list[EnvironmentVariable]]:
    arguments = sorted_launch_arguments(launch_data)
    topic_directions = extract_readme_topic_directions(package_readme)
    input_variables: list[EnvironmentVariable] = []
    output_variables: list[EnvironmentVariable] = []
    other_variables: list[EnvironmentVariable] = []

    for name in launch_data.remappable_topic_names:
        argument = arguments.get(name)
        if argument is None:
            continue
        direction = topic_directions.get(argument.default_value, "other")
        variable = EnvironmentVariable(name=env_name(argument.name), value=argument.default_value)

        if direction == "input":
            input_variables.append(variable)
        elif direction == "output":
            output_variables.append(variable)
        else:
            other_variables.append(variable)

    return input_variables, output_variables, other_variables


def build_template_environment() -> Environment:
    if Environment is None or FileSystemLoader is None:
        raise RuntimeError(
            "Missing dependency: jinja2. Install with "
            "`pip install -r .openads-dev-environment/scripts/requirements.txt`."
        )
    templates_dir = Path(__file__).resolve().parent / "templates"
    return Environment(
        loader=FileSystemLoader(str(templates_dir)),
        autoescape=False,
        keep_trailing_newline=True,
        trim_blocks=True,
        lstrip_blocks=True,
    )


def render_template(template_env: Environment, template_name: str, context: dict) -> str:
    try:
        template = template_env.get_template(template_name)
    except TemplateNotFound as exc:
        raise RuntimeError(f"Missing template: {template_name}") from exc
    return template.render(**context)


def installed_params_path(package_name: str) -> str:
    return f"/docker-ros/ws/install/{package_name}/share/{package_name}/config/params.yml"


def render_compose(
    repo_root: Path,
    package_metadata: PackageMetadata,
    launch_data: LaunchData,
    image_repository: str,
    *,
    multi_launch: bool,
) -> str:
    arguments = sorted_launch_arguments(launch_data)
    input_variables, output_variables, other_topic_variables = topic_environment_variables(
        launch_data, repo_root / package_metadata.name / "README.md"
    )
    if multi_launch:
        node_name = arguments.get("name", LaunchArgument("name", package_metadata.name, "")).default_value
        node_name = node_name or package_metadata.name
        params_argument = arguments.get("params")
        params_default_path = None
        if params_argument is not None:
            params_default_path = params_argument.installed_default_path or params_argument.default_value or None
    else:
        node_name = package_metadata.name
        params_default_path = installed_params_path(package_metadata.name) if "params" in arguments else None

    launch_argument_names = command_argument_names(launch_data)
    if multi_launch and "params" in launch_argument_names and params_default_path is None:
        launch_argument_names.remove("params")

    context = {
        "service_name": compose_service_name(package_metadata.name),
        "image": f"{image_repository}:v{package_metadata.version}",
        "namespace": DEFAULT_NAMESPACE,
        "node_name": node_name,
        "input_variables": input_variables,
        "output_variables": output_variables,
        "other_topic_variables": other_topic_variables,
        "other_environment_variables": _compose_environment_variables(launch_data, params_default_path),
        "params_default_path": params_default_path,
        "launch_package": package_metadata.name,
        "launch_file_name": launch_data.launch_file_name,
        "launch_arguments": [
            LaunchCommandArgument(
                name=argument_name,
                env_name=env_name(argument_name),
                only_if_set=arguments[argument_name].default_value == ""
                and argument_name not in STANDARD_LAUNCH_ARGUMENT_NAMES,
            )
            for argument_name in launch_argument_names
        ],
        "package_name": package_metadata.name,
    }
    return render_template(build_template_environment(), "docker_compose.yml.j2", context)


def build_compose(repo_root: Path, gitlab_registry: str | None = None) -> str:
    """Build the legacy single Compose output used by the Helm generator and callers."""
    package_metadata = find_default_package_metadata(repo_root)
    launch_data = parse_launch_file(find_default_launch_file(repo_root, package_metadata.name))
    image_repository = container_image_repository(repo_root, package_metadata.name, gitlab_registry)
    return render_compose(
        repo_root,
        package_metadata,
        launch_data,
        image_repository,
        multi_launch=False,
    )


def build_composes(repo_root: Path, gitlab_registry: str | None = None) -> dict[Path, str]:
    package_metadata = find_default_package_metadata(repo_root)
    launch_files = find_compose_launch_files(repo_root, package_metadata.name)
    compose_paths = compose_paths_for_launch_files(launch_files)
    multi_launch = len(launch_files) > 1
    image_repository = container_image_repository(repo_root, package_metadata.name, gitlab_registry)

    return {
        compose_paths[launch_file]: render_compose(
            repo_root,
            package_metadata,
            parse_launch_file(launch_file),
            image_repository,
            multi_launch=multi_launch,
        )
        for launch_file in launch_files
    }


def build_diff(expected: str, current: str, compose_path: Path, repo_root: Path) -> str:
    rel_path = compose_path.relative_to(repo_root).as_posix()
    return "".join(
        difflib.unified_diff(
            current.splitlines(keepends=True),
            expected.splitlines(keepends=True),
            fromfile=f"a/{rel_path}",
            tofile=f"b/{rel_path}",
        )
    )


def print_diff(old: str, new: str, compose_path: Path, repo_root: Path) -> None:
    red, green, cyan, reset = "\033[31m", "\033[32m", "\033[36m", "\033[0m"
    diff = build_diff(new, old, compose_path, repo_root)
    if not diff:
        print("(no changes)")
        return

    for line in diff.splitlines(keepends=True):
        if line.startswith("+") and not line.startswith("+++"):
            color = green
        elif line.startswith("-") and not line.startswith("---"):
            color = red
        elif line.startswith("@@"):
            color = cyan
        else:
            color = ""
        print(f"{color}{line}{reset if color else ''}", end="")


def resolve_repo_root(raw_repo_root: str | None) -> Path:
    if raw_repo_root:
        return Path(raw_repo_root).resolve()
    return Path(__file__).resolve().parents[2]


def managed_compose_paths(repo_root: Path) -> set[Path]:
    compose_dir = repo_root / COMPOSE_DIR
    paths = {repo_root / COMPOSE_PATH} if (repo_root / COMPOSE_PATH).is_file() else set()
    if compose_dir.is_dir():
        paths.update(path for path in compose_dir.glob("docker-compose.*.yml") if path.is_file())
    return paths


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Generate Docker Compose files from ROS 2 launch files")
    parser.add_argument("repo_root", nargs="?", help="Repository root (defaults to inferred top-level)")
    parser.add_argument("--check", action="store_true", help="Check whether docker compose output is up to date")
    parser.add_argument(
        "--gitlab-registry",
        default=os.environ.get(GITLAB_REGISTRY_ENV_NAME),
        help=(
            "GitLab container registry host, optionally including a port "
            f"(defaults to ${GITLAB_REGISTRY_ENV_NAME}, an existing Compose file, or <gitlab-host>:5050)"
        ),
    )
    return parser


def main() -> int:
    parser = build_arg_parser()
    args = parser.parse_args()
    repo_root = resolve_repo_root(args.repo_root)

    try:
        expected_files = build_composes(repo_root, gitlab_registry=args.gitlab_registry)
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2

    expected_paths = {repo_root / relative_path for relative_path in expected_files}
    obsolete_paths = managed_compose_paths(repo_root) - expected_paths

    if args.check:
        stale = False
        for relative_path, expected in sorted(expected_files.items()):
            compose_path = repo_root / relative_path
            current = compose_path.read_text(encoding="utf-8") if compose_path.exists() else ""
            if current == expected:
                continue
            stale = True
            diff = build_diff(expected, current, compose_path, repo_root)
            if diff:
                print(diff, end="" if diff.endswith("\n") else "\n")
            else:
                print(f"{relative_path} is stale", file=sys.stderr)
        for compose_path in sorted(obsolete_paths):
            stale = True
            current = compose_path.read_text(encoding="utf-8")
            diff = build_diff("", current, compose_path, repo_root)
            if diff:
                print(diff, end="" if diff.endswith("\n") else "\n")
            else:
                print(f"{compose_path.relative_to(repo_root)} is obsolete", file=sys.stderr)
        return 1 if stale else 0

    for relative_path, expected in sorted(expected_files.items()):
        compose_path = repo_root / relative_path
        compose_path.parent.mkdir(parents=True, exist_ok=True)
        current = compose_path.read_text(encoding="utf-8") if compose_path.exists() else ""
        compose_path.write_text(expected, encoding="utf-8")
        print(f"Updated {compose_path}")
        print_diff(current, expected, compose_path, repo_root)
    for compose_path in sorted(obsolete_paths):
        current = compose_path.read_text(encoding="utf-8")
        compose_path.unlink()
        print(f"Removed obsolete {compose_path}")
        print_diff(current, "", compose_path, repo_root)
    return 0


if __name__ == "__main__":
    sys.exit(main())