#!/usr/bin/python3
"""
Usage: update-deps.py [dep_name]
"""

import base64
import datetime
import hashlib
import io
import json
import os
import re
import subprocess
import sys
import tarfile
import urllib.request
import zipfile
from pathlib import Path

TARGET_FILTER = None if len(sys.argv) < 2 else sys.argv[1]

SCRIPT_DIR = Path(__file__).parent
if "BUILD_WORKSPACE_DIRECTORY" in os.environ:
    SCRIPT_DIR = Path(os.environ["BUILD_WORKSPACE_DIRECTORY"]) / "build" / "deps"

GEN_DIR = SCRIPT_DIR / "gen"
ALL_DEPS = ["deps.jsonc", "build_deps.jsonc", "shared_deps.jsonc"]


TOP = """# WARNING: THIS FILE IS AUTOGENERATED BY update-deps.py DO NOT EDIT

http = use_extension("@//:build/exts/http.bzl", "http")

git_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")

"""

GITHUB_TAR_URL_TEMPLATE = "https://github.com/{owner}/{repo}/tarball/{commit}"

GITHUB_RELEASE_FILE_URL_TEMPLATE = (
    "https://github.com/{owner}/{repo}/releases/download/v{version}/{file}"
)

EXT_DEP_TEMPLATE = """# {name}
{ext_name}.{rule_name}({attrs}
)
use_repo({ext_name}, "{name}")

"""


REPO_RULE_DEP_TEMPLATE = """# {name}
{rule_name}({attrs}
)

"""


BAZEL_DEP_TEMPLATE = """# {name}
bazel_dep({attrs})

"""

BAZEL_DEP_OVERRIDE_TEMPLATE = """# {name}
bazel_dep({bazel_dep_attrs})
{override_type}({override_attrs}
)

"""

GITHUB_ACCESS_TOKEN = ""


def format_attr_list(attrs, single_line=False):
    if not attrs:
        return ""

    # buildifier (Bazel build file formatter) requires keys to be sorted, except name and module_name go first
    attr_list = sorted(
        attrs.items(), key=lambda kv: "" if kv[0] in {"name", "module_name"} else kv[0]
    )
    attr_strs = (f"{k} = {format_attr(v)}" for k, v in attr_list)

    if single_line:
        # Print attributes on a single line (for short declarations)
        return ", ".join(attr_strs)
    else:
        # Print one attribute per line (for long declarations)
        return "\n" + "\n".join(f"    {kv}," for kv in attr_strs)


def format_attr(v):
    if isinstance(v, (bool, int)):
        return str(v)
    elif isinstance(v, list):
        # Format lists as multiline with proper indentation and trailing comma
        return json.dumps(v, indent=8).replace("\n]", ",\n    ]")
    else:
        return json.dumps(v)


def format_repo_rule_dep(repo, rule_name, attrs):
    return REPO_RULE_DEP_TEMPLATE.format(
        rule_name=rule_name,
        name=repo["name"],
        attrs=format_attr_list(repo_attributes(repo) | attrs),
    )


def format_ext_dep(repo, ext_name, rule_name, attrs):
    return EXT_DEP_TEMPLATE.format(
        ext_name=ext_name,
        rule_name=rule_name,
        name=repo["name"],
        attrs=format_attr_list(repo_attributes(repo) | attrs),
    )


class RateLimitedException(Exception):
    pass


class AssetsException(Exception):
    pass


class UnsupportedException(Exception):
    pass


