stem-launcher/workspace.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import re
import shlex
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import unquote, urlsplit


def expand_path(raw_path: str, base_dir: Path) -> Path:
    expanded = os.path.expandvars(os.path.expanduser(raw_path))
    path = Path(expanded)
    if not path.is_absolute():
        path = base_dir / path
    return path.resolve()


@dataclass
class GitConfig:
    base_url: str
    source_org: str
    answer_org: str
    source_remote: str
    answer_remote: str
    origin_remote: str
    fallback_branch: str


@dataclass
class WorkspaceConfig:
    config_path: Path
    workspace_root: Path
    series_root: Path
    socket_root: Path
    token_path: Path
    legacy_token_path: Path
    tools_root_candidates: list[Path]
    tools_root: Path
    defaults: dict[str, str]
    git: GitConfig


def load_config(config_path: Path) -> WorkspaceConfig:
    raw = json.loads(config_path.read_text(encoding="utf-8"))
    base_dir = config_path.parent

    workspace_root = expand_path(raw.get("workspace_root", "~/dev/workspace/rv"), base_dir)
    series_root = expand_path(raw.get("series_root", str(workspace_root / "series")), base_dir)
    socket_root = expand_path(raw.get("socket_root", str(workspace_root / "sockets")), base_dir)
    token_path = expand_path(raw.get("token_file", str(workspace_root / "tokens" / "tokens.json")), base_dir)
    legacy_token_path = expand_path(
        raw.get("legacy_token_file", str(workspace_root / "tokens" / "gitea_tokens.json")),
        base_dir,
    )

    candidate_values = raw.get("tools_root_candidates")
    if not candidate_values:
        candidate_values = [raw.get("tools_root", "~/dev/workspace/tools/rv32i-hazard3-env")]
    tools_root_candidates = [expand_path(value, base_dir) for value in candidate_values]

    tools_root = next((path for path in tools_root_candidates if path.exists()), tools_root_candidates[0])
    defaults = raw.get("defaults", {})
    git_raw = raw.get("git", {})
    git = GitConfig(
        base_url=git_raw.get("base_url", "http://77.90.8.171:3001").rstrip("/"),
        source_org=git_raw.get("source_org", "edu-inf"),
        answer_org=git_raw.get("answer_org", "zsl-inf"),
        source_remote=git_raw.get("source_remote", "r1"),
        answer_remote=git_raw.get("answer_remote", "a1"),
        origin_remote=git_raw.get("origin_remote", "origin"),
        fallback_branch=git_raw.get("fallback_branch", "build"),
    )

    return WorkspaceConfig(
        config_path=config_path,
        workspace_root=workspace_root,
        series_root=series_root,
        socket_root=socket_root,
        token_path=token_path,
        legacy_token_path=legacy_token_path,
        tools_root_candidates=tools_root_candidates,
        tools_root=tools_root,
        defaults=defaults,
        git=git,
    )


def series_directories(config: WorkspaceConfig) -> list[Path]:
    if not config.series_root.is_dir():
        raise SystemExit(f"Missing series root: {config.series_root}")
    return sorted(path for path in config.series_root.iterdir() if path.is_dir())


def card_directories(series_path: Path) -> list[Path]:
    return sorted(path for path in series_path.iterdir() if path.is_dir())


def read_card_title(card_path: Path) -> str:
    readme_path = card_path / "README.md"
    if not readme_path.is_file():
        return ""

    for line in readme_path.read_text(encoding="utf-8").splitlines():
        stripped = line.strip()
        if stripped.startswith("#"):
            return stripped.lstrip("#").strip()
    return ""


def resolve_selector(
    config: WorkspaceConfig,
    series_arg: str | None,
    card_arg: str | None,
) -> tuple[str, str]:
    default_series = config.defaults.get("series")
    default_card = config.defaults.get("card")

    if series_arg and "/" in series_arg:
        if card_arg:
            raise SystemExit("Pass either '<series> <card>' or '<series/card>', not both forms at once.")
        series_value, card_value = series_arg.split("/", 1)
        return series_value, card_value

    if series_arg and not card_arg:
        if (config.series_root / series_arg).is_dir():
            if not default_card:
                raise SystemExit("Missing default card in config.")
            return series_arg, default_card
        if default_series and (config.series_root / default_series / series_arg).is_dir():
            return default_series, series_arg

    series_value = series_arg or default_series
    card_value = card_arg or default_card
    if not series_value or not card_value:
        raise SystemExit("Missing series/card selector and no defaults are configured.")
    return series_value, card_value


