#!/usr/bin/env python3
# Private-overlay helper; not for Gentoo submission.
"""Update a Codex ebuild from a source tarball or Cargo.lock (Python 3.11+).

By default, a versioned source tarball selects dev-util/codex/codex-VERSION.ebuild
in this overlay. Use --ebuild for another target or --print for the old output
mode. Archives are read directly without extracting files.
"""

import argparse
import io
import os
import re
import sys
import tarfile
import tempfile
import tomllib
import urllib.parse
import urllib.request
from pathlib import Path, PurePosixPath


def read_source(source):
    versions = set()
    match = re.fullmatch(
        r"(?:codex-|rust-v)(\d+\.\d+\.\d+)\.tar(?:\.(?:gz|bz2|xz))?",
        source.name,
    )
    if match:
        versions.add(match[1])
    if tarfile.is_tarfile(source):
        with tarfile.open(source, "r:*") as archive:
            lockfiles = [
                member
                for member in archive
                if member.isfile()
                and (
                    member.name == "codex-rs/Cargo.lock"
                    or member.name.endswith("/codex-rs/Cargo.lock")
                )
            ]
            if len(lockfiles) != 1:
                raise ValueError(
                    f"{source}: expected one codex-rs/Cargo.lock, "
                    f"found {len(lockfiles)}"
                )
            parts = PurePosixPath(lockfiles[0].name).parts
            if len(parts) == 3:
                match = re.fullmatch(r"codex-rust-v(\d+\.\d+\.\d+)", parts[0])
                if match:
                    versions.add(match[1])
            with archive.extractfile(lockfiles[0]) as lockfile:
                packages = tomllib.load(lockfile)["package"]
    else:
        with source.open("rb") as lockfile:
            packages = tomllib.load(lockfile)["package"]
    if len(versions) > 1:
        raise ValueError("source filename and archive root have different versions")
    return packages, next(iter(versions), None)


def canonical_repo(uri):
    if ";" in uri:
        raise ValueError(f"Git repository URL contains a GIT_CRATES delimiter: {uri}")
    parsed = urllib.parse.urlsplit(uri)
    if parsed.scheme != "https" or not parsed.netloc or parsed.query or parsed.fragment:
        raise ValueError(f"unsupported Git repository URL: {uri}")
    if parsed.username or parsed.password:
        raise ValueError("Git repository URLs must not contain credentials")
    return uri.rstrip("/").removesuffix(".git")


def dependencies(packages):
    registry = []
    git = {}
    for package in packages:
        source = package.get("source", "")
        if not source:
            continue
        name, version = package["name"], package["version"]
        if not re.fullmatch(r"[A-Za-z0-9_-]+", name):
            raise ValueError(f"invalid crate name: {name}")
        if source == "registry+https://github.com/rust-lang/crates.io-index":
            if not re.fullmatch(
                r"\d+\.\d+\.\d+(?:-[A-Za-z0-9.-]+)?(?:\+[A-Za-z0-9.-]+)?",
                version,
            ):
                raise ValueError(f"invalid crate version: {version}")
            registry.append(f"{name}@{version}")
        elif source.startswith("git+"):
            parsed = urllib.parse.urlsplit(source[4:])
            commit = parsed.fragment
            if not re.fullmatch(r"[0-9a-f]{40}", commit):
                raise ValueError(f"Git crate {name} has no full locked commit")
            repo = canonical_repo(urllib.parse.urlunsplit(parsed._replace(query="", fragment="")))
            if name in git and git[name] != (repo, commit):
                raise ValueError(f"Git crate {name} uses multiple repositories or commits")
            git[name] = (repo, commit)
        else:
            raise ValueError(f"unsupported source: {source}")
    if len(registry) != len(set(registry)):
        raise ValueError("duplicate registry name/version entries")
    return sorted(registry, key=lambda item: tuple(item.split("@"))), git


def section(text, pattern, label):
    matches = list(re.finditer(pattern, text))
    if len(matches) != 1:
        raise ValueError(f"expected one literal {label} section, found {len(matches)}")
    return matches[0]


def read_git_mapping(body):
    entries = {}
    for line in body.splitlines():
        if not line.strip() or line.lstrip().startswith("#"):
            continue
        match = re.fullmatch(r"\s*\[([A-Za-z0-9_-]+)\]=(['\"])(.*?)\2\s*", line)
        if not match:
            raise ValueError("unsupported GIT_CRATES entry; expected a literal quoted value")
        name, _, value = match.groups()
        fields = value.split(";")
        if len(fields) not in (2, 3, 4) or any(char in value for char in "$`'\\\r\n"):
            raise ValueError(f"unsupported GIT_CRATES value for {name}")
        if name in entries:
            raise ValueError(f"duplicate GIT_CRATES entry: {name}")
        entries[name] = fields
    return entries