def github_urlopen(url):
    """
    A wrapper around urllib.request.urlopen() which parses GitHub rate limit errors and
    provides a more human-friendly explanation.
    """
    if GITHUB_ACCESS_TOKEN != "":
        url = urllib.request.Request(url)
        url.add_header("Authorization", f"Bearer {GITHUB_ACCESS_TOKEN}")
    try:
        return urllib.request.urlopen(url)
    except urllib.error.HTTPError as e:
        reset_ts = e.headers["x-ratelimit-reset"]
        if e.code != 403 or not reset_ts:
            raise
        reset_dt = datetime.datetime.fromtimestamp(int(reset_ts))
        reset_iso_utc = reset_dt.astimezone(datetime.UTC).isoformat(" ")
        reset_iso_local = reset_dt.isoformat(" ")
        raise RateLimitedException(
            f"""
We have been rate-limited by GitHub. We can make API calls again at:
  {reset_iso_utc} UTC ({reset_iso_local} local time).
"""
            + """
You can try re-running the script and specifying an access token since authenticated
GitHub API requests have a higher rate limit.
"""
            if GITHUB_ACCESS_TOKEN == ""
            else ""
        ) from e


def github_last_commit(repo):
    owner = repo["owner"]
    github_repo = repo["repo"]
    branch = repo.get("branch", "master")
    api_url = f"https://api.github.com/repos/{owner}/{github_repo}/commits/{branch}"
    commits = json.loads(github_urlopen(api_url).read())
    return commits["sha"]


def get_url_content_sha256(url):
    return hashlib.sha256(urllib.request.urlopen(url).read()).hexdigest()


def get_bcr_module_bazel_url(module_name, version):
    """Generate the URL for a MODULE.bazel file from BCR."""
    return f"https://raw.githubusercontent.com/bazelbuild/bazel-central-registry/refs/heads/main/modules/{module_name}/{version}/MODULE.bazel"


def get_bcr_module_bazel_integrity(module_name, version):
    """Fetch MODULE.bazel from BCR and compute its SHA256 integrity hash."""
    url = get_bcr_module_bazel_url(module_name, version)
    content = urllib.request.urlopen(url).read()
    sha256 = hashlib.sha256(content).digest()
    return f"sha256-{base64.b64encode(sha256).decode()}"


def repo_attributes(repo):
    repo_attrs = {}

    for option in (
        "name",
        "build_file",
        "repo_mapping",
        "downloaded_file_path",
        "build_file_content",
        "patches",
    ):
        if option in repo:
            repo_attrs[option] = repo[option]

    if "patches" in repo_attrs:
        repo_attrs["patch_strip"] = 1

    if "use_module_bazel_from_bcr" in repo:
        url = get_bcr_module_bazel_url(repo["name"], repo["use_module_bazel_from_bcr"])
        integrity = get_bcr_module_bazel_integrity(
            repo["name"], repo["use_module_bazel_from_bcr"]
        )
        repo_attrs["remote_file_urls"] = {"MODULE.bazel": [url]}
        repo_attrs["remote_file_integrity"] = {"MODULE.bazel": integrity}

    return repo_attrs


def format_bazel_dep_with_override(repo, override_type, override_attrs):
    """Format bazel_dep + override."""
    name = repo["name"]

    # bazel_dep attributes
    bazel_dep_attrs = {"name": name}
    if "repo_name" in repo:
        bazel_dep_attrs["repo_name"] = repo["repo_name"]

    # Override attributes - use repo_attributes but swap name for module_name
    base_attrs = repo_attributes(repo)
    base_attrs["module_name"] = name
    del base_attrs["name"]

    return BAZEL_DEP_OVERRIDE_TEMPLATE.format(
        name=name,
        bazel_dep_attrs=format_attr_list(bazel_dep_attrs, single_line=True),
        override_type=override_type,
        override_attrs=format_attr_list(base_attrs | override_attrs),
    )