def resolve_card_path(config: WorkspaceConfig, series_name: str, card_name: str) -> Path:
    card_path = config.series_root / series_name / card_name
    if not card_path.is_dir():
        raise SystemExit(f"Missing card path: {card_path}")
    return card_path


def print_series(config: WorkspaceConfig) -> None:
    for series_path in series_directories(config):
        card_count = len(card_directories(series_path))
        print(f"{series_path.name}\t{card_count}\t{series_path}")


def print_cards(config: WorkspaceConfig, series_arg: str | None) -> None:
    series_name = series_arg or config.defaults.get("series")
    if not series_name:
        raise SystemExit("Missing series name and no default series is configured.")

    series_path = config.series_root / series_name
    if not series_path.is_dir():
        raise SystemExit(f"Missing series path: {series_path}")

    for card_path in card_directories(series_path):
        title = read_card_title(card_path)
        if title:
            print(f"{card_path.name}\t{title}\t{card_path}")
        else:
            print(f"{card_path.name}\t{card_path}")


def slug_value(raw_value: str, field_name: str) -> str:
    slug = re.sub(r"[^a-z0-9._-]+", "-", raw_value.strip().lower()).strip("-")
    if not slug:
        raise SystemExit(f"Empty {field_name} after slug conversion: {raw_value!r}")
    return slug


