# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
from collections.abc import Mapping, Sequence
from typing import Any

COMMAND_RUN_CI = "/ci run"
COMMAND_RETRY_FAILED = "/ci retry"
READY_LABELS = {"ready", "ready-run-all-tests"}
TRUSTED_PERMISSIONS = {"admin", "maintain", "write"}
ACTIVE_BUILD_STATES = {
    "blocked",
    "creating",
    "scheduled",
    "running",
    "failing",
    "canceling",
    "waiting",
    "waiting_failed",
}
RETRY_STATES = "failed,timed_out,expired"


class ApiError(RuntimeError):
    def __init__(self, status: int | None, message: str) -> None:
        super().__init__(message)
        self.status = status


class HttpTransport:
    def request(
        self,
        url: str,
        *,
        body: Mapping[str, Any] | None = None,
        headers: Mapping[str, str] | None = None,
        method: str = "GET",
    ) -> Any:
        data = None if body is None else json.dumps(body).encode()
        request = urllib.request.Request(
            url,
            data=data,
            headers=dict(headers or {}),
            method=method,
        )
        try:
            with urllib.request.urlopen(request, timeout=30) as response:
                response_body = response.read().decode()
        except urllib.error.HTTPError as error:
            response_body = error.read().decode()
            message = self._error_message(response_body, error.reason)
            raise ApiError(
                error.code,
                f"API returned {error.code}: {message}",
            ) from error
        except urllib.error.URLError as error:
            raise ApiError(None, f"API request failed: {error.reason}") from error

        if not response_body:
            return None
        try:
            return json.loads(response_body)
        except json.JSONDecodeError as error:
            raise ApiError(None, "API returned a non-JSON response.") from error

    @staticmethod
    def _error_message(response_body: str, fallback: str) -> str:
        try:
            parsed = json.loads(response_body)
        except json.JSONDecodeError:
            return fallback
        return str(parsed.get("message", fallback))


class GitHubClient:
    def __init__(
        self,
        token: str,
        repository: str,
        transport: HttpTransport | None = None,
    ) -> None:
        if not token:
            raise RuntimeError("GH_TOKEN is not set.")
        self.owner, self.repo = repository.split("/", maxsplit=1)
        self.transport = transport or HttpTransport()
        self.headers = {
            "Accept": "application/vnd.github+json",
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
            "User-Agent": "vllm-ci-command",
            "X-GitHub-Api-Version": "2022-11-28",
        }

    def _request(
        self,
        path: str,
        *,
        body: Mapping[str, Any] | None = None,
        method: str = "GET",
    ) -> Any:
        return self.transport.request(
            f"https://api.github.com{path}",
            body=body,
            headers=self.headers,
            method=method,
        )

    def _repo_path(self, suffix: str) -> str:
        owner = urllib.parse.quote(self.owner, safe="")
        repo = urllib.parse.quote(self.repo, safe="")
        return f"/repos/{owner}/{repo}{suffix}"

    def _paginate(self, path: str) -> list[dict[str, Any]]:
        results: list[dict[str, Any]] = []
        separator = "&" if "?" in path else "?"
        for page in range(1, 101):
            response = self._request(f"{path}{separator}per_page=100&page={page}")
            if not isinstance(response, list):
                raise ApiError(None, "GitHub API returned an invalid list response.")
            results.extend(response)
            if len(response) < 100:
                return results
        raise ApiError(None, "GitHub API pagination exceeded 10,000 results.")

    def get_pr(self, number: int) -> dict[str, Any]:
        return self._request(self._repo_path(f"/pulls/{number}"))

    def get_permission(self, actor: str) -> str:
        username = urllib.parse.quote(actor, safe="")
        try:
            response = self._request(
                self._repo_path(f"/collaborators/{username}/permission")
            )
        except ApiError as error:
            if error.status == 404:
                return "none"
            raise
        return str(response["permission"])

    def get_review_decision(self, number: int) -> str | None:
        query = """
          query($owner: String!, $repo: String!, $number: Int!) {
            repository(owner: $owner, name: $repo) {
              pullRequest(number: $number) {
                reviewDecision
              }
            }
          }
        """
        response = self._request(
            "/graphql",
            body={
                "query": query,
                "variables": {
                    "number": number,
                    "owner": self.owner,
                    "repo": self.repo,
                },
            },
            method="POST",
        )
        return response["data"]["repository"]["pullRequest"]["reviewDecision"]

    def list_reviews(self, number: int) -> list[dict[str, Any]]:
        return self._paginate(self._repo_path(f"/pulls/{number}/reviews"))

    def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
        return self._paginate(
            self._repo_path(f"/issues/comments/{comment_id}/reactions")
        )

    def add_reaction(self, comment_id: int, content: str) -> None:
        self._request(
            self._repo_path(f"/issues/comments/{comment_id}/reactions"),
            body={"content": content},
            method="POST",
        )

    def add_comment(self, issue_number: int, body: str) -> None:
        self._request(
            self._repo_path(f"/issues/{issue_number}/comments"),
            body={"body": body},
            method="POST",
        )