def gen_github_tarball(repo):
    owner = repo["owner"]
    github_repo = repo["repo"]

    commit = github_last_commit(repo)
    if "freeze_commit" in repo:
        if repo["freeze_commit"] != commit:
            print(
                "frozen, update available ",
                repo["freeze_commit"][:7],
                " -> ",
                commit[:7],
                end="",
            )
        commit = repo["freeze_commit"]
    else:
        print(commit[:7], end="")

    prefix = f"{owner}-{github_repo}-{commit[:7]}"
    if "extra_strip_prefix" in repo:
        prefix = prefix + repo["extra_strip_prefix"]

    url = GITHUB_TAR_URL_TEMPLATE.format(
        owner=owner,
        repo=github_repo,
        commit=commit,
    )

    if "freeze_sha256" in repo:
        sha256 = repo["freeze_sha256"]
    else:
        sha256 = get_url_content_sha256(url)

    attrs = dict(
        url=url,
        strip_prefix=prefix,
        sha256=sha256,
        type="tgz",
    )

    if repo.get("use_bazel_dep"):
        return format_bazel_dep_with_override(
            repo,
            override_type="archive_override",
            override_attrs=attrs,
        )
    else:
        return format_ext_dep(
            repo,
            ext_name="http",
            rule_name="archive",
            attrs=attrs,
        )


def github_last_release(repo):
    owner = repo["owner"]
    github_repo = repo["repo"]
    api_url = f"https://api.github.com/repos/{owner}/{github_repo}/releases/latest"
    return json.loads(github_urlopen(api_url).read())


def github_release(repo, tag_name):
    owner = repo["owner"]
    github_repo = repo["repo"]
    api_url = (
        f"https://api.github.com/repos/{owner}/{github_repo}/releases/tags/{tag_name}"
    )
    return json.loads(github_urlopen(api_url).read())


def gen_github_release(repo):
    try:
        release = github_last_release(repo)
    except urllib.error.HTTPError as e:
        # If a repo only has pre-releases, github_last_release will throw a 404 error.
        # In that case, we must specify a "freeze_version".
        if e.code != 404 or "freeze_version" not in repo:
            raise
        release = None

    if "freeze_version" in repo:
        frozen_release = github_release(repo, repo["freeze_version"])
        if release is not None and frozen_release["tag_name"] != release["tag_name"]:
            print(
                "frozen, update available: {} -> {}".format(
                    frozen_release["tag_name"], release["tag_name"]
                ),
                end="",
            )
        release = frozen_release
    else:
        print(release["tag_name"], end="")

    if "file_regex" in repo:
        # Using file_regex to select a user-uploaded asset
        url = get_release_asset(repo, release)
    else:
        # Using Github-generated tarball
        url = release["tarball_url"]

    type = "tgz"
    if url.endswith(".zip"):
        type = "zip"
    elif url.endswith(".xz"):
        type = "xz"
    elif url.endswith(".tar.bz2"):
        type = "tar.bz2"

    content = urllib.request.urlopen(url).read()

    if "freeze_sha256" in repo:
        sha256 = repo["freeze_sha256"]
    else:
        sha256 = hashlib.sha256(content).hexdigest()

    file_type = repo.get("file_type", "archive")
    if file_type == "archive":
        if "strip_prefix" in repo:
            prefix = repo["strip_prefix"]
        elif url.endswith(".zip"):
            with zipfile.ZipFile(io.BytesIO(content)) as zip:
                prefix = os.path.commonprefix(zip.namelist())
        else:
            with tarfile.open(fileobj=io.BytesIO(content)) as tgz:
                prefix = os.path.commonprefix(tgz.getnames())

        attrs = dict(
            url=url,
            strip_prefix=prefix,
            sha256=sha256,
            type=type,
        )

        if repo.get("use_bazel_dep"):
            return format_bazel_dep_with_override(
                repo,
                override_type="archive_override",
                override_attrs=attrs,
            )
        else:
            return format_ext_dep(
                repo,
                ext_name="http",
                rule_name="archive",
                attrs=attrs,
            )
    elif file_type == "executable":
        if repo.get("use_bazel_dep"):
            raise UnsupportedException(
                "use_bazel_dep is not supported for executable file_type"
            )
        return format_ext_dep(
            repo,
            ext_name="http",
            rule_name="file",
            attrs=dict(url=url, sha256=sha256, executable=True),
        )
    else:
        raise UnsupportedException("Unsupported file_type: " + file_type)