def git_capture(repo_path: Path, git_args: list[str]) -> str | None:
    result = subprocess.run(
        ["git", "-C", str(repo_path)] + git_args,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        return None
    return result.stdout.strip()


def git_repo_root(repo_path: Path) -> Path | None:
    repo_root = git_capture(repo_path, ["rev-parse", "--show-toplevel"])
    if not repo_root:
        return None
    return Path(repo_root)


def migrate_legacy_token_store(config: WorkspaceConfig) -> None:
    if config.token_path.exists() or not config.legacy_token_path.exists():
        return

    config.token_path.parent.mkdir(parents=True, exist_ok=True)
    os.chmod(config.token_path.parent, 0o700)
    config.legacy_token_path.replace(config.token_path)
    os.chmod(config.token_path, 0o600)


def load_token_store(config: WorkspaceConfig) -> dict:
    migrate_legacy_token_store(config)

    if not config.token_path.exists():
        return {"users": {}}

    raw_data = json.loads(config.token_path.read_text(encoding="utf-8"))
    if not isinstance(raw_data, dict):
        raise SystemExit(f"Invalid token store format: {config.token_path}")
    raw_data.setdefault("users", {})
    return raw_data


def write_token_store(config: WorkspaceConfig, token_data: dict) -> None:
    config.token_path.parent.mkdir(parents=True, exist_ok=True)
    os.chmod(config.token_path.parent, 0o700)

    temp_path = config.token_path.with_suffix(config.token_path.suffix + ".tmp")
    temp_path.write_text(json.dumps(token_data, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    os.chmod(temp_path, 0o600)
    temp_path.replace(config.token_path)
    os.chmod(config.token_path, 0o600)


def next_token_name(tokens: dict[str, str]) -> str:
    index = 1
    while f"t{index}" in tokens:
        index += 1
    return f"t{index}"


def register_token(token_data: dict, user_name: str, token_value: str) -> bool:
    users = token_data.setdefault("users", {})
    user_entry = users.setdefault(user_name, {})
    tokens = user_entry.setdefault("tokens", {})

    for existing_token in tokens.values():
        if existing_token == token_value:
            return False

    tokens[next_token_name(tokens)] = token_value
    return True


def remote_urls(repo_path: Path) -> list[str]:
    remote_output = git_capture(repo_path, ["remote"])
    if not remote_output:
        return []

    urls: list[str] = []
    seen: set[str] = set()
    for remote_name in remote_output.splitlines():
        url_output = git_capture(repo_path, ["remote", "get-url", "--all", remote_name])
        if not url_output:
            continue
        for remote_url in url_output.splitlines():
            if remote_url in seen:
                continue
            seen.add(remote_url)
            urls.append(remote_url)
    return urls


def remote_credentials(remote_url: str) -> tuple[str, str] | None:
    split_url = urlsplit(remote_url)
    if split_url.scheme not in {"http", "https"}:
        return None
    if split_url.username is None or split_url.password is None:
        return None

    return unquote(split_url.username), unquote(split_url.password)


def sync_tokens_from_repo(config: WorkspaceConfig, repo_path: Path) -> bool:
    repo_root = git_repo_root(repo_path)
    if repo_root is None:
        migrate_legacy_token_store(config)
        return False

    found_credentials: list[tuple[str, str]] = []
    for remote_url in remote_urls(repo_root):
        credentials = remote_credentials(remote_url)
        if credentials is not None:
            found_credentials.append(credentials)

    if not found_credentials:
        migrate_legacy_token_store(config)
        return False

    token_data = load_token_store(config)
    changed = False
    for user_name, token_value in found_credentials:
        if register_token(token_data, user_name, token_value):
            changed = True

    if changed or not config.token_path.exists():
        write_token_store(config, token_data)

    return changed


def git_current_branch(config: WorkspaceConfig, repo_path: Path) -> str:
    branch_name = git_capture(repo_path, ["branch", "--show-current"])
    if branch_name:
        return branch_name

    head_ref = git_capture(repo_path, ["symbolic-ref", "--short", f"refs/remotes/{config.git.origin_remote}/HEAD"])
    if head_ref and "/" in head_ref:
        return head_ref.split("/", 1)[1]

    return config.git.fallback_branch


def fallback_source_repo_name(series_name: str, card_name: str) -> str:
    return "-".join([slug_value(series_name, "series"), slug_value(card_name, "card")])


def fallback_source_url(config: WorkspaceConfig, series_name: str, card_name: str) -> str:
    repo_name = fallback_source_repo_name(series_name, card_name)
    return f"{config.git.base_url}/{config.git.source_org}/{repo_name}.git"


def discover_source_url(config: WorkspaceConfig, repo_path: Path, series_name: str, card_name: str) -> str:
    source_url = git_capture(repo_path, ["remote", "get-url", config.git.origin_remote])
    if source_url:
        return source_url

    source_url = git_capture(repo_path, ["remote", "get-url", config.git.source_remote])
    if source_url:
        return source_url

    return fallback_source_url(config, series_name, card_name)


def repo_name_from_url(remote_url: str) -> str:
    repo_name = remote_url.rstrip("/").rsplit("/", 1)[-1]
    if repo_name.endswith(".git"):
        repo_name = repo_name[:-4]
    return slug_value(repo_name, "source repo")


def build_answer_repo_name(source_repo_name: str, class_name: str, date_value: str) -> str:
    return "-".join([source_repo_name, slug_value(class_name, "class"), date_value])


def build_answer_url(config: WorkspaceConfig, answer_repo_name: str) -> str:
    return f"{config.git.base_url}/{config.git.answer_org}/{answer_repo_name}.git"


def set_git_remote(repo_path: Path, remote_name: str, remote_url: str) -> str:
    existing_url = git_capture(repo_path, ["remote", "get-url", remote_name])
    if existing_url == remote_url:
        return "unchanged"
    if existing_url:
        subprocess.run(
            ["git", "-C", str(repo_path), "remote", "set-url", remote_name, remote_url],
            check=True,
        )
        return "updated"

    subprocess.run(
        ["git", "-C", str(repo_path), "remote", "add", remote_name, remote_url],
        check=True,
    )
    return "added"


def print_submission_plan(config: WorkspaceConfig, args: argparse.Namespace) -> None:
    series_name, card_name = resolve_selector(config, args.series, args.card)
    selector = f"{series_name}/{card_name}"
    card_path = resolve_card_path(config, series_name, card_name)
    sync_tokens_from_repo(config, card_path)
    class_name = slug_value(args.class_name, "class")
    branch_name = slug_value(args.branch or args.nick, "branch")
    source_url = args.source_url or discover_source_url(config, card_path, series_name, card_name)
    source_repo_name = repo_name_from_url(source_url)
    answer_repo_name = build_answer_repo_name(source_repo_name, class_name, args.date)
    answer_url = build_answer_url(config, answer_repo_name)
    source_branch = git_current_branch(config, card_path)

    print(f"selector\t{selector}")
    print(f"card_path\t{card_path}")
    print(f"source_repo\t{config.git.source_org}/{source_repo_name}")
    print(f"source_remote\t{config.git.source_remote}")
    print(f"source_url\t{source_url}")
    print(f"source_branch\t{source_branch}")
    print(f"answer_remote\t{config.git.answer_remote}")
    print(f"answer_repo\t{config.git.answer_org}/{answer_repo_name}")
    print(f"answer_url\t{answer_url}")
    print(f"student_branch\t{branch_name}")

    if args.apply:
        source_status = set_git_remote(card_path, config.git.source_remote, source_url)
        answer_status = set_git_remote(card_path, config.git.answer_remote, answer_url)
        print(f"apply_{config.git.source_remote}\t{source_status}")
        print(f"apply_{config.git.answer_remote}\t{answer_status}")

    print("commands")
    print(f"  git -C {shlex.quote(str(card_path))} remote add {config.git.source_remote} {shlex.quote(source_url)}")
    print(f"  git -C {shlex.quote(str(card_path))} remote add {config.git.answer_remote} {shlex.quote(answer_url)}")
    print(f"  git -C {shlex.quote(str(card_path))} fetch {config.git.source_remote} {shlex.quote(source_branch)}")
    print(f"  git -C {shlex.quote(str(card_path))} switch -c {shlex.quote(branch_name)}")
    print(f"  git -C {shlex.quote(str(card_path))} push -u {config.git.answer_remote} {shlex.quote(branch_name)}")


def tmux_target_exists(session_name: str) -> bool:
    result = subprocess.run(
        ["tmux", "has-session", "-t", session_name],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        check=False,
    )
    return result.returncode == 0


def build_container_command(
    config: WorkspaceConfig,
    selector: str,
    card_path: Path,
    instance: str,
) -> str:
    tool_root = config.tools_root
    rv_script = tool_root / "rv"
    if not rv_script.is_file():
        candidates = "\n".join(f"  - {path}" for path in config.tools_root_candidates)
        raise SystemExit(f"Missing rv launcher. Checked:\n{candidates}")

    env_parts = [
        f"WORKSPACE_ROOT={shlex.quote(str(config.workspace_root))}",
        f"RV_REPO_PATH={shlex.quote(str(card_path))}",
        f"RV_CARD={shlex.quote(selector)}",
        f"RV_INSTANCE={shlex.quote(instance)}",
    ]
    return " && ".join(
        [
            f"cd {shlex.quote(str(tool_root))}",
            " ".join(env_parts + ["bash", "./rv", "shell"]),
        ]
    )


def run_tmux_container(config: WorkspaceConfig, args: argparse.Namespace) -> None:
    series_name, card_name = resolve_selector(config, args.series, args.card)
    selector = f"{series_name}/{card_name}"
    card_path = resolve_card_path(config, series_name, card_name)

    session_name = args.session or config.defaults.get("tmux_session", "rv-workspace")
    window_name = args.window or config.defaults.get("tmux_window", "rv")
    instance = args.instance or config.defaults.get("instance", "shell")

    if tmux_target_exists(session_name):
        raise SystemExit(
            f"tmux session already exists: {session_name}\n"
            "Use another --session name or remove the existing session first."
        )

    shell_command = build_container_command(config, selector, card_path, instance)
    tmux_command = [
        "tmux",
        "new-session",
        "-d",
        "-s",
        session_name,
        "-n",
        window_name,
        f"bash -lc {shlex.quote(shell_command)}",
    ]

    if args.dry_run:
        print("Selector:", selector)
        print("Card path:", card_path)
        print("Tools root:", config.tools_root)
        print("Pane 0 command:")
        print(shell_command)
        print("tmux command:")
        print(shlex.join(tmux_command))
        return

    subprocess.run(tmux_command, check=True)

    if args.attach:
        subprocess.run(["tmux", "attach", "-t", session_name], check=True)
        return

    print(f"Created tmux session: {session_name}")
    print(f"Window: {window_name}")
    print("Pane 0: container shell")
    print(f"Attach with: tmux attach -t {session_name}")


def print_config(config: WorkspaceConfig) -> None:
    print(f"config_path\t{config.config_path}")
    print(f"workspace_root\t{config.workspace_root}")
    print(f"series_root\t{config.series_root}")
    print(f"socket_root\t{config.socket_root}")
    print(f"token_path\t{config.token_path}")
    print(f"tools_root\t{config.tools_root}")
    print(f"git_base_url\t{config.git.base_url}")
    print(f"git_source_org\t{config.git.source_org}")
    print(f"git_answer_org\t{config.git.answer_org}")
    print(f"git_source_remote\t{config.git.source_remote}")
    print(f"git_answer_remote\t{config.git.answer_remote}")
    print(f"git_origin_remote\t{config.git.origin_remote}")
    print(f"git_fallback_branch\t{config.git.fallback_branch}")
    print("tools_root_candidates")
    for path in config.tools_root_candidates:
        print(f"  {path}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Workspace launcher for rv cards and tools.")
    parser.add_argument(
        "--config",
        default=Path(__file__).with_name("workspace.json"),
        type=Path,
        help="Path to workspace.json",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    subparsers.add_parser("show-config", help="Print resolved workspace paths.")
    subparsers.add_parser("list-series", help="List series directories from series_root.")

    cards_parser = subparsers.add_parser("list-cards", help="List cards in a series.")
    cards_parser.add_argument("series", nargs="?", help="Series id, for example 'inf'.")

    tmux_parser = subparsers.add_parser(
        "tmux-container",
        help="Create a tmux session with pane 0 running the container shell for a selected card.",
    )
    tmux_parser.add_argument("series", nargs="?", help="Series id or full selector like 'inf/03'.")
    tmux_parser.add_argument("card", nargs="?", help="Card id, for example '03'.")
    tmux_parser.add_argument("--session", help="tmux session name.")
    tmux_parser.add_argument("--window", help="tmux window name.")
    tmux_parser.add_argument("--instance", help="RV_INSTANCE passed to the tool launcher.")
    tmux_parser.add_argument("--attach", action="store_true", help="Attach to the session after creation.")
    tmux_parser.add_argument("--dry-run", action="store_true", help="Print commands instead of running tmux.")

    submission_parser = subparsers.add_parser(
        "submission",
        help="Print or apply source/answer git remotes for a selected card.",
    )
    submission_parser.add_argument("series", nargs="?", help="Series id or full selector like 'inf/03'.")
    submission_parser.add_argument("card", nargs="?", help="Card id, for example '03'.")
    submission_parser.add_argument("--class", dest="class_name", required=True, help="Class id, for example '4i'.")
    submission_parser.add_argument("--nick", required=True, help="Student nick used as the branch name.")
    submission_parser.add_argument(
        "--branch",
        help="Override the submission branch name. Defaults to the value of --nick.",
    )
    submission_parser.add_argument(
        "--date",
        default=dt.date.today().isoformat(),
        help="Lesson date used in the answer repo name. Default: today in YYYY-MM-DD.",
    )
    submission_parser.add_argument(
        "--source-url",
        help="Override the source repo URL. By default launcher reads the current origin from the card repo.",
    )
    submission_parser.add_argument(
        "--apply",
        action="store_true",
        help="Write remotes into the selected card repo instead of only printing the plan.",
    )

    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    config = load_config(args.config)
    sync_tokens_from_repo(config, Path.cwd())

    if args.command == "show-config":
        print_config(config)
        return 0
    if args.command == "list-series":
        print_series(config)
        return 0
    if args.command == "list-cards":
        print_cards(config, args.series)
        return 0
    if args.command == "tmux-container":
        run_tmux_container(config, args)
        return 0
    if args.command == "submission":
        print_submission_plan(config, args)
        return 0

    parser.error(f"Unsupported command: {args.command}")
    return 2


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