class BuildkiteClient:
    def __init__(
        self,
        token: str,
        organization: str,
        pipeline: str,
        transport: HttpTransport | None = None,
    ) -> None:
        self.token = token
        self.transport = transport or HttpTransport()
        organization = urllib.parse.quote(organization, safe="")
        pipeline = urllib.parse.quote(pipeline, safe="")
        self.base_url = (
            "https://api.buildkite.com/v2/organizations/"
            f"{organization}/pipelines/{pipeline}/builds"
        )

    def _request(
        self,
        *,
        body: Mapping[str, Any] | None = None,
        method: str = "GET",
        path: str = "",
        query: Sequence[tuple[str, str]] = (),
    ) -> Any:
        if not self.token:
            raise RuntimeError("The BUILDKITE_API_TOKEN repository secret is not set.")
        url = f"{self.base_url}{path}"
        if query:
            url = f"{url}?{urllib.parse.urlencode(query)}"
        return self.transport.request(
            url,
            body=body,
            headers={
                "Authorization": f"Bearer {self.token}",
                "Content-Type": "application/json",
                "User-Agent": "vllm-ci-command",
            },
            method=method,
        )

    def list_builds(
        self,
        commit: str,
        *,
        metadata: tuple[str, str] | None = None,
    ) -> list[dict[str, Any]]:
        query = [
            ("commit", commit),
            ("exclude_jobs", "true"),
            ("exclude_pipeline", "true"),
            ("per_page", "100"),
        ]
        if metadata:
            key, value = metadata
            query.append((f"meta_data[{key}]", value))
        response = self._request(query=query)
        if not isinstance(response, list):
            raise ApiError(None, "Buildkite API returned an invalid build list.")
        return response

    def create_build(self, body: Mapping[str, Any]) -> dict[str, Any]:
        return self._request(body=body, method="POST")

    def retry_failed_jobs(
        self,
        build_number: int,
        states: str,
    ) -> dict[str, Any]:
        number = urllib.parse.quote(str(build_number), safe="")
        return self._request(
            body={"states": states},
            method="PUT",
            path=f"/{number}/retry_failed_jobs",
        )


def parse_command(body: str) -> str | None:
    if body in {COMMAND_RUN_CI, COMMAND_RETRY_FAILED}:
        return body
    return None


def parse_trusted_users(value: str = "") -> set[str]:
    return {
        user.casefold() for item in value.split(",") for user in item.split() if user
    }


def has_ready_label(pr: Mapping[str, Any]) -> bool:
    return any(label["name"] in READY_LABELS for label in pr["labels"])


def is_trusted_permission(permission: str) -> bool:
    return permission in TRUSTED_PERMISSIONS


def authorize(
    *,
    actor: str,
    permission: str,
    pr: Mapping[str, Any],
    trusted_approval: bool = False,
    trusted_users: set[str] | None = None,
) -> tuple[bool, str]:
    trusted_users = trusted_users or set()
    if is_trusted_permission(permission):
        return True, f"repository {permission} permission"
    if actor.casefold() in trusted_users:
        return True, "configured trusted contributor"
    if actor.casefold() != pr["user"]["login"].casefold():
        return (
            False,
            "Only reviewers with write access can run CI before it is "
            "delegated to the PR author.",
        )
    if pr["draft"]:
        return False, "PR authors cannot run CI while the PR is a draft."
    if has_ready_label(pr):
        return True, "ready label"
    if trusted_approval:
        return True, "approval from a trusted reviewer"
    return (
        False,
        "A reviewer with write access must run `/ci run`, approve the PR, "
        "or add the `ready` label first.",
    )


