"""Build, lint, test and launch a generated ROS 2 package, the way a developer would on their machine.

Runs inside the ROS 2 image (worker/ros/Dockerfile) as a sandbox: the Modal function has no network and no
secrets, and every ROS command gets a scrubbed environment. Standard library only, so it runs with the image's
own Python as well as Modal's.

    python3 -m worker.ros_runner PACKAGE_DIR_OR_JSON --distro jazzy [--out OUT_DIR]

Steps (each one's output is kept as a log file):
    1. colcon build                      build.log
    2. ament linters on the sources      lint.log   (flake8/pep257 for Python, cpplint/cppcheck for C++, xmllint)
    3. colcon test, if it has tests      test.log   (launch_testing, pytest, gtest, ament_lint_auto)
    4. launch it for a few seconds       launch.log (ros2 launch, or ros2 run for each executable without one)
       while ros2 bag records every topic to MCAP  bag.mcap
"""

from __future__ import annotations

import argparse
import json
import os
import random
import re
import shutil
import signal
import subprocess
import tempfile
import time
import xml.etree.ElementTree as ET
from pathlib import Path, PurePosixPath

DISTROS = ("humble", "jazzy")
BUILD_TIMEOUT = 600
TEST_TIMEOUT = 300
LINT_TIMEOUT = 120
LAUNCH_SECONDS = 10
MAX_BAG_BYTES = 20_000_000
MAX_LOG_BYTES = 400_000
# DDS discovery over unicast on loopback only. Containers (and Modal's no-network sandbox) have no multicast on
# `lo`, which the default localhost-only discovery needs, so nodes would never see each other.
CYCLONE_CONFIG = """<CycloneDDS xmlns="https://cdds.io/config">
  <Domain Id="any">
    <General>
      <Interfaces><NetworkInterface name="lo"/></Interfaces>
      <AllowMulticast>false</AllowMulticast>
    </General>
    <Discovery>
      <ParticipantIndex>auto</ParticipantIndex>
      <MaxAutoParticipantIndex>60</MaxAutoParticipantIndex>
      <Peers><Peer address="127.0.0.1"/></Peers>
    </Discovery>
  </Domain>
</CycloneDDS>
"""
LAUNCH_SUFFIXES = (".launch.py", "_launch.py", ".launch.xml", "_launch.xml", ".launch.yaml", "_launch.yaml")


def _env(distro: str, home: Path) -> dict[str, str]:
    """A clean environment: nothing from the host (no credentials), ROS traffic kept on localhost."""
    (home / "cyclonedds.xml").write_text(CYCLONE_CONFIG)
    env = {
        "PATH": "/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin",
        "HOME": str(home),
        "LANG": "C.UTF-8",
        "LC_ALL": "C.UTF-8",
        "ROS_DISTRO": distro,
        "ROS_DOMAIN_ID": str(random.randint(1, 100)),
        "RMW_IMPLEMENTATION": "rmw_cyclonedds_cpp",
        "CYCLONEDDS_URI": f"file://{home / 'cyclonedds.xml'}",
        "PYTHONUNBUFFERED": "1",
        "RCUTILS_COLORIZED_OUTPUT": "0",
        "QT_QPA_PLATFORM": "offscreen",
        "MAKEFLAGS": f"-j{os.cpu_count() or 2}",
    }
    return env


def _bash(script: str, ws: Path, env: dict, timeout: float) -> tuple[int | None, str]:
    """Run `script` in bash with ROS (and the workspace, once built) sourced. Returns (exit code or None, output)."""
    distro = env["ROS_DISTRO"]
    prelude = f"source /opt/ros/{distro}/setup.bash; [ -f install/setup.bash ] && source install/setup.bash; "
    try:
        p = subprocess.run(
            ["bash", "-c", prelude + script],
            cwd=ws,
            env=env,
            capture_output=True,
            text=True,
            timeout=timeout,
            errors="replace",
        )
        return p.returncode, p.stdout + p.stderr
    except subprocess.TimeoutExpired as e:
        out = (e.stdout or b"").decode(errors="replace") if isinstance(e.stdout, bytes) else (e.stdout or "")
        return None, out + f"\n[rosdiff] timed out after {timeout:.0f}s\n"


def write_workspace(files: dict[str, str], ws: Path) -> list[str]:
    """Write the generated files under ws/src. Returns the package names found."""
    src = ws / "src"
    src.mkdir(parents=True, exist_ok=True)
    for rel, content in files.items():
        p = PurePosixPath(rel)
        if p.is_absolute() or ".." in p.parts:
            raise ValueError(f"unsafe path in package: {rel}")
        dest = src / p
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_text(content)
    names = []
    for xml in sorted(src.rglob("package.xml")):
        try:
            name = ET.parse(xml).getroot().findtext("name")
        except ET.ParseError:
            name = None
        if name:
            names.append(name.strip())
    return names