def archive_paths(repo, commit, names, distdir):
    parsed = urllib.parse.urlsplit(repo)
    project = parsed.path.rsplit("/", 1)[1]
    filename = f"{project}-{commit}.tar.gz"
    cached = distdir / filename
    if cached.is_file():
        archive = tarfile.open(cached, "r:*")
    else:
        if parsed.netloc == "github.com":
            url = f"{repo}/archive/{commit}.tar.gz"
        elif parsed.netloc == "gitlab.com":
            url = f"{repo}/-/archive/{commit}/{filename}"
        else:
            raise ValueError(f"cannot discover new Git crate paths for {repo}")
        print(f"Reading Git crate paths from {url}", file=sys.stderr)
        request = urllib.request.Request(url, headers={"User-Agent": "Codex-private-overlay-helper/1.0"})
        with urllib.request.urlopen(request, timeout=30) as response:
            data = response.read(32 * 1024 * 1024 + 1)
        if len(data) > 32 * 1024 * 1024:
            raise ValueError(f"Git source archive exceeds 32 MiB: {repo}")
        archive = tarfile.open(fileobj=io.BytesIO(data), mode="r:*")
    paths = {name: [] for name in names}
    with archive:
        for member in archive:
            path = PurePosixPath(member.name)
            if not member.isfile() or path.name != "Cargo.toml":
                continue
            if path.is_absolute() or ".." in path.parts:
                raise ValueError(f"unsafe Cargo.toml path in Git archive: {member.name}")
            with archive.extractfile(member) as manifest:
                package = tomllib.load(manifest).get("package", {})
            name = package.get("name")
            if name in paths:
                crate_path = path.parent.as_posix().replace(commit, "%commit%")
                if ";" in crate_path:
                    raise ValueError(f"Git crate path contains a GIT_CRATES delimiter: {name}")
                paths[name].append(crate_path)
    for name, candidates in paths.items():
        if len(candidates) != 1:
            raise ValueError(f"Git crate {name}: expected one Cargo.toml, found {len(candidates)}")
    return {name: candidates[0] for name, candidates in paths.items()}


def render_git(git, existing, distdir, newline):
    values = {}
    missing = {}
    for name, (repo, commit) in git.items():
        previous = existing.get(name)
        if previous and canonical_repo(previous[0]) == repo and previous[1] == commit:
            values[name] = ";".join(previous)
        else:
            missing.setdefault((repo, commit), []).append(name)
    for (repo, commit), names in missing.items():
        paths = archive_paths(repo, commit, names, distdir)
        for name in names:
            values[name] = f"{repo};{commit};{paths[name]}"
    lines = ["declare -A GIT_CRATES=("]
    for name, value in sorted(values.items()):
        if any(char in value for char in "$`'\\\r\n"):
            raise ValueError(f"unsupported characters in Git crate path: {name}")
        lines.append(f"\t[{name}]='{value}'")
    lines.append(")")
    return newline.join(lines)


def replace_sections(ebuild, registry, git, distdir):
    original = ebuild.read_bytes()
    text = original.decode("utf-8")
    crates = section(text, r'(?m)^CRATES="[^"]*"[ \t]*(?=\r?$)', "CRATES")
    git_crates = section(
        text,
        r"(?m)^declare -A GIT_CRATES=\((?:\r?\n(?P<body>[\s\S]*?)^)?\)[ \t]*(?=\r?$)",
        "GIT_CRATES",
    )
    newline = "\r\n" if "\r\n" in text else "\n"
    existing = read_git_mapping(git_crates.group("body") or "")
    rendered_crates = newline.join(['CRATES="', *("\t" + crate for crate in registry), '"'])
    rendered_git = render_git(git, existing, distdir, newline)
    for match, replacement in sorted(
        [(crates, rendered_crates), (git_crates, rendered_git)],
        key=lambda item: item[0].start(),
        reverse=True,
    ):
        text = text[:match.start()] + replacement + text[match.end():]
    updated = text.encode("utf-8")
    if updated == original:
        return False
    temporary = None
    try:
        with tempfile.NamedTemporaryFile(dir=ebuild.parent, prefix="." + ebuild.name + ".", delete=False) as output:
            temporary = Path(output.name)
            os.fchmod(output.fileno(), ebuild.stat().st_mode & 0o7777)
            output.write(updated)
        if ebuild.read_bytes() != original:
            raise ValueError(f"{ebuild} changed while preparing the update; retry")
        temporary.replace(ebuild)
        temporary = None
    finally:
        if temporary is not None:
            temporary.unlink()
    return True


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source", type=Path, help="Codex source tarball or extracted Cargo.lock")
    modes = parser.add_mutually_exclusive_group()
    modes.add_argument("--ebuild", type=Path, help="ebuild to update instead of selecting it by version")
    modes.add_argument("--print", action="store_true", dest="print_only", help="print CRATES and report Git dependencies without editing")
    parser.add_argument("--distdir", type=Path, help="cached Git source archives (default: source file's directory)")
    args = parser.parse_args()
    try:
        packages, version = read_source(args.source)
        registry, git = dependencies(packages)
        if args.print_only:
            print('CRATES="')
            for crate in registry:
                print("\t" + crate)
            print('"')
            if git:
                print("Git dependencies:", file=sys.stderr)
                for name, (repo, commit) in sorted(git.items()):
                    print(f"{name}: {repo}#{commit}", file=sys.stderr)
            return
        ebuild = args.ebuild
        if ebuild is None:
            if version is None:
                raise ValueError("cannot infer the Codex version; specify --ebuild or --print")
            ebuild = Path(__file__).resolve().parents[1] / "dev-util/codex" / f"codex-{version}.ebuild"
        match = re.fullmatch(r"codex-(\d+\.\d+\.\d+)(?:-r\d+)?\.ebuild", ebuild.name)
        if match and version is not None and match[1] != version:
            raise ValueError(f"source version {version} does not match {ebuild.name}")
        changed = replace_sections(ebuild, registry, git, args.distdir or args.source.parent)
        print(f"{'Updated' if changed else 'Already current'}: {ebuild} ({len(registry)} registry crates, {len(git)} Git crates)")
    except (OSError, ValueError, KeyError, TypeError, tarfile.TarError) as error:
        parser.error(str(error))


if __name__ == "__main__":
    main()