import glob
import json
import os
import re
import sys
import time
import unicodedata
from datetime import datetime, timezone

import requests
from github import Auth, Github

# Import scripts/ci/runner_configs.py (sibling-up dir) for runner_config -> runs_on lookup.
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
import runner_configs as _runner_configs  # noqa: E402

# rerun-test workflow doesn't build sgl-kernel, so b200 stages always use the
# non-kernel pool when resolving the `$b200_runner` sentinel from runner_configs.yml.
_B200_DEFAULT_RUNNER = "4-gpu-b200"

# install_script values from runner_configs.yml are passed verbatim into a
# `bash ${{ inputs.install_script }}` step in rerun-test.yml. GHA expression
# substitution happens before bash parses, so shell metacharacters in the
# string would inject. Restrict the allowed shape to `scripts/ci/cuda/*.sh`
# (single path component under that dir, no whitespace/operators).
_ALLOWED_INSTALL_SCRIPT = re.compile(r"^scripts/ci/cuda/[\w.-]+\.sh$")

# Configuration
PERMISSIONS_FILE_PATH = ".github/CI_PERMISSIONS.json"


MAINTENANCE_ISSUE_NUMBER = 21065


def _check_rebase_gate(gh_repo, pr, token):
    """
    Pre-dispatch gate mirroring `.github/actions/check-maintenance/action.yml`.

    Without this, /rerun-test would dispatch a workflow_run
    on a PR that's behind a required base, the action would catch it, and
    every job in the run would fail at the gate — wasting runner time and
    producing N error annotations instead of one comment. Pre-checking here
    short-circuits the dispatch and posts a single explanatory comment.

    Mirrors the action's two independent modes driven by issue #21065:
      (1) Full-pause: maintenance issue is OPEN
      (2) Rebase-required: issue body contains `MIN_BASE_SHA: <sha>`
    Both bypassed by the `bypass-maintenance` PR label.

    Returns (allowed: bool, message: Optional[str]). When allowed=False,
    caller MUST post `message` to the PR and skip dispatch.
    Fail-open on API errors (matches the action's behavior).
    """
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/vnd.github+json",
        "X-GitHub-Api-Version": "2022-11-28",
    }
    repo_full_name = gh_repo.full_name

    try:
        issue_resp = requests.get(
            f"https://api.github.com/repos/{repo_full_name}/issues/{MAINTENANCE_ISSUE_NUMBER}",
            headers=headers,
            timeout=15,
        )
        if issue_resp.status_code != 200:
            print(
                f"check_rebase_gate: issue fetch returned {issue_resp.status_code}; fail-open"
            )
            return True, None
        issue_data = issue_resp.json()
    except Exception as e:
        print(f"check_rebase_gate: issue fetch failed ({e}); fail-open")
        return True, None

    issue_state = (issue_data.get("state") or "").lower()
    issue_body = issue_data.get("body") or ""

    min_base_sha = None
    # First MIN_BASE_SHA: <sha> line wins. Match the action's parser:
    # tolerate optional backticks and either ':' or '=' separator.
    for line in issue_body.replace("\r", "").split("\n"):
        m = re.match(
            r"^\s*`?MIN_BASE_SHA`?\s*[:=]\s*`?([A-Fa-f0-9]+)`?",
            line,
        )
        if m:
            candidate = m.group(1)
            if 7 <= len(candidate) <= 40:
                min_base_sha = candidate
            break

    gate_active = (issue_state == "open") or bool(min_base_sha)
    if not gate_active:
        return True, None

    bypass = any(
        (lbl.name if hasattr(lbl, "name") else lbl.get("name")) == "bypass-maintenance"
        for lbl in pr.get_labels()
    )
    if bypass:
        print("check_rebase_gate: PR has bypass-maintenance label; allowing dispatch")
        return True, None

    if issue_state == "open":
        msg = (
            "## ⚠️ CI Maintenance Mode is Active\n"
            "The CI infrastructure is currently under maintenance. "
            "All PR CI runs are paused until maintenance is complete. "
            "**Merging non-CI-fix PRs is prohibited during maintenance mode.**\n\n"
            f"Follow [issue #{MAINTENANCE_ISSUE_NUMBER}]"
            f"(https://github.com/{repo_full_name}/issues/{MAINTENANCE_ISSUE_NUMBER}) "
            "for status updates. Re-run was not dispatched."
        )
        return False, msg

    # MIN_BASE_SHA set, issue not OPEN — check rebase status.
    pr_head_sha = pr.head.sha
    try:
        compare_resp = requests.get(
            f"https://api.github.com/repos/{repo_full_name}/compare/{min_base_sha}...{pr_head_sha}",
            headers=headers,
            timeout=15,
        )
        if compare_resp.status_code != 200:
            print(
                f"check_rebase_gate: compare API returned {compare_resp.status_code}; fail-open"
            )
            return True, None
        status = compare_resp.json().get("status", "unknown")
    except Exception as e:
        print(f"check_rebase_gate: compare API failed ({e}); fail-open")
        return True, None

    if status in ("ahead", "identical"):
        return True, None

    msg = (
        "## ⚠️ Rebase Required Before Re-run\n"
        f"A major update has landed on `main`. Your PR is `{status}` relative "
        f"to required base commit `{min_base_sha[:12]}`.\n\n"
        "**Re-run was not dispatched.** What to do:\n"
        "- Rebase your branch onto the latest `main` and push again\n"
        f"- Follow [issue #{MAINTENANCE_ISSUE_NUMBER}]"
        f"(https://github.com/{repo_full_name}/issues/{MAINTENANCE_ISSUE_NUMBER}) for context\n"
        "- CI-fix PRs may request the `bypass-maintenance` label to skip this check"
    )
    return False, msg