# --------------------------------------------------------------------------- steps


_MISSING = [
    re.compile(r'provided by "([\w-]+)"'),  # CMake find_package
    re.compile(r"No module named '([\w.]+)'"),  # Python import
    re.compile(r"Package '([\w-]+)' not found"),  # ros2 launch / ament index
]


def missing_dependencies(text: str) -> list[str]:
    found: list[str] = []
    for pattern in _MISSING:
        for name in pattern.findall(text):
            if name not in found:
                found.append(name)
    return found


def first_errors(text: str, limit: int = 8) -> list[str]:
    """The first compiler/CMake/Python errors, each on one line, for the summary."""
    lines = text.splitlines()
    out: list[str] = []
    used: set[int] = set()
    for i, raw in enumerate(lines):
        ln = raw.strip()
        if not ln or i in used or ln.startswith("Traceback (most recent call last)"):
            continue
        if "CMake Error" in ln:
            following = [(j, x.strip()) for j, x in enumerate(lines[i + 1 : i + 4], i + 1) if x.strip()]
            if following:
                used.add(following[0][0])
                ln = f"{ln} {following[0][1]}"
        elif not (
            "error:" in ln
            or re.match(r"^[\w.]*(Error|Exception)\b.*:", ln)
            or "SyntaxError" in ln
            or "ImportError" in ln
            or "ModuleNotFoundError" in ln
        ):
            continue
        if ln not in out:
            out.append(ln[:300])
        if len(out) >= limit:
            break
    return out


def build(ws: Path, env: dict) -> dict:
    code, log = _bash(
        "colcon build --event-handlers console_direct+ --cmake-args -DBUILD_TESTING=ON -DCMAKE_BUILD_TYPE=Release",
        ws,
        env,
        BUILD_TIMEOUT,
    )
    (ws / "build.log").write_text(log)
    return {
        "ok": code == 0,
        "exit_code": code,
        "errors": first_errors(log) if code != 0 else [],
        "missing_dependencies": missing_dependencies(log) if code != 0 else [],
    }


def lint(ws: Path, env: dict, packages: dict[str, Path]) -> dict:
    """ament linters, as `ament_lint_auto` would run them. Reported, never a failure on their own."""
    results: dict[str, int] = {}
    logs = []
    for name, pkg in packages.items():
        py = any(pkg.rglob("*.py"))
        cpp = any(p.suffix in (".cpp", ".hpp", ".cc", ".h") for p in pkg.rglob("*"))
        tools = ["ament_xmllint"]
        if py:
            tools += ["ament_flake8", "ament_pep257"]
        if cpp:
            tools += ["ament_cpplint", "ament_cppcheck"]
        for tool in tools:
            code, out = _bash(f"{tool} {pkg}", ws, env, LINT_TIMEOUT)
            issues = _count_issues(tool, out) if code else 0
            results[tool] = results.get(tool, 0) + issues
            logs.append(f"===== {name}: {tool} (exit {code})\n{out}")
    (ws / "lint.log").write_text("\n".join(logs))
    return {"issues": results, "total": sum(results.values())}


def _count_issues(tool: str, out: str) -> int:
    m = re.search(r"(?:errors found: |^)(\d+)(?: (?:errors?|problems?))?\s*$", out, re.MULTILINE)
    if m:
        return int(m.group(1))
    return max(1, sum(1 for ln in out.splitlines() if re.match(r"^\S+:\d+", ln)))


def test(ws: Path, env: dict) -> dict:
    code, log = _bash(
        "colcon test --event-handlers console_direct+ --return-code-on-test-failure; "
        "echo '===== colcon test-result'; colcon test-result --all --verbose",
        ws,
        env,
        TEST_TIMEOUT,
    )
    (ws / "test.log").write_text(log)
    m = re.search(r"Summary: (\d+) tests?, (\d+) errors?, (\d+) failures?, (\d+) skipped", log)
    if not m:
        return {"ran": False, "ok": code == 0}
    total, errors, failures, skipped = map(int, m.groups())
    return {
        "ran": total > 0,
        "ok": errors == 0 and failures == 0,
        "tests": total,
        "errors": errors,
        "failures": failures,
        "skipped": skipped,
    }


