#!/usr/bin/env python3
"""Enforce the workspace dependency-management conventions.

The rules mirror the comment above `[workspace.dependencies]` in the root
`Cargo.toml`. This check fails (exit 1) if either is violated:

  1. Centralize shared deps: a registry dependency declared *inline* (i.e. not
     via `workspace = true`) in two or more workspace members must instead live
     in `[workspace.dependencies]` and be inherited.
  2. No shadowing: a member must not re-declare a `[workspace.dependencies]`
     dependency inline; it must inherit it with `name.workspace = true`
     (features / `optional` may still be added per-crate).

Exemptions:
  - Protocol-compiler-generated manifests (`Code generated by ...` header).
  - `sample_connector_scaffolding` (kept standalone so release tooling only has
    to strip `path` keys; it deliberately pins its own versions).
  - Local `path` dependencies and sibling `azure_iot_operations_*` crates, which
    are declared per-crate by design.

The rule is workspace-global, so all members are checked regardless of the
current directory.
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import sys

# Matches `[dependencies]`, `[dev-dependencies]`, `[build-dependencies]`, and
# their target-specific forms, e.g. `[target.'cfg(unix)'.dependencies]`.
SECTION_RE = re.compile(
    r"^\[(?:.*\.)?(?:dependencies|dev-dependencies|build-dependencies)\]$"
)
KEY_RE = re.compile(r"^\s*([A-Za-z0-9_.-]+)\s*=(.*)$")
WORKSPACE_TRUE_RE = re.compile(r"workspace\s*=\s*true")
PATH_KEY_RE = re.compile(r"\bpath\s*=")
SIBLING_RE = re.compile(r"^azure_iot_operations_")

GENERATED_MARKER = "Code generated by Azure.Iot.Operations.ProtocolCompiler"
SCAFFOLDING_SUFFIX = (
    "sample_applications/sample_connector_scaffolding/Cargo.toml"
)


def load_metadata() -> dict:
    out = subprocess.run(
        ["cargo", "metadata", "--format-version=1", "--no-deps"],
        check=True,
        capture_output=True,
        text=True,
    ).stdout
    return json.loads(out)


def parse_workspace_dep_keys(root_manifest: str) -> set[str]:
    keys: set[str] = set()
    in_section = False
    with open(root_manifest, encoding="utf-8") as handle:
        for raw in handle:
            stripped = raw.strip()
            if stripped == "[workspace.dependencies]":
                in_section = True
                continue
            if in_section and stripped.startswith("["):
                break
            if not in_section or not stripped or stripped.startswith("#"):
                continue
            match = KEY_RE.match(raw)
            if match:
                keys.add(match.group(1))
    return keys


def is_generated(manifest: str) -> bool:
    try:
        with open(manifest, encoding="utf-8") as handle:
            return GENERATED_MARKER in handle.readline()
    except OSError:
        return False


def iter_dependencies(manifest: str):
    """Yield (name, inherited, has_path) for each dependency table entry.

    Handles multi-line inline tables by tracking brace/bracket depth so that
    continuation lines are not mistaken for new entries.
    """
    in_deps = False
    pending_name: str | None = None
    buf = ""
    depth = 0

    def classify(name: str, value: str):
        inherited = name.endswith(".workspace") or bool(
            WORKSPACE_TRUE_RE.search(value)
        )
        base = name[: -len(".workspace")] if name.endswith(".workspace") else name
        has_path = bool(PATH_KEY_RE.search(value))
        return base, inherited, has_path

    with open(manifest, encoding="utf-8") as handle:
        for raw in handle:
            line = raw.rstrip("\n")
            if pending_name is not None:
                buf += "\n" + line
                depth += line.count("{") + line.count("[")
                depth -= line.count("}") + line.count("]")
                if depth <= 0:
                    yield classify(pending_name, buf)
                    pending_name = None
                    buf = ""
                continue

            stripped = line.strip()
            if SECTION_RE.match(stripped):
                in_deps = True
                continue
            if stripped.startswith("["):
                in_deps = False
                continue
            if not in_deps or not stripped or stripped.startswith("#"):
                continue

            match = KEY_RE.match(line)
            if not match:
                continue
            name, value = match.group(1), match.group(2)
            depth = value.count("{") + value.count("[")
            depth -= value.count("}") + value.count("]")
            if depth > 0:
                pending_name, buf = name, value
                continue
            yield classify(name, value)


def main() -> int:
    metadata = load_metadata()
    root = metadata["workspace_root"]
    ws_keys = parse_workspace_dep_keys(os.path.join(root, "Cargo.toml"))

    member_ids = set(metadata["workspace_members"])
    manifests = [
        pkg["manifest_path"]
        for pkg in metadata["packages"]
        if pkg["id"] in member_ids
    ]

    inline_registry: dict[str, set[str]] = {}
    shadows: set[tuple[str, str]] = set()

    for manifest in manifests:
        normalized = manifest.replace(os.sep, "/")
        if normalized.endswith(SCAFFOLDING_SUFFIX) or is_generated(manifest):
            continue
        for name, inherited, has_path in iter_dependencies(manifest):
            if name in ws_keys and not inherited:
                shadows.add((manifest, name))
            if inherited or has_path:
                continue
            if SIBLING_RE.match(name) or name == "envoy":
                continue
            inline_registry.setdefault(name, set()).add(manifest)

    errors: list[str] = []
    for name, mans in sorted(inline_registry.items()):
        if len(mans) >= 2 and name not in ws_keys:
            crates = ", ".join(
                sorted(os.path.relpath(m, root) for m in mans)
            )
            errors.append(
                f"'{name}' is declared inline in {len(mans)} workspace crates "
                f"({crates}); add it to [workspace.dependencies] and inherit "
                f"with `{name}.workspace = true`."
            )
    for manifest, name in sorted(shadows):
        rel = os.path.relpath(manifest, root)
        errors.append(
            f"{rel}: '{name}' is in [workspace.dependencies] but declared "
            f"inline; inherit it with `{name}.workspace = true` "
            f"(features/optional may be added per-crate)."
        )

    if errors:
        print(
            "ERROR: workspace dependency policy violations "
            "(see [workspace.dependencies] conventions in rust/Cargo.toml):",
            file=sys.stderr,
        )
        for err in errors:
            print(f"  - {err}", file=sys.stderr)
        return 1
    return 0


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