tlane

#!/usr/bin/env python3

"""Create task-specific Temporal development lanes.

A new lane is a directory directly beneath the lane root (the home directory
by default) containing an independent, colocated jj+Git clone of temporal:

    tlane new my-task
    cd ~/my-task

Promoting a lane adds api, api-go, sdk-go, and cli, then creates a lane-local
go.work tying all five repositories together:

    tlane promote

From a promoted lane, synchronize the current API working copy. A colocated jj
repository stores @ as a Git commit, so it can be fetched directly by commit ID
without creating or moving a bookmark:

    tlane api-sync

Commands that operate on an existing lane infer it from the current directory.
From elsewhere, pass its name explicitly:

    tlane promote my-task
    tlane api-sync my-task

Clones are seeded from the repositories under ~/t by default. Git's local
clone mode hard-links object files when possible, while each lane receives its
own .git and .jj directories. Set TLANE_SOURCE_ROOT and TLANE_ROOT, or use the
matching command-line options, to override the defaults.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import re
import shlex
import subprocess
import sys
from typing import Any, Iterable, Sequence


STATE_FILE = ".tlane.json"
STATE_VERSION = 1
ALL_REPOS = ("api", "api-go", "sdk-go", "cli", "temporal")
PROMOTION_REPOS = ("api", "api-go", "sdk-go", "cli")
GO_MODULES = (
    "./api-go",
    "./sdk-go",
    "./temporal",
    "./cli",
    "./api/cmd/protoc-gen-nexus-rpc-yaml",
    "./api/cmd/protoc-gen-system-nexus-wit",
)
MAIN_BRANCHES = ("main", "master")
PREFERRED_REMOTES = ("up", "origin")
LANE_NAME_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*")


class TlaneError(Exception):
    pass


def log_command(args: Sequence[object], cwd: Path | None = None) -> None:
    prefix = f"cd {shlex.quote(str(cwd))} && " if cwd is not None else ""
    print(f"+++ {prefix}{shlex.join(str(arg) for arg in args)}", flush=True)


def run(
    *args: object,
    cwd: Path | None = None,
    env_extra: dict[str, str] | None = None,
    capture_output: bool = False,
    check: bool = True,
) -> subprocess.CompletedProcess[str]:
    command = [str(arg) for arg in args]
    log_command(command, cwd)
    env = None
    if env_extra:
        env = {**os.environ, **env_extra}
    return subprocess.run(
        command,
        cwd=cwd,
        env=env,
        text=True,
        capture_output=capture_output,
        check=check,
    )


def output(*args: object, cwd: Path | None = None) -> str:
    return run(*args, cwd=cwd, capture_output=True).stdout.strip()


def git(repo: Path, *args: object, **kwargs: Any) -> subprocess.CompletedProcess[str]:
    return run("git", *args, cwd=repo, **kwargs)


def git_output(repo: Path, *args: object) -> str:
    return output("git", *args, cwd=repo)


def jj(repo: Path, *args: object, **kwargs: Any) -> subprocess.CompletedProcess[str]:
    return run("jj", *args, cwd=repo, **kwargs)


def jj_output(repo: Path, *args: object) -> str:
    return output("jj", *args, cwd=repo)


def validate_lane_name(name: str) -> str:
    if not LANE_NAME_RE.fullmatch(name) or name in {".", ".."}:
        raise TlaneError(
            f"invalid lane name {name!r}; use letters, digits, '.', '_', or '-'"
        )
    return name


def lane_path(lane_root: Path, name: str) -> Path:
    return lane_root / validate_lane_name(name)


def state_path(lane: Path) -> Path:
    return lane / STATE_FILE


def read_state(lane: Path) -> dict[str, Any]:
    path = state_path(lane)
    try:
        state = json.loads(path.read_text())
    except FileNotFoundError as err:
        raise TlaneError(f"{lane} is not a tlane lane (missing {STATE_FILE})") from err
    except (OSError, json.JSONDecodeError) as err:
        raise TlaneError(f"cannot read {path}: {err}") from err
    if state.get("version") != STATE_VERSION:
        raise TlaneError(f"unsupported state version in {path}")
    return state


def write_state(lane: Path, state: dict[str, Any]) -> None:
    path = state_path(lane)
    temporary = lane / f"{STATE_FILE}.tmp"
    temporary.write_text(json.dumps(state, indent=2, sort_keys=True) + "\n")
    temporary.replace(path)


def infer_lane(lane_root: Path) -> Path:
    root = lane_root.resolve()
    current = Path.cwd().resolve()
    try:
        relative = current.relative_to(root)
    except ValueError as err:
        raise TlaneError(
            f"current directory is not beneath lane root {lane_root}; pass a lane name"
        ) from err
    if not relative.parts:
        raise TlaneError("current directory is the lane root; pass a lane name")
    lane = root / relative.parts[0]
    read_state(lane)
    return lane


def resolve_lane(lane_root: Path, name: str | None) -> tuple[Path, dict[str, Any]]:
    lane = lane_path(lane_root, name) if name else infer_lane(lane_root)
    if not lane.is_dir():
        raise TlaneError(f"lane does not exist: {lane}")
    return lane, read_state(lane)


def git_ref(repo: Path, ref: str) -> str | None:
    result = git(
        repo,
        "rev-parse",
        "--verify",
        f"{ref}^{{commit}}",
        capture_output=True,
        check=False,
    )
    if result.returncode:
        return None
    return result.stdout.strip()


def remote_names(repo: Path) -> list[str]:
    return [line for line in git_output(repo, "remote").splitlines() if line]


def ordered_remotes(repo: Path) -> list[str]:
    names = remote_names(repo)
    preferred = [name for name in PREFERRED_REMOTES if name in names]
    return preferred + [name for name in names if name not in preferred]


def find_main(repo: Path) -> tuple[str, str | None, str]:
    for remote in ordered_remotes(repo):
        for branch in MAIN_BRANCHES:
            commit = git_ref(repo, f"refs/remotes/{remote}/{branch}")
            if commit:
                return branch, remote, commit
    for branch in MAIN_BRANCHES:
        commit = git_ref(repo, f"refs/heads/{branch}")
        if commit:
            return branch, None, commit
    raise TlaneError(f"cannot find a main/master ref in {repo}")


def copy_remotes(source: Path, destination: Path) -> None:
    for name in remote_names(destination):
        git(destination, "remote", "remove", name)

    for name in remote_names(source):
        fetch_urls = git_output(source, "remote", "get-url", "--all", name).splitlines()
        if not fetch_urls:
            continue
        git(destination, "remote", "add", name, fetch_urls[0])
        for url in fetch_urls[1:]:
            git(destination, "remote", "set-url", "--add", name, url)

        push_urls = git_output(
            source, "remote", "get-url", "--push", "--all", name
        ).splitlines()
        if push_urls and push_urls != fetch_urls:
            for url in push_urls:
                git(destination, "remote", "set-url", "--add", "--push", name, url)


def copy_remote_refs(source: Path, destination: Path) -> None:
    # A normal clone does not copy the source repo's remote-tracking refs.
    # Seed only main/master from the local repo. Importing every upstream PR
    # ref makes jj's initial operation unnecessarily large; feature refs can
    # be fetched normally if api-sync needs one.
    refspecs: list[str] = []
    for remote in remote_names(source):
        for branch in MAIN_BRANCHES:
            ref = f"refs/remotes/{remote}/{branch}"
            if git_ref(source, ref):
                refspecs.append(f"+{ref}:{ref}")
    if refspecs:
        git(destination, "fetch", "--quiet", "--no-tags", source, *refspecs)


def clone_repo(source_root: Path, name: str, destination: Path) -> None:
    source = source_root / name
    if not (source / ".git").exists():
        raise TlaneError(f"clone source is not a Git repository: {source}")
    if destination.exists():
        if (destination / ".git").exists() and (destination / ".jj").exists():
            print(f"+++ keeping existing repository {destination}")
            return
        raise TlaneError(f"refusing to replace existing path: {destination}")

    _, _, base_commit = find_main(source)
    # --local gives us hard-linked Git objects on the normal same-filesystem
    # path. Git cannot make those links across filesystems, so use its regular
    # file transport when (for example) a test lane lives under /tmp.
    local_mode = (
        "--local"
        if source.stat().st_dev == destination.parent.stat().st_dev
        else "--no-local"
    )
    run(
        "git",
        "clone",
        local_mode,
        "--no-checkout",
        "--",
        source,
        destination,
    )
    copy_remotes(source, destination)
    copy_remote_refs(source, destination)
    git(destination, "checkout", "--detach", base_commit)
    run("jj", "git", "init", "--colocate", destination)
    jj(destination, "new", base_commit)


def ensure_symlink(lane: Path, source_root: Path, name: str) -> None:
    source = (source_root / name).resolve()
    if not source.is_file():
        raise TlaneError(f"{name} does not exist: {source}")
    destination = lane / name
    if destination.is_symlink() and destination.resolve() == source:
        return
    if destination.exists() or destination.is_symlink():
        raise TlaneError(f"refusing to replace existing {destination}")
    destination.symlink_to(source)
    print(f"+++ symlink {destination} -> {source}")


def new_lane(lane_root: Path, source_root: Path, name: str) -> None:
    lane = lane_path(lane_root, name)
    if lane.exists():
        raise TlaneError(f"lane already exists: {lane}")
    if not lane_root.is_dir():
        raise TlaneError(f"lane root does not exist: {lane_root}")

    # Validate inputs before leaving behind a partial lane.
    if not (source_root / "temporal" / ".git").exists():
        raise TlaneError(f"missing source repository: {source_root / 'temporal'}")
    if not (source_root / "shell.nix").is_file():
        raise TlaneError(f"missing shell.nix: {source_root / 'shell.nix'}")
    if not (source_root / ".envrc").is_file():
        raise TlaneError(f"missing .envrc: {source_root / '.envrc'}")

    lane.mkdir()
    state = {
        "name": name,
        "promoted": False,
        "source_root": str(source_root.resolve()),
        "version": STATE_VERSION,
    }
    write_state(lane, state)
    ensure_symlink(lane, source_root, "shell.nix")
    ensure_symlink(lane, source_root, ".envrc")
    clone_repo(source_root, "temporal", lane / "temporal")
    print(f"lane ready: {lane}")


def ensure_go_work(lane: Path) -> None:
    work_file = lane / "go.work"
    versions: list[tuple[tuple[int, ...], str]] = []
    for module in GO_MODULES:
        go_mod = lane / module / "go.mod"
        try:
            contents = go_mod.read_text()
        except OSError as err:
            raise TlaneError(f"cannot read {go_mod}: {err}") from err
        match = re.search(r"(?m)^go\s+([0-9]+(?:\.[0-9]+)+)\s*$", contents)
        if not match:
            raise TlaneError(f"cannot find a go directive in {go_mod}")
        version = match.group(1)
        versions.append((tuple(int(part) for part in version.split(".")), version))

    go_version = max(versions)[1]
    uses = "\n".join(f"\t{module}" for module in GO_MODULES)
    contents = f"go {go_version}\n\nuse (\n{uses}\n)\n"
    work_file.write_text(contents)
    print(f"+++ write {work_file}")


def promote_lane(lane_root: Path, name: str | None) -> None:
    lane, state = resolve_lane(lane_root, name)
    source_root = Path(state["source_root"])
    ensure_symlink(lane, source_root, "shell.nix")
    ensure_symlink(lane, source_root, ".envrc")
    if not (lane / "temporal" / ".git").exists():
        raise TlaneError(f"lane has no temporal repository: {lane}")

    for repo in PROMOTION_REPOS:
        clone_repo(source_root, repo, lane / repo)
    ensure_go_work(lane)
    state["promoted"] = True
    write_state(lane, state)
    print(f"promoted lane ready: {lane}")


def diff_for_paths(repo: Path, paths: Iterable[str]) -> bool:
    return bool(jj_output(repo, "diff", "--summary", "--", *paths))


def jj_commit(repo: Path, message: str, *paths: str) -> None:
    if not diff_for_paths(repo, paths):
        return
    jj(repo, "commit", "-m", f"NOT FOR MERGE: {message}", "--", *paths)


def make_proto(
    repo: Path,
    *commit_paths: str,
    gowork_off: bool = False,
    make_vars: dict[str, str] | None = None,
) -> None:
    env_extra = {"GOWORK": "off"} if gowork_off else None
    variables = [f"{key}={value}" for key, value in (make_vars or {}).items()]
    run("make", "proto", *variables, cwd=repo, env_extra=env_extra)
    if commit_paths:
        jj_commit(repo, "make proto", *commit_paths)


def current_jj_commit(repo: Path) -> str:
    commit = jj_output(repo, "log", "--no-graph", "-r", "@", "-T", "commit_id")
    git_commit = git_output(repo, "rev-parse", f"{commit}^{{commit}}")
    if git_commit != commit:
        raise TlaneError(
            f"jj @ is not available from Git in {repo}; is this repository colocated?"
        )
    return commit


def set_api_submodule(api_go: Path, api: Path) -> None:
    api_commit = current_jj_commit(api)
    submodule = api_go / "proto/api"
    # Override the tracked URL only in this clone's config. .gitmodules stays
    # untouched and no branch is involved in synchronization.
    git(api_go, "config", "submodule.proto/api.url", api.resolve())
    if not (submodule / ".git").exists():
        git(
            api_go,
            "-c",
            "protocol.file.allow=always",
            "submodule",
            "update",
            "--init",
            "--recursive",
            "proto/api",
        )
    git(
        submodule,
        "-c",
        "protocol.file.allow=always",
        "fetch",
        "--quiet",
        "--no-tags",
        api.resolve(),
        api_commit,
    )
    fetched_commit = git_output(submodule, "rev-parse", "FETCH_HEAD^{commit}")
    if fetched_commit != api_commit:
        raise TlaneError(
            f"fetched API commit {fetched_commit} does not match jj @ {api_commit}"
        )
    # This submodule is a disposable view of ../api. Resetting also makes a
    # half-finished previous sync harmless; API edits belong in the API repo.
    git(submodule, "reset", "--hard", api_commit)
    git(
        submodule,
        "-c",
        "protocol.file.allow=always",
        "submodule",
        "update",
        "--init",
        "--recursive",
    )


def go_module_version(repo: Path, module: str) -> str:
    version = run(
        "go",
        "list",
        "-m",
        "-f",
        "{{.Version}}",
        module,
        cwd=repo,
        env_extra={"GOWORK": "off"},
        capture_output=True,
    ).stdout.strip()
    if not version:
        raise TlaneError(f"empty version for {module} in {repo}")
    return version


def require_promoted(lane: Path, state: dict[str, Any]) -> None:
    missing = [repo for repo in ALL_REPOS if not (lane / repo / ".git").exists()]
    if not state.get("promoted") or missing or not (lane / "go.work").exists():
        detail = f"; missing: {', '.join(missing)}" if missing else ""
        raise TlaneError(f"lane is not promoted; run 'tlane promote'{detail}")


def api_sync(lane_root: Path, name: str | None) -> None:
    lane, state = resolve_lane(lane_root, name)
    source_root = Path(state["source_root"])
    ensure_symlink(lane, source_root, "shell.nix")
    ensure_symlink(lane, source_root, ".envrc")
    require_promoted(lane, state)
    ensure_go_work(lane)

    api = lane / "api"
    make_proto(api, "openapi", "buf.lock")

    api_go = lane / "api-go"
    set_api_submodule(api_go, api)
    make_proto(api_go, ".", gowork_off=True)

    temporal = lane / "temporal"
    api_version = go_module_version(temporal, "go.temporal.io/api")
    make_proto(
        temporal,
        "api",
        "proto/api.binpb",
        make_vars={"GO_API_VER": api_version},
    )

    state.pop("api_branch", None)
    state.pop("api_bookmark", None)
    write_state(lane, state)
    print(f"synchronized API @ in {lane}")


def default_path(environment_name: str, fallback: Path) -> Path:
    return Path(os.environ.get(environment_name, fallback)).expanduser().resolve()


def make_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--lane-root",
        type=Path,
        default=default_path("TLANE_ROOT", Path.home()),
        help="directory whose immediate children are lanes (default: home directory)",
    )
    parser.add_argument(
        "--source-root",
        type=Path,
        default=default_path("TLANE_SOURCE_ROOT", Path.home() / "t"),
        help="directory containing the source repositories and shell.nix",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    new_parser = subparsers.add_parser("new", help="create a Temporal-only lane")
    new_parser.add_argument("name", help="name of the new directory beneath the lane root")

    promote_parser = subparsers.add_parser(
        "promote", help="add api, api-go, sdk-go, cli, and go.work"
    )
    promote_parser.add_argument(
        "name", nargs="?", help="lane name (inferred when run from a lane)"
    )

    sync_parser = subparsers.add_parser(
        "api-sync", help="synchronize API changes in a promoted lane"
    )
    sync_parser.add_argument(
        "name", nargs="?", help="lane name (inferred when run from a lane)"
    )
    return parser


def main() -> int:
    parser = make_parser()
    args = parser.parse_args()
    lane_root = args.lane_root.expanduser().resolve()
    source_root = args.source_root.expanduser().resolve()
    try:
        if args.command == "new":
            new_lane(lane_root, source_root, args.name)
        elif args.command == "promote":
            promote_lane(lane_root, args.name)
        elif args.command == "api-sync":
            api_sync(lane_root, args.name)
        else:
            parser.error(f"unknown command: {args.command}")
    except (TlaneError, subprocess.CalledProcessError) as err:
        print(f"tlane: {err}", file=sys.stderr)
        return 1
    return 0


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