def find_workflow_run_url(
    gh_repo,
    workflow_id,
    ref,
    target_stage,
    token,
    dispatch_time,
    pr_head_sha=None,
    max_wait=30,
    test_command=None,
):
    """
    Poll for the workflow run URL after dispatch.

    Uses the dynamic run-name feature to identify runs:
    - Fork PRs: display_title = "[stage-name] sha"
    - Non-fork PRs: display_title = "[stage-name]"

    Args:
        gh_repo: PyGithub repository object
        workflow_id: ID of the workflow that was dispatched
        ref: Branch/ref the workflow was dispatched on
        target_stage: The stage name we're looking for
        token: GitHub API token
        dispatch_time: Unix timestamp when dispatch was triggered
        pr_head_sha: PR head SHA (for fork PRs, used to match display_title)
        max_wait: Maximum seconds to wait for the run to appear

    Returns:
        The workflow run URL if found, None otherwise.
    """
    # Build expected display_title based on workflow's run-name.
    # rerun-test includes test_command: "[rerun-test] <test_command> [<sha>]"
    # Other workflows: "[stage-name] [<sha>]"
    suffix = f" {test_command}" if test_command else ""
    if pr_head_sha:
        expected_title = f"[{target_stage}]{suffix} {pr_head_sha}"
    else:
        expected_title = f"[{target_stage}]{suffix}"

    print(f"Looking for workflow run with display_title: {expected_title}")

    for attempt in range(max_wait // 5):
        time.sleep(5)

        # Get recent workflow_dispatch runs for this workflow
        runs_url = f"https://api.github.com/repos/{gh_repo.full_name}/actions/workflows/{workflow_id}/runs"
        runs_resp = requests.get(
            runs_url,
            params={"event": "workflow_dispatch", "branch": ref, "per_page": 10},
            headers={
                "Authorization": f"Bearer {token}",
                "Accept": "application/vnd.github+json",
            },
        )

        if runs_resp.status_code != 200:
            print(f"Failed to fetch workflow runs: {runs_resp.status_code}")
            continue

        for run in runs_resp.json().get("workflow_runs", []):
            # Skip runs created before our dispatch (with 10s tolerance)
            run_created = datetime.fromisoformat(
                run["created_at"].replace("Z", "+00:00")
            ).timestamp()
            if run_created < dispatch_time - 10:
                continue

            # Match by display_title (set by workflow's run-name directive)
            # This is immediately available, unlike job names which require waiting.
            #
            # GitHub caps display_title at 512 chars and appends "..." when the
            # workflow's run-name is longer (e.g. /rerun-test with many files
            # newline-joins them into one long command). Accept the un-truncated
            # prefix when the rest matches our submission verbatim — concurrent
            # dispatches use different commands (different files per group), so
            # the prefix is enough to disambiguate.
            display_title = run.get("display_title", "")
            if display_title == expected_title or (
                display_title.endswith("...")
                and expected_title.startswith(display_title[:-3])
            ):
                print(
                    f"Found matching workflow run: {run['id']} with title '{display_title}'"
                )
                return run["html_url"]

    print(f"Could not find workflow run after {max_wait} seconds")
    return None


def get_env_var(name):
    val = os.getenv(name)
    if not val:
        print(f"Error: Environment variable {name} not set.")
        sys.exit(1)
    return val


def _strip_format_chars(s):
    """Remove Unicode format characters (category Cf: LRM/RLM U+200E/200F,
    zero-width space/joiners U+200B-200D, word joiner U+2060, BOM U+FEFF).

    GitHub's copy-path button and rich-text copy inject these invisibly;
    a pasted `/rerun-test foo.py<U+200E>` then never matches any test file
    (see PR #31059). They are display hints and never legitimate in a
    command or path, so dropping them is always safe.
    """
    return "".join(c for c in s if unicodedata.category(c) != "Cf")


def load_permissions(user_login):
    """
    Reads the permissions JSON from the local file system and returns
    the permissions dict for the specific user.
    """
    try:
        print(f"Loading permissions from {PERMISSIONS_FILE_PATH}...")
        if not os.path.exists(PERMISSIONS_FILE_PATH):
            print(f"Error: Permissions file not found at {PERMISSIONS_FILE_PATH}")
            return None

        with open(PERMISSIONS_FILE_PATH, "r") as f:
            data = json.load(f)

        user_perms = data.get(user_login)

        if not user_perms:
            print(f"User '{user_login}' not found in permissions file.")
            return None

        return user_perms

    except Exception as e:
        print(f"Failed to load or parse permissions file: {e}")
        sys.exit(1)


def has_sgl_kernel_changes(pr):
    """
    Check if the PR has changes to the AOT kernel directory.
    This is used to determine if we need a full workflow rerun
    (to rebuild the kernel) vs just rerunning failed jobs.
    """
    try:
        files = pr.get_files()
        for f in files:
            if f.filename.startswith("python/sglang/kernels/aot/"):
                return True
        return False
    except Exception as e:
        print(f"Warning: Could not check PR files for sgl-kernel changes: {e}")
        # Default to False to avoid unnecessary full reruns
        return False


def handle_tag_run_ci(
    gh_repo, pr, comment, user_perms, react_on_success=True, tag_extra=False
):
    """
    Handles the /tag-run-ci-label command.

    When tag_extra is True (triggered by the `extra` argument), also adds the
    `run-ci-extra` label. pr-test-extra.yml gates on BOTH `run-ci` and
    `run-ci-extra`, so both must be present for the extra workflow to run —
    we always add `run-ci` alongside `run-ci-extra`. Reuses the same
    `can_tag_run_ci_label` permission.

    How fresh runs get dispatched: pr-test.yml and pr-test-extra.yml both
    include `labeled` in `on.pull_request.types`, so adding a label fires a
    new `pull_request.labeled` event with the up-to-date label set in its
    payload, which spawns a fresh workflow run that satisfies the
    `check-changes.if` gate. Note that this is the ONLY way to "un-skip" a
    label-gated run — `run.rerun()` on a previously-skipped pull_request run
    reuses the original event payload (frozen labels), so it would skip
    again. handle_rerun_failed_ci can't recover label-skipped runs; the
    labeled event is the recovery mechanism.

    Returns True if action was taken, False otherwise.
    """
    if not user_perms.get("can_tag_run_ci_label", False):
        print("Permission denied: can_tag_run_ci_label is false.")
        return False

    labels = ["run-ci"]
    if tag_extra:
        labels.append("run-ci-extra")
    print(f"Permission granted. Adding labels: {labels}.")
    for label in labels:
        pr.add_to_labels(label)

    if react_on_success:
        comment.create_reaction("+1")
        print("Labels added and comment reacted.")
    else:
        print("Labels added (reaction suppressed).")

    return True


def handle_rerun_failed_ci(gh_repo, pr, comment, user_perms, react_on_success=True):
    """
    Handles the /rerun-failed-ci command.
    Reruns workflows with 'failure' or 'skipped' conclusions.
    Returns True if action was taken, False otherwise.
    """
    if not user_perms.get("can_rerun_failed_ci", False):
        print("Permission denied: can_rerun_failed_ci is false.")
        return False

    print("Permission granted. Triggering rerun of failed or skipped workflows.")

    # Check if PR has sgl-kernel changes - if so, we may need full reruns
    # to ensure sgl-kernel-build-wheels runs and produces fresh artifacts.
    # However, if the wheel already built successfully for this commit,
    # we can just rerun failed jobs — the artifact is already there.
    sgl_kernel_changes = has_sgl_kernel_changes(pr)
    if sgl_kernel_changes:
        print("PR has sgl-kernel changes - checking if kernel wheel already built")

    # Get the SHA of the latest commit in the PR
    head_sha = pr.head.sha
    print(f"Checking workflows for commit: {head_sha}")

    # If PR has sgl-kernel changes, check whether ALL wheel builds already
    # succeeded for this commit (CUDA + ARM). If so, we can use
    # rerun_failed_jobs and avoid retriggering all tests. If any wheel
    # build is pending/failed, a dependent job could fail for missing
    # artifacts, so fall back to full rerun.
    # Check-runs display names: "Build Wheel (<python>, <cuda>)" (CUDA) and
    # "Build Wheel Arm (<python>, <cuda>)" (ARM). The YAML job ids
    # sgl-kernel-build-wheels{,-arm} are NOT what the check-runs API
    # returns — it returns the job's `name:` field.
    kernel_wheel_built = False
    if sgl_kernel_changes:
        try:
            wheel_builds = [
                cr
                for cr in gh_repo.get_commit(head_sha).get_check_runs()
                if cr.name.startswith("Build Wheel")
            ]
            kernel_wheel_built = bool(wheel_builds) and all(
                cr.conclusion == "success" for cr in wheel_builds
            )
            print(
                f"All {len(wheel_builds)} kernel wheel build(s) passed - using rerun_failed_jobs"
                if kernel_wheel_built
                else f"Kernel wheel not fully built "
                f"({sum(1 for c in wheel_builds if c.conclusion == 'success')}"
                f"/{len(wheel_builds)} success) - will use full rerun"
            )
        except Exception as e:
            print(
                f"Failed to check kernel wheel status: {e} - falling back to full rerun"
            )

    # Rerun workflows with conclusion=failure or conclusion=skipped.
    #
    # - failure: use rerun_failed_jobs() which reruns failed jobs *and their
    #   dependent jobs* (GitHub API). Fast-fail cascades call
    #   core.setFailed(...) so their conclusion is "failure" and are covered.
    # - skipped: the entire run was skipped (no jobs ran), so there are no
    #   failed jobs for rerun_failed_jobs() to target. Use run.rerun().
    #
    #   Caveat: GitHub's `run.rerun()` reuses the original event payload, so
    #   reruns of `pull_request`-event runs that were skipped because their
    #   `if` evaluated to false (e.g. missing label) will skip again — the
    #   label set in the frozen payload doesn't update. To un-skip a
    #   label-gated workflow, add the missing label (the `labeled` event
    #   dispatches a fresh run with the current label set); this function
    #   cannot recover those by rerun alone.
    # - kernel wheel escape: if the PR touches sgl-kernel and not all wheel
    #   builds are success yet, full-rerun failure runs too — Build Wheel
    #   lives in pr-test-sgl-kernel.yml, consumers in pr-test.yml, and
    #   rerun_failed_jobs() is scoped to a single workflow run.
    runs = gh_repo.get_workflow_runs(head_sha=head_sha)

    rerun_count = 0
    for run in runs:
        if run.status != "completed":
            continue
        if run.conclusion not in ("failure", "skipped"):
            continue

        print(f"Processing {run.conclusion} workflow: {run.name} (ID: {run.id})")
        try:
            if run.conclusion == "skipped" or (
                sgl_kernel_changes and not kernel_wheel_built
            ):
                print("  Full rerun")
                run.rerun()
            else:
                print("  rerun_failed_jobs")
                run.rerun_failed_jobs()
            rerun_count += 1
        except Exception as e:
            print(f"Failed to rerun workflow {run.id}: {e}")

    if rerun_count > 0:
        print(f"Triggered rerun for {rerun_count} workflows.")
        if react_on_success:
            comment.create_reaction("+1")
        return True
    else:
        print("No failed or skipped workflows found to rerun.")
        return False


MULTIMODAL_TEST_DIR = "python/sglang/multimodal_gen/test"

MULTIMODAL_PATH_TO_RUNNER = {
    "2_gpu": "2-gpu-h100",
    "2-gpu": "2-gpu-h100",
    "b200": _B200_DEFAULT_RUNNER,
}
MULTIMODAL_DEFAULT_RUNNER = "1-gpu-h100"


def _known_test_groups():
    groups = []
    for group_dir in glob.glob("test/registered/*"):
        if os.path.isdir(group_dir):
            groups.append(os.path.basename(group_dir))
    return sorted(groups)


def resolve_test_group_specs(group_name):
    """
    Resolve a test group name into /rerun-test specs.

    A group maps to a directory under test/registered/. For example,
    "hicache" maps to all test_*.py files under test/registered/hicache/.

    Returns (test_specs, error_message). On success error_message is None.
    """
    group_name = group_name.strip().strip("/")
    if (
        not group_name
        or group_name.startswith(".")
        or "/." in group_name
        or ".." in group_name.split("/")
    ):
        return [], f"Invalid test group `{group_name}`."

    group_dir = os.path.join("test", "registered", group_name)
    if not os.path.isdir(group_dir):
        known = ", ".join(f"`{g}`" for g in _known_test_groups())
        return (
            [],
            f"Unknown test group `{group_name}`.\n\nKnown groups: {known}",
        )

    test_files = sorted(
        glob.glob(os.path.join(group_dir, "**", "test_*.py"), recursive=True)
    )
    if not test_files:
        return [], f"No registered test files found in `{group_dir}`."

    return [os.path.relpath(path, "test") for path in test_files], None


# A spec is treated as a wildcard pattern (expanding to many files) when its
# file part contains a glob metacharacter. Plain specs keep the existing
# single-file resolution, which requires a unique match.
_GLOB_METACHARS = ("*", "?", "[")


def _is_glob_pattern(file_part):
    return any(ch in file_part for ch in _GLOB_METACHARS)


def expand_glob_spec(file_part):
    """
    Expand a wildcard file_part into matching test files (repo-relative paths).

    Globs are matched against the same locations resolve_test_file() searches
    — test/registered/ and the multimodal_gen test dir — so e.g.
    `test_*backend*.py` reruns every backend test without hand-enumerating
    each file. Two constraints keep a broad pattern from pulling in non-tests:
    a match must live under a known test root and be named `test_*.py`.

    glob's `*` matches path separators only via `**`, so a bare pattern is
    searched recursively under each root; a path-ful pattern is anchored.

    Returns (sorted_repo_relative_paths, error). On success error is None.
    """
    pat = file_part
    if pat.startswith("test/"):
        pat = pat[len("test/") :]

    matches = set()
    if "/" in pat:
        # Path-ful pattern. Glob from the repo root (handles fully qualified
        # multimodal paths like python/sglang/multimodal_gen/test/**/test_*.py)
        # and under test/ (handles test/-relative patterns like
        # registered/attention/test_*.py).
        for base in (".", "test"):
            matches.update(glob.glob(os.path.join(base, pat), recursive=True))
    else:
        # Bare pattern: search recursively under each known test root.
        for root in ("test/registered", MULTIMODAL_TEST_DIR):
            matches.update(glob.glob(os.path.join(root, "**", pat), recursive=True))

    # If the literal glob matched any directories — e.g. `unittest/*` matching
    # subdirs `dense/`, `dsa/`, … — descend recursively into them for
    # test_*.py. glob's `*` doesn't cross `/`, so without this a tree of
    # subdirs containing tests would return only the top-level entries (and
    # then get filtered out as non-test files).
    expanded = set()
    for p in matches:
        if os.path.isdir(p):
            expanded.update(
                glob.glob(os.path.join(p, "**", "test_*.py"), recursive=True)
            )
        else:
            expanded.add(p)
    matches = expanded

    def _under_test_root(path):
        return path.startswith("test/registered/") or path.startswith(
            MULTIMODAL_TEST_DIR + "/"
        )

    files = sorted(
        {
            os.path.normpath(p)
            for p in matches
            if os.path.isfile(p)
            and os.path.basename(p).startswith("test_")
            and p.endswith(".py")
            and _under_test_root(os.path.normpath(p))
        }
    )
    if not files:
        return [], (
            f"No test files matched wildcard `{file_part}` under "
            f"`test/registered/` or `{MULTIMODAL_TEST_DIR}/` "
            f"(patterns only match files named `test_*.py`)."
        )
    return files, None


def resolve_test_file(file_part):
    """
    Resolve a user-provided file path to a path relative to test/ or full path for multimodal.

    Supports:
    - Full path: test/registered/core/test_srt_endpoint.py
    - Relative to test/: registered/core/test_srt_endpoint.py
    - Bare filename: test_srt_endpoint.py (glob-matched, must be unique)
    - Multimodal paths: python/sglang/multimodal_gen/test/server/test_server_a.py

    Returns (resolved_path, is_multimodal, error_message). On success error_message is None.
    """
    # Check if it's explicitly a multimodal path
    multimodal_prefixes = [
        "python/sglang/multimodal_gen/test/",
        "sglang/multimodal_gen/test/",
        "multimodal_gen/test/",
    ]
    for prefix in multimodal_prefixes:
        if file_part.startswith(prefix):
            full_path = (
                file_part
                if file_part.startswith("python/")
                else f"python/sglang/multimodal_gen/test/{file_part[len(prefix):]}"
            )
            if not os.path.isfile(full_path):
                return None, False, f"File not found: `{full_path}`"
            return full_path, True, None

    # Existing logic for test/registered/ paths
    if file_part.startswith("test/"):
        file_part = file_part[len("test/") :]

    if "/" not in file_part:
        # Try test/registered/ first
        matches = glob.glob(f"test/registered/**/{file_part}", recursive=True)

        # Try multimodal test directory
        mm_matches = glob.glob(f"{MULTIMODAL_TEST_DIR}/**/{file_part}", recursive=True)
        # Filter to only test files
        mm_matches = [m for m in mm_matches if os.path.basename(m).startswith("test_")]

        if len(matches) == 1 and len(mm_matches) == 0:
            return matches[0][len("test/") :], False, None
        if len(matches) == 0 and len(mm_matches) == 1:
            return mm_matches[0], True, None

        all_matches = matches + mm_matches
        if len(all_matches) == 0:
            return (
                None,
                False,
                f"No test file found matching `{file_part}` under `test/registered/` or `{MULTIMODAL_TEST_DIR}/`.",
            )
        if len(all_matches) > 1:
            match_list = "\n".join(f"- `{m}`" for m in sorted(all_matches))
            return (
                None,
                False,
                (
                    f"Ambiguous filename `{file_part}` — matched {len(all_matches)} files:\n\n"
                    f"{match_list}\n\n"
                    f"Please provide the full path, e.g. `/rerun-test {all_matches[0]}`"
                ),
            )
        # Shouldn't reach here, but handle gracefully
        if mm_matches:
            return mm_matches[0], True, None
        return matches[0][len("test/") :], False, None

    # Path with directory - check test/ location
    full_path = f"test/{file_part}"
    if os.path.isfile(full_path):
        return file_part, False, None

    return None, False, f"File not found: `{full_path}`"


def detect_multimodal_suite(file_path):
    """
    Determine runner for a multimodal gen test file based on its path.

    Returns (runner_label, error_message).
    """
    # Check path components and basename for GPU count hints
    for pattern, runner in MULTIMODAL_PATH_TO_RUNNER.items():
        if pattern in file_path:
            return runner, None
    return MULTIMODAL_DEFAULT_RUNNER, None


def _extract_runner_configs(content):
    """Pull `(runner_config, args_str)` from EVERY `register_cuda_ci(...)` call.

    A test file can register itself on multiple pools (e.g. both
    `4-gpu-b200` and `1-gpu-large`). The earlier `re.search` variant
    returned only the first match, so /rerun-test silently dropped every
    registration after the first — multi-pool files only ever ran on the
    pool listed first. `re.finditer` is what makes the fan-out happen.
    """
    out = []
    for args in re.finditer(
        r"^[^#\n]*register_cuda_ci\s*\(([^)]*)\)", content, re.MULTILINE
    ):
        m = re.search(r'runner_config\s*=\s*["\']([^"\']+)["\']', args.group(1))
        if m:
            out.append((m.group(1), args.group(1)))
    return out


def _extract_legacy_suites(content):
    """Pull every legacy single-string `suite=` from `register_cuda_ci(...)` calls.

    Mirrors _extract_runner_configs for the legacy nightly/weekly shape: a file
    may register on multiple pools, so collect all of them rather than the first.
    """
    out = []
    for args in re.finditer(
        r"^[^#\n]*register_cuda_ci\s*\(([^)]*)\)", content, re.MULTILINE
    ):
        m = re.search(r'suite\s*=\s*["\']([^"\']+)["\']', args.group(1))
        if m:
            out.append(m.group(1))
    return out


# Legacy nightly/weekly CUDA suites register with a single-string `suite=`
# instead of `runner_config=`, so they carry no runner metadata of their own.
# Map each to the runner_config in scripts/ci/runner_configs.yml whose hardware
# matches the runner the nightly/weekly pipeline actually uses (see
# .github/workflows/{nightly,weekly}-test-nvidia.yml), so /rerun-test can still
# dispatch a single nightly/weekly test. The runner label, install script,
# timeout, grace_blackwell, and rdma_devices are then resolved from
# runner_configs.yml as usual, keeping that file the single source of truth for
# runner details.
#
# Suites on hardware with no matching runner_config (e.g. nightly-4-gpu-gb300)
# and non-CUDA suites (npu/amd) are intentionally absent and stay
# non-dispatchable until a matching runner_config exists.
_LEGACY_SUITE_TO_RUNNER_CONFIG = {
    "nightly-1-gpu": "1-gpu-large",
    "nightly-kernel-1-gpu": "1-gpu-large",
    "nightly-eval-text-2-gpu": "2-gpu-large",
    "nightly-perf-text-2-gpu": "2-gpu-large",
    "nightly-eval-vlm-2-gpu": "2-gpu-large",
    "nightly-perf-vlm-2-gpu": "2-gpu-large",
    "nightly-4-gpu": "4-gpu-h100",
    "nightly-4-gpu-b200": "4-gpu-b200",
    "nightly-8-gpu-common": ["8-gpu-h200", "8-gpu-b200"],
    "nightly-8-gpu-h200": "8-gpu-h200",
    "nightly-kernel-8-gpu-h200": "8-gpu-h200",
    "nightly-precision-8-gpu-h200": "8-gpu-h200",
    "nightly-8-gpu-h20": "8-gpu-h20",
    "nightly-8-gpu-b200": "8-gpu-b200",
    "weekly-8-gpu-h200": "8-gpu-h200",
}


def _dispatch_err(suite, msg):
    """Build a detect_suite error result for the given suite."""
    return {
        "suite": suite,
        "runner_label": None,
        "install_script": "",
        "install_timeout": "",
        "grace_blackwell": "0",
        "rdma_devices": "",
        "is_cpu": False,
        "error": msg,
    }


def _resolve_runner_config(rc, full_path, suite):
    """Resolve a runner_config key into a detect_suite dispatch dict.

    Returns the dispatch dict on success, or an error dict on failure.
    """
    configs = _runner_configs.load()
    cfg = configs.get(rc)
    if cfg is None:
        known = ", ".join(f"`{k}`" for k in sorted(configs))
        return _dispatch_err(
            suite,
            f"Unknown runner_config `{rc}` in `{full_path}` "
            f"— not in scripts/ci/runner_configs.yml.\n\n"
            f"Known runner_configs: {known}",
        )
    install_script = cfg["install"]
    if not _ALLOWED_INSTALL_SCRIPT.match(install_script):
        return _dispatch_err(
            suite,
            f"Disallowed `install` value `{install_script}` for runner_config "
            f"`{rc}` in scripts/ci/runner_configs.yml. The slash handler "
            f"passes this string verbatim into a shell step, so it must "
            f"match `scripts/ci/cuda/*.sh`.",
        )
    runs_on = cfg.get("runs_on")
    # Resolve $b200_runner sentinel: rerun-test never builds sgl-kernel,
    # so always pick the non-kernel b200 pool.
    if runs_on == "$b200_runner":
        runs_on = _B200_DEFAULT_RUNNER
    return {
        "suite": suite,
        "runner_label": runs_on,
        "install_script": install_script,
        "install_timeout": str(cfg["install_timeout"]),
        "grace_blackwell": str(cfg.get("grace_blackwell", "0")),
        "rdma_devices": cfg.get("rdma_devices", ""),
        "is_cpu": False,
        "error": None,
    }


def detect_suite(file_path_from_test):
    """
    Read a test file and extract dispatch info from register_cuda_ci or
    register_cpu_ci.

    A CUDA file can carry multiple `register_cuda_ci(...)` calls — one per
    pool it should run on — so this returns a *list* of dispatch dicts, one
    per registration. Runner label, install script, timeout, grace_blackwell,
    and rdma_devices are all resolved from scripts/ci/runner_configs.yml — the
    same single source of truth that drives the main PR test pipeline.

    Legacy nightly/weekly CUDA suites (single-string `suite=`) are dispatchable
    too: each suite name is mapped to the matching runner_config via
    _LEGACY_SUITE_TO_RUNNER_CONFIG, then resolved the same way.

    CPU files yield a single-element list. A file with no recognised (or no
    dispatchable) registration yields a one-element list whose dict has an
    `error` set.

    Each dict has keys: suite, runner_label, install_script,
    install_timeout, grace_blackwell, rdma_devices, is_cpu, error.
    """
    full_path = f"test/{file_path_from_test}"
    with open(full_path, "r") as f:
        content = f.read()

    cuda_calls = _extract_runner_configs(content)
    if cuda_calls:
        results = []
        for rc, args_str in cuda_calls:
            stage_m = re.search(r'stage\s*=\s*["\']([^"\']+)["\']', args_str)
            suite = f"{stage_m.group(1)}-test-{rc}" if stage_m else rc
            results.append(_resolve_runner_config(rc, full_path, suite))
        return results

    # Legacy nightly/weekly CUDA suites: single-string `suite=`, no
    # runner_config. Map each mappable suite to its runner_config and resolve.
    legacy_suites = _extract_legacy_suites(content)
    mappable = [s for s in legacy_suites if s in _LEGACY_SUITE_TO_RUNNER_CONFIG]
    if mappable:
        results = []
        for s in mappable:
            rcs = _LEGACY_SUITE_TO_RUNNER_CONFIG[s]
            if isinstance(rcs, str):
                rcs = [rcs]
            for rc in rcs:
                results.append(_resolve_runner_config(rc, full_path, s))
        return results

    if re.search(r"^[^#\n]*register_cpu_ci\s*\(", content, re.MULTILINE):
        return [
            {
                "suite": "cpu",
                "runner_label": "ubuntu-latest",
                "install_script": "",
                "install_timeout": "",
                "grace_blackwell": "0",
                "rdma_devices": "",
                "is_cpu": True,
                "error": None,
            }
        ]

    if legacy_suites:
        suite = legacy_suites[0]
        return [
            _dispatch_err(
                suite,
                f"Suite `{suite}` in `{full_path}` is not dispatchable via "
                f"/rerun-test. It has no entry in _LEGACY_SUITE_TO_RUNNER_CONFIG "
                f"— either it is a non-CUDA suite (npu/amd) or it runs on "
                f"hardware with no matching runner_config in "
                f"scripts/ci/runner_configs.yml.",
            )
        ]

    return [
        _dispatch_err(
            None,
            f"No `register_cuda_ci(runner_config=...)` or `register_cpu_ci()` "
            f"found in `{full_path}`. This file may not be a registered CI test.",
        )
    ]


def _resolve_test_spec(test_spec):
    """
    Resolve a single test spec into one or more dispatch entries.

    A file registered on N pools (multiple `register_cuda_ci(...)` calls)
    yields N entries — handle_rerun_test's grouping then sends one workflow
    per pool. Multimodal and CPU files yield a single entry. Resolution
    errors yield a one-element list with {"spec","error"}.
    """
    if "::" in test_spec:
        file_part, test_selector = test_spec.split("::", 1)
    else:
        file_part = test_spec
        test_selector = None

    file_part = file_part.strip()
    if test_selector:
        test_selector = test_selector.strip()

    resolved_path, is_multimodal, err = resolve_test_file(file_part)
    if err:
        return [{"spec": test_spec, "error": err}]

    if is_multimodal:
        runner_label, err = detect_multimodal_suite(resolved_path)
        if err:
            return [{"spec": test_spec, "error": err}]

        # For multimodal pytest tests, use :: separator for test selection
        test_command = resolved_path
        if test_selector:
            test_command = f"{resolved_path}::{test_selector}"

        print(
            f"Resolved (multimodal_gen): file={resolved_path}, "
            f"selector={test_selector}, runner={runner_label}, "
            f"command='{test_command}'"
        )
        return [
            {
                "spec": test_spec,
                "test_command": test_command,
                "mode": "multimodal_gen",
                "runs_on": runner_label,
                "install_script": "",
                "install_timeout": "",
                "grace_blackwell": "0",
                "rdma_devices": "",
                "error": None,
            }
        ]

    test_command = resolved_path
    if test_selector:
        test_command = f"{resolved_path} {test_selector}"

    out = []
    for info in detect_suite(resolved_path):
        if info["error"]:
            out.append({"spec": test_spec, "error": info["error"]})
            continue
        mode = "cpu" if info["is_cpu"] else "cuda"
        print(
            f"Resolved: file={resolved_path}, selector={test_selector}, "
            f"suite={info['suite']}, mode={mode}, runs_on={info['runner_label']}, "
            f"install={info['install_script']}, grace_blackwell={info['grace_blackwell']}, "
            f"rdma={info['rdma_devices']}, "
            f"command='{test_command}'"
        )
        out.append(
            {
                "spec": test_spec,
                "test_command": test_command,
                "mode": mode,
                "runs_on": info["runner_label"],
                "install_script": info["install_script"],
                "install_timeout": info["install_timeout"],
                "grace_blackwell": info["grace_blackwell"],
                "rdma_devices": info["rdma_devices"],
                "error": None,
            }
        )
    return out


def _dispatch_batch(gh_repo, pr, batch, token, reply_comment_id="", reply_marker=""):
    """
    Dispatch a single workflow run for a batch of resolved test specs that
    share the same dispatch shape (mode + runs_on + install_script +
    install_timeout + grace_blackwell + rdma_devices).

    Returns a dict with keys: specs, success, test_commands, runner_label, run_url, error.
    """
    test_commands = [r["test_command"] for r in batch]
    mode = batch[0]["mode"]
    runs_on = batch[0]["runs_on"]
    install_script = batch[0]["install_script"]
    install_timeout = batch[0]["install_timeout"]
    grace_blackwell = batch[0]["grace_blackwell"]
    rdma_devices = batch[0]["rdma_devices"]

    # Join multiple commands with newlines for the workflow to iterate over
    combined_command = "\n".join(test_commands)

    try:
        workflow_name = "Rerun Test"
        workflows = gh_repo.get_workflows()
        target_workflow = None
        for wf in workflows:
            if wf.name == workflow_name:
                target_workflow = wf
                break

        if not target_workflow:
            return {
                "specs": [r["spec"] for r in batch],
                "success": False,
                "error": f"{workflow_name} workflow not found",
            }

        is_fork = (
            pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login
        )

        pr_head_sha = None
        inputs = {
            "mode": mode,
            "test_command": combined_command,
            "runs_on": runs_on or "",
            "install_script": install_script,
            "install_timeout": install_timeout or "20",
            "grace_blackwell": grace_blackwell or "0",
            "rdma_devices": rdma_devices,
            "reply_comment_id": str(reply_comment_id) if reply_comment_id else "",
            "reply_marker": reply_marker,
        }
        if is_fork:
            ref = "main"
            pr_head_sha = pr.head.sha
            inputs["pr_head_sha"] = pr_head_sha
        else:
            ref = pr.head.ref

        dispatch_time = time.time()

        dispatch_url = f"https://api.github.com/repos/{gh_repo.full_name}/actions/workflows/{target_workflow.id}/dispatches"
        dispatch_resp = requests.post(
            dispatch_url,
            json={"ref": ref, "inputs": inputs},
            headers={
                "Authorization": f"Bearer {token}",
                "Accept": "application/vnd.github+json",
            },
        )
        success = dispatch_resp.status_code in (200, 204)
        if not success:
            print(f"Dispatch failed: {dispatch_resp.status_code} {dispatch_resp.text}")
            return {
                "specs": [r["spec"] for r in batch],
                "success": False,
                "error": f"Dispatch failed: {dispatch_resp.status_code}",
            }

        print(f"Successfully triggered rerun-test: {combined_command}")

        run_url = find_workflow_run_url(
            gh_repo,
            target_workflow.id,
            ref,
            "rerun-test",
            token,
            dispatch_time,
            pr_head_sha=pr_head_sha,
            max_wait=30,
            test_command=combined_command,
        )
        return {
            "specs": [r["spec"] for r in batch],
            "success": True,
            "test_commands": test_commands,
            "mode": mode,
            "runs_on": runs_on,
            "run_url": run_url,
            "reply_marker": reply_marker,
        }

    except Exception as e:
        print(f"Error triggering rerun-test for batch: {e}")
        return {
            "specs": [r["spec"] for r in batch],
            "success": False,
            "error": str(e),
        }


def _check_rerun_test_permissions(gh_repo, pr, comment, user_perms, command_name):
    """
    Check permissions shared by /rerun-test and /rerun-group.
    """
    # SECURITY: These commands check out and execute code from the PR branch on
    # self-hosted GPU runners, so fork PRs require a trusted collaborator.
    is_fork = pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login
    if is_fork:
        commenter = comment.user.login
        perm = gh_repo.get_collaborator_permission(commenter)
        if perm not in ("admin", "write"):
            print(f"Permission denied: /{command_name} on fork PR by {commenter}.")
            comment.create_reaction("confused")
            pr.create_issue_comment(
                f"⛔ `/{command_name}` is not available for fork PRs unless the commenter "
                "has write permission on the repo.\n\n"
                "Please ask a maintainer to run this command, or use the normal CI flow."
            )
            return False
        print(f"Fork PR, but commenter {commenter} has write+ permission. Proceeding.")

    if not (
        user_perms.get("can_rerun_test", False)
        or user_perms.get("can_rerun_stage", False)
    ):
        print("Permission denied: neither can_rerun_test nor can_rerun_stage is true.")
        return False

    return True


def handle_rerun_test(
    gh_repo,
    pr,
    comment,
    user_perms,
    test_specs,
    token,
    skip_permission_check=False,
    command_label=None,
):
    """
    Handles the /rerun-test command. Resolves all test specs, groups them by
    dispatch shape (mode + runs_on + install_script + install_timeout +
    grace_blackwell + rdma_devices), and dispatches one workflow per group.
    """
    if not skip_permission_check and not _check_rerun_test_permissions(
        gh_repo, pr, comment, user_perms, "rerun-test"
    ):
        return False

    if not test_specs:
        comment.create_reaction("confused")
        pr.create_issue_comment(
            "⛔ Please specify a test: `/rerun-test <file>::<TestClass.test_method>`\n\n"
            "Examples:\n"
            "- `/rerun-test test/registered/core/test_srt_endpoint.py::TestSRTEndpoint.test_simple_decode`\n"
            "- `/rerun-test registered/core/test_srt_endpoint.py::TestSRTEndpoint`\n"
            "- `/rerun-test test_srt_endpoint.py`\n"
            "- `/rerun-test test_a.py test_b.py test_c.py` (multiple tests)\n"
            "- `/rerun-test test_*backend*.py` (wildcard — reruns every matching "
            "file; wrap the pattern in backticks so GitHub keeps the `*` literal)"
        )
        return False

    allowed, gate_msg = _check_rebase_gate(gh_repo, pr, token)
    if not allowed:
        comment.create_reaction("confused")
        pr.create_issue_comment(gate_msg)
        return False

    # Phase 0: Expand wildcard specs into concrete test files. A spec whose
    # file part contains a glob metacharacter (* ? [) expands to every
    # matching file; plain specs pass through to single-file resolution.
    resolve_failures = []
    seen_failures = set()

    def _record_failure(spec, error):
        # De-dupe failures by canonical path so the same un-dispatchable file
        # reported two ways (explicit + glob, or two globs) yields one line.
        key = spec.strip().strip("\"'`")
        if key.startswith("test/"):
            key = key[len("test/") :]
        if key not in seen_failures:
            seen_failures.add(key)
            resolve_failures.append({"spec": spec, "error": error})

    expanded_specs = []
    for spec in test_specs:
        # Quotes/backticks are never meaningful here (the command isn't
        # shell-parsed), so strip them. Backtick-wrapping is in fact the
        # recommended way to write a glob: plain `*backend*` renders as italics
        # in a GitHub comment, but `` `test_*backend*.py` `` stays literal — and
        # either way the handler reads the raw body, so the `*` survives.
        file_part = spec.split("::", 1)[0].strip().strip("\"'`")
        if not _is_glob_pattern(file_part):
            expanded_specs.append(spec)
            continue
        if "::" in spec:
            _record_failure(
                spec,
                "Wildcard patterns can't be combined with a `::test` "
                "selector — drop the `::...` to rerun whole files.",
            )
            continue
        matched, err = expand_glob_spec(file_part)
        if err:
            _record_failure(spec, err)
            continue
        expanded_specs.extend(matched)

    # Phase 1: Resolve all specs, de-duping by the *resolved* identity. A glob
    # expands to `test/`-prefixed paths while an explicit spec keeps the form
    # the user typed, so the same file requested both ways resolves to two
    # different raw strings — keying de-dup on the resolved identity is what
    # collapses them into a single dispatch instead of running twice.
    # `runs_on` is included in the key so a file registered on N pools fans
    # out into one dispatch per pool, instead of all-but-the-first collapsing.
    resolved = []
    seen_commands = set()
    for spec in expanded_specs:
        for r in _resolve_test_spec(spec):
            if r.get("error"):
                _record_failure(r["spec"], r["error"])
                continue
            key = (r["mode"], r["runs_on"], r["test_command"])
            if key in seen_commands:
                continue
            seen_commands.add(key)
            resolved.append(r)

    # Phase 2: Group by dispatch shape.
    groups = {}
    for r in resolved:
        key = (
            r["mode"],
            r["runs_on"],
            r["install_script"],
            r["install_timeout"],
            r["grace_blackwell"],
            r["rdma_devices"],
        )
        groups.setdefault(key, []).append(r)

    # Phase 3a: Create placeholder reply comment so we have its ID before
    # dispatching workflows. This lets each dispatched run write its
    # success/failure result back to the right line in this comment.
    dispatching = f"`{command_label}`" if command_label else "rerun-test workflow(s)"
    reply_comment = pr.create_issue_comment(f"🚀 Dispatching {dispatching}...")

    # Phase 3b: Dispatch one workflow per group, with a unique per-batch
    # marker each. The marker is an HTML comment that the writeback step
    # uses to locate the line and replace 🚀 with ✅/❌.
    dispatch_results = []
    for idx, batch in enumerate(groups.values()):
        marker = f"<!--rrt:{idx}-->"
        dispatch_results.append(
            _dispatch_batch(
                gh_repo,
                pr,
                batch,
                token,
                reply_comment_id=reply_comment.id,
                reply_marker=marker,
            )
        )

    # Build consolidated comment body (markers placed at line ends so the
    # writeback step can locate and update each line).
    lines = []
    for dr in dispatch_results:
        if dr["success"]:
            if dr["mode"] == "multimodal_gen":
                cmds = "\n".join(
                    f"python3 -m pytest {cmd} -x" for cmd in dr["test_commands"]
                )
            else:
                cmds = "\n".join(
                    f"cd test/ && python3 {cmd}" for cmd in dr["test_commands"]
                )
            marker = dr.get("reply_marker", "")
            label = dr["runs_on"] or dr["mode"]
            if dr.get("run_url"):
                lines.append(
                    f"🚀 `{label}` ({len(dr['test_commands'])} test{'s' if len(dr['test_commands']) > 1 else ''}): "
                    f"⏳ [View workflow run]({dr['run_url']}) {marker}\n"
                    f"```\n{cmds}\n```"
                )
            else:
                lines.append(
                    f"🚀 `{label}` ({len(dr['test_commands'])} test{'s' if len(dr['test_commands']) > 1 else ''}): ⏳ {marker}\n"
                    f"```\n{cmds}\n```\n"
                    f"⚠️ Could not retrieve workflow run URL. "
                    f"Check the [Actions tab](https://github.com/{gh_repo.full_name}/actions) for progress."
                )
        else:
            specs_str = ", ".join(f"`{s}`" for s in dr["specs"])
            lines.append(f"⛔ {specs_str}: {dr['error']}")

    for r in resolve_failures:
        lines.append(f"⛔ `{r['spec']}`: {r['error']}")

    body = "\n\n".join(lines)
    # Echo the originating command so each reply is self-identifying when
    # several /rerun-test commands are in flight at once. Backtick-wrapping
    # also keeps any `*` in the pattern from rendering as italics.
    if command_label:
        body = f"Results for `{command_label}`:\n\n{body}"

    successes = [dr for dr in dispatch_results if dr["success"]]
    if successes:
        comment.create_reaction("+1")
    if not successes and (resolve_failures or dispatch_results):
        comment.create_reaction("confused")

    reply_comment.edit(body)
    return len(successes) > 0


def handle_rerun_group(
    gh_repo, pr, comment, user_perms, group_names, token, command_label=None
):
    """
    Handles the /rerun-group command. Expands one or more registered test
    groups into test file specs, then reuses /rerun-test dispatch behavior.
    """
    if not _check_rerun_test_permissions(
        gh_repo, pr, comment, user_perms, "rerun-group"
    ):
        return False

    if not group_names:
        comment.create_reaction("confused")
        pr.create_issue_comment(
            "⛔ Please specify a test group: `/rerun-group <group>`\n\n"
            "Example:\n"
            "- `/rerun-group hicache`"
        )
        return False

    test_specs = []
    failures = []
    seen = set()
    for group_name in group_names:
        specs, err = resolve_test_group_specs(group_name)
        if err:
            failures.append((group_name, err))
            continue

        for spec in specs:
            if spec not in seen:
                test_specs.append(spec)
                seen.add(spec)

    if failures:
        comment.create_reaction("confused")
        lines = [f"⛔ `{group}`: {err}" for group, err in failures]
        pr.create_issue_comment("\n\n".join(lines))
        return False

    return handle_rerun_test(
        gh_repo,
        pr,
        comment,
        user_perms,
        test_specs,
        token,
        skip_permission_check=True,
        command_label=command_label,
    )


def main():
    # 1. Load Environment Variables
    token = get_env_var("GITHUB_TOKEN")
    repo_name = get_env_var("REPO_FULL_NAME")
    pr_number = int(get_env_var("PR_NUMBER"))
    comment_id = int(get_env_var("COMMENT_ID"))
    comment_body = _strip_format_chars(get_env_var("COMMENT_BODY")).strip()
    user_login = get_env_var("USER_LOGIN")

    # 2. Load Permissions (local file check first to avoid unnecessary API calls)
    user_perms = load_permissions(user_login)

    # 3. Initialize GitHub API with Auth
    auth = Auth.Token(token)
    g = Github(auth=auth)

    repo = g.get_repo(repo_name)
    pr = repo.get_pull(pr_number)
    comment = repo.get_issue(pr_number).get_comment(comment_id)

    # PR authors can always rerun failed CI and rerun individual UTs on their own PRs,
    # even if they are not listed in CI_PERMISSIONS.json.
    # Note: /tag-run-ci-label still requires CI_PERMISSIONS.json.
    # Note: /rerun-test is blocked entirely for fork PRs in handle_rerun_test() itself.
    if pr.user.login == user_login:
        if user_perms is None:
            print(
                f"User {user_login} is the PR author (not in CI_PERMISSIONS.json). "
                "Granting CI rerun permissions."
            )
            user_perms = {}
        else:
            print(
                f"User {user_login} is the PR author and has existing CI permissions."
            )
        user_perms["can_rerun_failed_ci"] = True
        user_perms["can_rerun_test"] = True

    if not user_perms:
        print(f"User {user_login} does not have any configured permissions. Exiting.")
        return

    # 4. Parse Command and Execute
    first_line = comment_body.split("\n")[0].strip()
    # `extra` argument opts in to also tagging `run-ci-extra`. Both
    # `/tag-run-ci-label extra` and `/tag-and-rerun-ci extra` share this
    # parser so the surface is symmetric.
    tokens = first_line.split()
    tag_extra = len(tokens) > 1 and "extra" in tokens[1:]

    if first_line.startswith("/tag-run-ci-label"):
        handle_tag_run_ci(repo, pr, comment, user_perms, tag_extra=tag_extra)

    elif first_line.startswith("/rerun-failed-ci"):
        handle_rerun_failed_ci(repo, pr, comment, user_perms)

    elif first_line.startswith("/tag-and-rerun-ci"):
        # Perform both actions, but suppress individual reactions
        print(f"Processing combined command: /tag-and-rerun-ci (tag_extra={tag_extra})")

        tagged = handle_tag_run_ci(
            repo, pr, comment, user_perms, react_on_success=False, tag_extra=tag_extra
        )

        # Wait for the label to propagate before triggering rerun
        if tagged:
            print("Waiting 5 seconds for label to propagate...")
            time.sleep(5)

        rerun = handle_rerun_failed_ci(
            repo, pr, comment, user_perms, react_on_success=False
        )

        # If at least one action was successful, add the reaction here
        if tagged or rerun:
            comment.create_reaction("+1")
            print("Combined command processed successfully; reaction added.")
        else:
            print("Combined command finished, but no actions were taken.")

    elif first_line.startswith("/rerun-stage"):
        print("/rerun-stage is deprecated; posting deprecation notice.")
        comment.create_reaction("-1")
        pr.create_issue_comment(
            "⚠️ **`/rerun-stage` has been deprecated.**\n\n"
            "Stage granularity is too coarse — a stage usually doesn't map to one "
            "feature, so rerunning a stage re-pays the cost of unrelated tests. "
            "If you don't know which exact test files to rerun, you shouldn't be "
            "using `/rerun-stage` or `/rerun-test` in the first place.\n\n"
            "**Use one of these instead:**\n"
            "- **Selective tests** (you know exactly which files to rerun):\n"
            "  ```\n"
            "  /rerun-test test_foo.py test_bar.py\n"
            "  ```\n"
            "- **Rerun only failed jobs**:\n"
            "  ```\n"
            "  /rerun-failed-ci\n"
            "  ```\n"
            "- **Full CI rerun** (with extra coverage): add the `run-ci` or "
            "`run-ci-extra` label and push a new commit (or use `/tag-and-rerun-ci`).\n\n"
            "**AMD CI**: stage-level dispatch is still available via "
            "Actions UI → *PR Test (AMD)* / *PR Test ROCm 7.2 (AMD)* → "
            "*Run workflow* → pick a stage from the dropdown."
        )

    elif first_line.startswith("/rerun-group"):
        group_names = first_line.split()[1:]
        handle_rerun_group(
            repo,
            pr,
            comment,
            user_perms,
            group_names or None,
            token,
            command_label=first_line,
        )

    elif first_line.startswith("/rerun-test"):
        test_specs = first_line.split()[1:]
        handle_rerun_test(
            repo,
            pr,
            comment,
            user_perms,
            test_specs or None,
            token,
            command_label=first_line,
        )

    else:
        print(f"Unknown or ignored command: {first_line}")


if __name__ == "__main__":
    main()