def has_trusted_approval(
    github: GitHubClient,
    number: int,
    trusted_users: set[str],
) -> bool:
    if github.get_review_decision(number) != "APPROVED":
        return False

    latest_review_states: dict[str, tuple[str, str]] = {}
    for review in github.list_reviews(number):
        user = review.get("user") or {}
        login = user.get("login")
        state = review.get("state")
        if login and state in {"APPROVED", "CHANGES_REQUESTED", "DISMISSED"}:
            latest_review_states[login.casefold()] = (login, state)

    for login, state in latest_review_states.values():
        if state != "APPROVED":
            continue
        if login.casefold() in trusted_users:
            return True
        if is_trusted_permission(github.get_permission(login)):
            return True
    return False


def is_build_for_pr(build: Mapping[str, Any], pr_number: int) -> bool:
    pull_request = build.get("pull_request")
    if isinstance(pull_request, Mapping):
        build_pr_number = pull_request.get("id", pull_request.get("number"))
        if build_pr_number is not None:
            return str(build_pr_number) == str(pr_number)
    metadata = build.get("meta_data") or {}
    return str(metadata.get("github-pr-number")) == str(pr_number)


def is_active_build(build: Mapping[str, Any]) -> bool:
    return bool(build.get("blocked")) or build.get("state") in ACTIVE_BUILD_STATES


def select_latest_build(
    builds: Sequence[dict[str, Any]],
    pr_number: int,
) -> dict[str, Any] | None:
    matching = [build for build in builds if is_build_for_pr(build, pr_number)]
    return max(matching, key=lambda build: build.get("created_at", ""), default=None)


def create_build_payload(
    *,
    actor: str,
    comment_id: int,
    pr: Mapping[str, Any],
) -> dict[str, Any]:
    return {
        "commit": pr["head"]["sha"],
        "branch": pr["head"]["ref"],
        "message": f"PR #{pr['number']} {COMMAND_RUN_CI} by @{actor}",
        "pull_request_id": pr["number"],
        "pull_request_base_branch": pr["base"]["ref"],
        "pull_request_repository": pr["head"]["repo"]["clone_url"],
        "pull_request_labels": [label["name"] for label in pr["labels"]],
        "ignore_pipeline_branch_filters": True,
        "env": {
            "VLLM_CI_GITHUB_COMMENT_ID": str(comment_id),
            "VLLM_CI_TRIGGERED_BY": actor,
        },
        "meta_data": {
            "github-comment-id": str(comment_id),
            "github-pr-number": str(pr["number"]),
            "github-triggered-by": actor,
        },
    }


def add_reaction_safely(
    github: GitHubClient,
    comment_id: int,
    content: str,
) -> None:
    try:
        github.add_reaction(comment_id, content)
    except Exception as error:
        print(f"Could not add {content} reaction: {error}", file=sys.stderr)


def is_already_handled(github: GitHubClient, comment_id: int) -> bool:
    return any(
        reaction.get("content") in {"rocket", "-1"}
        and (reaction.get("user") or {}).get("login") == "github-actions[bot]"
        for reaction in github.list_reactions(comment_id)
    )


def handle_run_ci(
    *,
    actor: str,
    buildkite: BuildkiteClient,
    comment_id: int,
    github: GitHubClient,
    pr: Mapping[str, Any],
) -> str:
    duplicate_builds = buildkite.list_builds(
        pr["head"]["sha"],
        metadata=("github-comment-id", str(comment_id)),
    )
    duplicate = select_latest_build(duplicate_builds, pr["number"])
    if duplicate:
        return f"CI was already requested by this comment: {duplicate['web_url']}"

    current_builds = buildkite.list_builds(pr["head"]["sha"])
    active_build = next(
        (
            build
            for build in current_builds
            if is_build_for_pr(build, pr["number"]) and is_active_build(build)
        ),
        None,
    )
    if active_build:
        return f"CI is already running for this commit: {active_build['web_url']}"

    current_pr = github.get_pr(pr["number"])
    if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]:
        return (
            "The PR head changed while processing the command. Comment `/ci run` again."
        )

    build = buildkite.create_build(
        create_build_payload(
            actor=actor,
            comment_id=comment_id,
            pr=current_pr,
        )
    )
    return (
        f"Triggered [Buildkite CI #{build['number']}]({build['web_url']}) "
        f"for commit `{current_pr['head']['sha'][:12]}`."
    )