def get_release_asset(repo, release):
    file_regex = re.compile(repo["file_regex"])
    assets = [a for a in release["assets"] if file_regex.match(a["name"])]

    if len(assets) == 0:
        raise AssetsException("No assets found: " + json.dumps(release))

    if len(assets) > 1:
        raise AssetsException(
            "Too many assets, use more specific file_regex: "
            + str([a["name"] for a in assets])
        )

    return assets[0]["browser_download_url"]


def gen_git_clone(repo):
    url = repo["url"]

    # We used to clone the repository here to get a shallow_since timestamp, but based
    # on # https://github.com/bazelbuild/bazel/issues/12857 it is unclear if this is
    # actually helpful.
    ls_remote = subprocess.run(
        ["git", "ls-remote", url, repo["branch"]], capture_output=True, text=True
    )
    ls_remote.check_returncode()
    commit = ls_remote.stdout.strip().split()[0]

    if "freeze_commit" in repo:
        freeze_commit = repo["freeze_commit"]
        if freeze_commit != commit:
            print(
                "frozen, update available ",
                repo["freeze_commit"][:7],
                " -> ",
                commit[:7],
                end="",
            )
            commit = freeze_commit
        else:
            print(commit[:7], end="")

    attrs = dict(
        remote=url,
        commit=commit,
    )

    if repo.get("use_bazel_dep"):
        return format_bazel_dep_with_override(
            repo,
            override_type="git_override",
            override_attrs=attrs,
        )
    else:
        return format_repo_rule_dep(
            repo,
            rule_name="git_repository",
            attrs=attrs,
        )


def get_bcr_version(name: str) -> str:
    module_versions_url = f"https://bcr.bazel.build/modules/{name}/metadata.json"
    with urllib.request.urlopen(module_versions_url) as res:
        meta = json.load(res)

        # Assuming the newer versions are appended to the end of the list
        for version in reversed(meta["versions"]):
            # Do not recommend a yanked version
            if version in meta["yanked_versions"]:
                continue

            # Make a best efforts attempt to exclude pre-releases
            if re.search(r"(beta|rc|alpha|dev)", version):
                continue

            # TODO: Consider parsing versions with `packaging`. However, there is no standardized
            # version scheme for BCR dependencies, so this would definitely fail on some deps.

            return version


def gen_bazel_dep(repo):
    name = repo["name"]

    latest_version = get_bcr_version(name)

    if "freeze_version" in repo:
        frozen_version = repo["freeze_version"]
        if frozen_version != latest_version:
            print(
                f"frozen, update available: {frozen_version} -> {latest_version}",
                end="",
            )
        version = frozen_version
    else:
        print(latest_version, end="")
        version = latest_version

    if "patches" in repo:
        return BAZEL_DEP_OVERRIDE_TEMPLATE.format(
            name=name,
            bazel_dep_attrs=format_attr_list(
                dict(name=name, version=version), single_line=True
            ),
            override_type="single_version_override",
            override_attrs=format_attr_list(
                {
                    "module_name": name,
                    "patch_strip": 1,
                    "patches": repo["patches"],
                }
            ),
        )

    return BAZEL_DEP_TEMPLATE.format(
        name=name,
        attrs=format_attr_list(dict(name=name, version=version), single_line=True),
    )


def gen_repo_str(repo):
    if repo["type"] == "github_tarball":
        return gen_github_tarball(repo)
    elif repo["type"] == "github_release":
        return gen_github_release(repo)
    elif repo["type"] == "git_clone":
        return gen_git_clone(repo)
    elif repo["type"] == "bazel_dep":
        return gen_bazel_dep(repo)
    else:
        raise UnsupportedException(f"Unsupported repo type: {repo['type']}")


