#!/usr/bin/env python3

from __future__ import annotations

import re
import sys
import tomllib
from pathlib import Path
from typing import Any


def extract_workspace_source_roots(flake_path: Path) -> dict[str, list[str]]:
    source = flake_path.read_text()
    match = re.search(r"workspaceSourceRoots\s*=\s*\{(.*?)\n\s*\};", source, re.S)
    if match is None:
        raise ValueError(f"Could not find workspaceSourceRoots in {flake_path}")

    roots: dict[str, list[str]] = {}
    for name, body in re.findall(r"\n\s*(\w+)\s*=\s*\[(.*?)\];", match.group(1), re.S):
        roots[name] = re.findall(r'"([^"]+)"', body)
    return roots


def collect_path_dependencies(value: Any) -> list[str]:
    found: list[str] = []

    if isinstance(value, dict):
        path = value.get("path")
        if isinstance(path, str):
            found.append(path)
        for nested in value.values():
            found.extend(collect_path_dependencies(nested))
    elif isinstance(value, list):
        for nested in value:
            found.extend(collect_path_dependencies(nested))

    return found


def workspace_manifests(repo_root: Path, workspace_name: str) -> list[Path]:
    workspace_manifest = repo_root / workspace_name / "Cargo.toml"
    manifests = [workspace_manifest]
    workspace_data = tomllib.loads(workspace_manifest.read_text())
    members = workspace_data.get("workspace", {}).get("members", [])

    for member in members:
        for candidate in (workspace_manifest.parent).glob(member):
            manifest = candidate if candidate.name == "Cargo.toml" else candidate / "Cargo.toml"
            if manifest.is_file():
                manifests.append(manifest)

    unique_manifests: list[Path] = []
    seen: set[Path] = set()
    for manifest in manifests:
        resolved = manifest.resolve()
        if resolved in seen:
            continue
        seen.add(resolved)
        unique_manifests.append(manifest)
    return unique_manifests


def required_root(dep_rel: Path) -> str:
    parts = dep_rel.parts
    if not parts:
        return ""
    if parts[0] == "crates" and len(parts) >= 2:
        return "/".join(parts[:2])
    return parts[0]


def is_covered(dep_rel: str, configured_roots: list[str]) -> bool:
    return any(dep_rel == root or dep_rel.startswith(f"{root}/") for root in configured_roots)


def main() -> int:
    repo_root = Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else Path.cwd().resolve()
    flake_path = repo_root / "flake.nix"
    workspace_roots = extract_workspace_source_roots(flake_path)

    failures: list[str] = []

    for workspace_name, configured_roots in sorted(workspace_roots.items()):
        workspace_manifest = repo_root / workspace_name / "Cargo.toml"
        if not workspace_manifest.is_file():
            continue

        for manifest in workspace_manifests(repo_root, workspace_name):
            manifest_data = tomllib.loads(manifest.read_text())
            for dep_path in collect_path_dependencies(manifest_data):
                dependency_dir = (manifest.parent / dep_path).resolve()
                try:
                    dep_rel = dependency_dir.relative_to(repo_root)
                except ValueError:
                    continue

                dep_rel_str = dep_rel.as_posix()
                if is_covered(dep_rel_str, configured_roots):
                    continue

                needed = required_root(dep_rel)
                manifest_rel = manifest.relative_to(repo_root).as_posix()
                failures.append(
                    f"{workspace_name}: missing source root '{needed}' for dependency "
                    f"'{dep_rel_str}' referenced by {manifest_rel}"
                )

    if failures:
        print("workspaceSourceRoots is missing path dependencies:", file=sys.stderr)
        for failure in failures:
            print(f"  - {failure}", file=sys.stderr)
        return 1

    print("workspaceSourceRoots covers all workspace path dependencies.")
    return 0


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