def _launch_targets(ws: Path, env: dict, packages: dict[str, Path]) -> list[tuple[str, list[str]]]:
    """What to start: each package's launch file, else each of its executables (at most 3 per package)."""
    targets = []
    for name, pkg in packages.items():
        launches = sorted(
            p for p in pkg.rglob("*") if p.is_file() and p.name.endswith(LAUNCH_SUFFIXES) and "test" not in p.parts
        )
        if launches:
            targets.append((f"ros2 launch {name} {launches[0].name}", [name, launches[0].name]))
            continue
        _, out = _bash(f"ros2 pkg executables {name}", ws, env, 60)
        for line in out.splitlines()[:3]:
            parts = line.split()
            if len(parts) == 2 and parts[0] == name:
                targets.append((f"ros2 run {name} {parts[1]}", parts))
    return targets


def _start(cmd: str, ws: Path, env: dict, log: Path) -> subprocess.Popen:
    distro = env["ROS_DISTRO"]
    script = f"source /opt/ros/{distro}/setup.bash; source install/setup.bash; exec {cmd}"
    return subprocess.Popen(
        ["bash", "-c", script],
        cwd=ws,
        env=env,
        stdout=log.open("w"),
        stderr=subprocess.STDOUT,
        start_new_session=True,  # its own process group, so every node it spawns can be stopped
    )


def _stop(proc: subprocess.Popen, grace: float = 8.0) -> int | None:
    if proc.poll() is not None:
        return proc.returncode
    for sig, wait in ((signal.SIGINT, grace), (signal.SIGTERM, 3.0), (signal.SIGKILL, 2.0)):
        try:
            os.killpg(proc.pid, sig)
        except ProcessLookupError:
            break
        try:
            proc.wait(wait)
            break
        except subprocess.TimeoutExpired:
            continue
    return None  # still running when we stopped it: that is what we want