def handle_retry_failed(
    *,
    buildkite: BuildkiteClient,
    pr: Mapping[str, Any],
) -> str:
    builds = buildkite.list_builds(pr["head"]["sha"])
    build = select_latest_build(builds, pr["number"])
    if not build:
        return "No CI build exists for the current PR commit. Use `/ci run` first."
    if not build.get("finished_at") or is_active_build(build):
        return f"CI is still running for this commit: {build['web_url']}"

    retried = buildkite.retry_failed_jobs(build["number"], RETRY_STATES)
    if retried["retried_jobs_count"] == 0:
        return (
            f"No failed, timed-out, or expired jobs need retrying: {build['web_url']}"
        )
    return (
        f"Queued {retried['retried_jobs_count']} failed job(s) for retry in "
        f"[Buildkite CI #{build['number']}]({build['web_url']})."
    )


def run(
    event: Mapping[str, Any],
    github: GitHubClient,
    buildkite: BuildkiteClient,
    trusted_users_value: str = "",
) -> None:
    command = parse_command(event["comment"]["body"])
    if not command or "pull_request" not in event["issue"]:
        return

    issue_number = event["issue"]["number"]
    comment_id = event["comment"]["id"]
    actor = event["comment"]["user"]["login"]

    if is_already_handled(github, comment_id):
        print(f"Comment {comment_id} was already handled.")
        return
    add_reaction_safely(github, comment_id, "eyes")

    try:
        pr = github.get_pr(issue_number)
        permission = github.get_permission(actor)
        if pr["state"] != "open":
            github.add_comment(issue_number, "CI commands require an open PR.")
            return

        trusted_users = parse_trusted_users(trusted_users_value)
        should_check_approval = (
            not is_trusted_permission(permission)
            and actor.casefold() not in trusted_users
            and actor.casefold() == pr["user"]["login"].casefold()
            and not pr["draft"]
            and not has_ready_label(pr)
        )
        trusted_approval = should_check_approval and has_trusted_approval(
            github,
            issue_number,
            trusted_users,
        )
        allowed, reason = authorize(
            actor=actor,
            permission=permission,
            pr=pr,
            trusted_approval=trusted_approval,
            trusted_users=trusted_users,
        )
        if not allowed:
            add_reaction_safely(github, comment_id, "-1")
            github.add_comment(issue_number, f"@{actor}, {reason}")
            return

        print(f"Authorized @{actor}: {reason}")
        if command == COMMAND_RUN_CI:
            message = handle_run_ci(
                actor=actor,
                buildkite=buildkite,
                comment_id=comment_id,
                github=github,
                pr=pr,
            )
        else:
            message = handle_retry_failed(buildkite=buildkite, pr=pr)
        add_reaction_safely(github, comment_id, "rocket")
        github.add_comment(issue_number, message)
    except Exception:
        add_reaction_safely(github, comment_id, "confused")
        raise


def main() -> None:
    event_path = os.environ["GITHUB_EVENT_PATH"]
    with open(event_path, encoding="utf-8") as event_file:
        event = json.load(event_file)

    if not parse_command(event["comment"]["body"]):
        return

    github = GitHubClient(
        os.environ.get("GH_TOKEN", ""),
        os.environ["GITHUB_REPOSITORY"],
    )
    buildkite = BuildkiteClient(
        os.environ.get("BUILDKITE_API_TOKEN", ""),
        os.environ.get("BUILDKITE_ORGANIZATION", "vllm"),
        os.environ.get("BUILDKITE_PIPELINE", "ci"),
    )
    run(
        event,
        github,
        buildkite,
        os.environ.get("CI_TRUSTED_USERS", ""),
    )


if __name__ == "__main__":
    main()