def gen_repo_bzl(repo, current_dep):
    """Generate and return the content for a repository dependency."""
    print("Checking", repo["name"], "... ", end="", flush=True)
    if TARGET_FILTER is not None and not repo["name"].startswith(TARGET_FILTER):
        print("skipped")
        return current_dep

    content = gen_repo_str(repo)
    print()
    return content


def gen_deps_bzl(repo_contents, deps_bzl):
    """Generate the index file by concatenating all repository contents."""
    deps_bzl_content = TOP

    # Concatenate all repository contents
    for content in repo_contents:
        if content:
            deps_bzl_content += content

    with deps_bzl.open("w") as f:
        # Strip extra trailing newline but keep one for POSIX compliance
        f.write(deps_bzl_content.rstrip("\n") + "\n")


def process_deps(deps, current_deps, deps_bzl):
    # Sort repositories by name for consistent ordering (buildifier preference)
    sorted_repos = sorted(deps["repositories"], key=lambda r: r["name"])

    # Generate content for each repository
    repo_contents = []
    for repo in sorted_repos:
        content = gen_repo_bzl(repo, current_deps.get(repo["name"]))
        repo_contents.append(content)

    gen_deps_bzl(repo_contents, deps_bzl)


def split_bzl_file(file: Path) -> dict[str, str]:
    # Creates a map from dependency name to generated code for that dep
    deps = {}
    text = file.read_text()

    blocks = iter(re.finditer(r"^# (.*)$", text, re.MULTILINE))
    a = next(blocks)

    for b in blocks:
        # Key: dependency name from comment line
        # Value: generated code (all text until the next comment line)
        deps[a.groups()[0]] = text[a.start() : b.start()]
        a = b

    deps[a.groups()[0]] = text[a.start() :]

    return deps


def strip_comments(text):
    # capture string literals first, comments send
    regex = re.compile(r"(\".*\")|(//.*$)", re.MULTILINE)
    return regex.sub(
        lambda match: "" if match.group(2) is not None else match.group(1), text
    )


def read_access_token():
    if not sys.stdin.isatty():
        return ""

    # 1. Try to obtain token from the gh tool
    try:
        res = subprocess.run(["gh", "auth", "token"], capture_output=True)
        if res.returncode == 0:
            return res.stdout.decode().strip()
        else:
            # User has gh but is not logged in
            print("Please log in to Github")
            print("$ gh auth login")
            raise SystemExit
    except FileNotFoundError:
        pass  # User does not have gh tool installed

        print(
            """Follow these steps to obtain a GitHub API access token with
appropriate permissions:

1. On github.com, go to
Settings > Developer Settings > Personal access tokens > Fine-grained tokens.
2. Generate a new token with default settings.

Alternatively, install the gh CLI tool to save time <https://github.com/cli/cli#installation>.
"""
        )
        print(
            "Please enter GitHub API access token (or empty to skip): ",
            end="",
            flush=True,
        )

        return sys.stdin.readline().strip("\n")


def process_config(deps_file):
    deps_path = SCRIPT_DIR / deps_file
    bzl_path = GEN_DIR / Path(deps_file).with_suffix(".MODULE.bazel").name

    # Create output directory if it doesn't exist
    GEN_DIR.mkdir(parents=True, exist_ok=True)

    # Load existing generated deps file (if any)
    try:
        current_deps = split_bzl_file(bzl_path)
    except FileNotFoundError:
        current_deps = {}

    try:
        with deps_path.open() as fp:
            json_text = strip_comments(fp.read())
            process_deps(json.loads(json_text), current_deps, bzl_path)
    except FileNotFoundError:
        pass


def run():
    if TARGET_FILTER is None:
        global GITHUB_ACCESS_TOKEN
        GITHUB_ACCESS_TOKEN = read_access_token()

        # Clean all generated .bazel files
        for f in GEN_DIR.glob("*.bazel"):
            f.unlink()

    for deps in ALL_DEPS:
        process_config(deps)


run()