def launch(ws: Path, env: dict, packages: dict[str, Path], seconds: float = LAUNCH_SECONDS) -> dict:
    targets = _launch_targets(ws, env, packages)
    if not targets:
        (ws / "launch.log").write_text("nothing to launch: no launch file and no executables\n")
        return {"ok": False, "started": [], "error": "no launch file and no executables to start"}

    logs_dir = ws / "launch_logs"
    logs_dir.mkdir()
    bag = subprocess.Popen(
        [
            "bash",
            "-c",
            f"source /opt/ros/{env['ROS_DISTRO']}/setup.bash; exec ros2 bag record -a -s mcap -o {ws / 'bag'}",
        ],
        cwd=ws,
        env=env,
        stdout=(logs_dir / "bag.log").open("w"),
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    time.sleep(2.0)  # let the recorder subscribe before the nodes start publishing
    procs = [(cmd, _start(cmd, ws, env, logs_dir / f"{i}.log")) for i, (cmd, _) in enumerate(targets)]
    exited_early: list[tuple[str, int]] = []
    deadline = time.monotonic() + seconds
    graph: dict = {}
    while time.monotonic() < deadline:
        time.sleep(0.5)
        for cmd, p in procs:
            if p.poll() is not None and (cmd, p.returncode) not in exited_early:
                exited_early.append((cmd, p.returncode))
        if not graph and time.monotonic() > deadline - seconds / 2:
            graph = _graph(ws, env)
    if not graph:
        graph = _graph(ws, env)
    # Crashes are judged on what happened before we stop everything (stopping makes launch report exit -2).
    before_stop = "\n".join((logs_dir / f"{i}.log").read_text(errors="replace") for i in range(len(procs)))
    for _, p in procs:
        _stop(p)
    _stop(bag, grace=10.0)

    texts = []
    for i, (cmd, _) in enumerate(procs):
        texts.append(f"===== {cmd}\n" + (logs_dir / f"{i}.log").read_text(errors="replace"))
    texts.append("===== ros2 bag record\n" + (logs_dir / "bag.log").read_text(errors="replace"))
    log = "\n".join(texts)
    (ws / "launch.log").write_text(log)

    died = re.findall(r"\[ERROR\] \[([^\]]+)\]: process has died \[pid \d+, exit code (-?\d+)", before_stop)
    crashed = [f"{name} (exit {code})" for name, code in died]
    crashed += [f"{cmd} (exit {code})" for cmd, code in exited_early if code != 0]
    tracebacks = first_errors(before_stop)
    ok = not crashed and bool(graph.get("nodes"))
    return {
        "ok": ok,
        "started": [cmd for cmd, _ in procs],
        "seconds": seconds,
        "nodes": graph.get("nodes", []),
        "topics": graph.get("topics", []),
        "crashed": crashed,
        "errors": tracebacks,
        "missing_dependencies": missing_dependencies(log),
        "error": None if ok else ("a node crashed" if crashed else "no ROS node came up"),
    }


def _graph(ws: Path, env: dict) -> dict:
    _, nodes = _bash("ros2 node list --no-daemon --spin-time 2", ws, env, 30)
    _, topics = _bash("ros2 topic list -t --no-daemon --spin-time 2", ws, env, 30)
    hidden = ("/rosbag2", "/_ros2cli", "/launch_ros")
    node_list = [n.strip() for n in nodes.splitlines() if n.startswith("/") and not n.strip().startswith(hidden)]
    topic_list = [t.strip() for t in topics.splitlines() if t.startswith("/")]
    return {"nodes": node_list, "topics": topic_list}


# --------------------------------------------------------------------------- job


def check(files: dict[str, str], distro: str, out: Path, launch_seconds: float = LAUNCH_SECONDS) -> dict:
    if distro not in DISTROS:
        raise ValueError(f"distro must be one of {DISTROS}")
    started = time.monotonic()
    root = Path(tempfile.mkdtemp(prefix="ros_check_"))
    ws, home = root / "ws", root / "home"
    home.mkdir(parents=True)
    ws.mkdir()
    env = _env(distro, home)
    try:
        names = write_workspace(files, ws)
        if not names:
            raise ValueError("no package.xml with a <name> in the generated files")
        packages = {}
        for xml in (ws / "src").rglob("package.xml"):
            n = ET.parse(xml).getroot().findtext("name")
            if n:
                packages[n.strip()] = xml.parent
        result: dict = {"distro": distro, "packages": names}
        result["build"] = build(ws, env)
        result["lint"] = lint(ws, env, packages)
        if result["build"]["ok"]:
            result["test"] = test(ws, env)
            result["launch"] = launch(ws, env, packages, launch_seconds)
        result["headline"] = headline(result)
        result["passed"] = bool(result["build"]["ok"] and result.get("launch", {}).get("ok"))

        out.mkdir(parents=True, exist_ok=True)
        for name in ("build.log", "lint.log", "test.log", "launch.log"):
            if (ws / name).is_file():
                data = (ws / name).read_bytes()
                (out / name).write_bytes(data[-MAX_LOG_BYTES:])
        bags = sorted((ws / "bag").glob("*.mcap")) if (ws / "bag").is_dir() else []
        if bags and bags[0].stat().st_size <= MAX_BAG_BYTES:
            shutil.copy(bags[0], out / "bag.mcap")
            result["bag_bytes"] = bags[0].stat().st_size
        result["worker_seconds"] = round(time.monotonic() - started, 1)
        (out / "summary.json").write_text(json.dumps(result, indent=2))
        return result
    finally:
        shutil.rmtree(root, ignore_errors=True)


def headline(r: dict) -> str:
    if not r["build"]["ok"]:
        missing = r["build"].get("missing_dependencies")
        return "does not build" + (f" (missing: {', '.join(missing[:4])})" if missing else "")
    parts = ["builds"]
    la = r.get("launch") or {}
    if la.get("ok"):
        parts.append(f"{len(la['nodes'])} node(s) up, {len(la['topics'])} topic(s)")
    else:
        parts.append(f"launch failed: {la.get('error')}")
    t = r.get("test") or {}
    if t.get("ran"):
        parts.append(f"tests {t['tests'] - t['failures'] - t['errors']}/{t['tests']} passed")
    lint_total = (r.get("lint") or {}).get("total", 0)
    parts.append(f"lint: {lint_total} issue(s)" if lint_total else "lint clean")
    return " · ".join(parts)


def run_check_job(job_input: dict, distro: str) -> dict:
    """Worker entry point: returns the summary plus the logs/bag as inline artifacts (the API stores them)."""
    files = job_input["files"]
    if not isinstance(files, dict) or not files:
        raise ValueError("files must be a non-empty {path: content} mapping")
    with tempfile.TemporaryDirectory() as tmp:
        out = Path(tmp)
        summary = check(files, distro, out)
        artifacts = {p.name: p.read_bytes() for p in sorted(out.iterdir()) if p.is_file()}
    return {"summary": summary, "artifacts": artifacts, "worker_seconds": summary["worker_seconds"]}


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("source", type=Path, help="a directory containing package(s), or a JSON {path: content} file")
    parser.add_argument("--distro", default=os.environ.get("ROS_DISTRO", "jazzy"), choices=DISTROS)
    parser.add_argument("--out", type=Path, default=Path("ros_check_out"))
    parser.add_argument("--launch-seconds", type=float, default=LAUNCH_SECONDS)
    args = parser.parse_args(argv)
    if args.source.is_dir():
        files = {p.relative_to(args.source).as_posix(): p.read_text() for p in args.source.rglob("*") if p.is_file()}
    else:
        files = json.loads(args.source.read_text())
    result = check(files, args.distro, args.out, args.launch_seconds)
    print(json.dumps(result, indent=2))
    return 0 if result["passed"] else 1


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