import argparse
import hashlib
import json
from pathlib import Path, PurePosixPath
import re
import shutil
import subprocess
import tarfile


MAX_BYTES = 256 * 1024 * 1024
TARGETS = {
    "linux-x86_64": ("x86_64", "Advanced Micro Devices X86-64"),
    "linux-aarch64": ("aarch64", "AArch64"),
}
OMIT = {"__pycache__", "site-packages", "test", "tests", "tkinter", "turtledemo", "idlelib", "ensurepip"}
PROBE = """import ctypes, hashlib, importlib.machinery, json, platform, sqlite3, ssl, sys, zlib
print(json.dumps({"version": platform.python_version(), "machine": platform.machine(),
                  "executable": sys.executable, "prefix": sys.prefix,
                  "suffixes": importlib.machinery.EXTENSION_SUFFIXES}))
"""


def checksum(path):
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def selected(name, minor):
    path = PurePosixPath(name)
    if path == PurePosixPath(f"python/bin/python{minor}"):
        return True
    if len(path.parts) == 3 and path.parts[:2] == ("python", "lib"):
        return bool(re.fullmatch(r"libpython[\d.]+\.so(?:\.\d+)*", path.name))
    prefix = ("python", "lib", f"python{minor}")
    return (path.parts[:3] == prefix and not OMIT.intersection(path.parts[3:])
            and not any(part.startswith("config-") for part in path.parts[3:])
            and not path.name.startswith("_test") and path.suffix not in (".pyc", ".pyo", ".a"))


def extract(archive, output, minor):
    total = 0
    seen = set()
    with tarfile.open(archive, "r:gz") as source:
        for member in source:
            name = member.name.rstrip("/")
            parts = name.split("/")
            if (parts[0] != "python" or any(part in ("", ".", "..") for part in parts)
                    or "\\" in name or name in seen or len(seen) >= 65536):
                raise ValueError("invalid, repeated or excessive archive paths")
            seen.add(name)
            total += member.size
            if total > MAX_BYTES:
                raise ValueError("Python source exceeds the runtime byte budget")
            if not member.isfile() or not selected(name, minor):
                continue
            destination = output / name
            destination.parent.mkdir(parents=True, exist_ok=True)
            with source.extractfile(member) as incoming, destination.open("xb") as outgoing:
                shutil.copyfileobj(incoming, outgoing)
            destination.chmod(0o755 if name == f"python/bin/python{minor}" else 0o644)


def check_elf(path, readelf, target, baseline):
    with path.open("rb") as stream:
        if stream.read(4) != b"\x7fELF":
            return
    result = subprocess.run([str(readelf), "--wide", "--file-header", "--version-info", str(path)],
                            check=True, capture_output=True, text=True, timeout=15)
    if len(result.stdout) > 256 * 1024 or TARGETS[target][1] not in result.stdout:
        raise ValueError("runtime ELF has the wrong target or excessive metadata")
    if "Name: GLIBC_PRIVATE" in result.stdout:
        raise ValueError("runtime ELF depends on private glibc interfaces")
    versions = [tuple(map(int, value.split(".")))
                for value in re.findall(r"Name: GLIBC_(\d+(?:\.\d+)+)", result.stdout)]
    if versions and max(versions) > baseline:
        raise ValueError(f"{path.name} requires GLIBC_{'.'.join(map(str, max(versions)))}")


def wheel_tags(minor, target, baseline):
    architecture = TARGETS[target][0]
    oldest = 5 if architecture == "x86_64" else 17
    if baseline[0] != 2 or baseline[1] < oldest:
        raise ValueError("unsupported target glibc baseline")
    aliases = {17: "manylinux2014", 12: "manylinux2010", 5: "manylinux1"}
    platforms = []
    for glibc_minor in range(baseline[1], oldest - 1, -1):
        platforms.append(f"manylinux_2_{glibc_minor}_{architecture}")
        if glibc_minor in aliases:
            platforms.append(f"{aliases[glibc_minor]}_{architecture}")
    platforms.append(f"linux_{architecture}")
    python_minor = int(minor.split(".")[1])
    abi = "cp" + minor.replace(".", "")
    tags = [f"{abi}-{abi}-{platform}" for platform in platforms]
    tags += [f"{abi}-abi3-{platform}" for platform in platforms]
    tags += [f"{abi}-none-{platform}" for platform in platforms]
    tags += [f"cp3{version}-abi3-{platform}"
             for version in range(python_minor - 1, 1, -1) for platform in platforms]
    interpreters = [f"py3{python_minor}", "py3"]
    interpreters += [f"py3{version}" for version in range(python_minor - 1, -1, -1)]
    tags += [f"{interpreter}-none-{platform}" for interpreter in interpreters for platform in platforms]
    tags += [f"{abi}-none-any"]
    tags += [f"{interpreter}-none-any" for interpreter in interpreters]
    return tags


def prepare(archive, expected, version, target, baseline, readelf, output, runner=None, sysroot=None):
    if (not re.fullmatch(r"[0-9a-f]{64}", expected) or archive.stat().st_size > MAX_BYTES
            or checksum(archive) != expected):
        raise ValueError("Python archive SHA-256 or size mismatch")
    if not re.fullmatch(r"3\.\d+\.\d+", version) or target not in TARGETS:
        raise ValueError("invalid Python version or target")
    if not readelf.is_absolute() or not readelf.is_file():
        raise ValueError("readelf must be an explicit absolute executable")
    if ((runner is None) != (sysroot is None)
            or runner is not None and (not runner.is_absolute() or not runner.is_file()
                                       or not sysroot.is_absolute() or not sysroot.is_dir())):
        raise ValueError("cross probing requires an explicit QEMU runner and target sysroot")
    minor = ".".join(version.split(".")[:2])
    output.mkdir(mode=0o700)
    try:
        extract(archive, output, minor)
        root = output / "python"
        binary = root / "bin" / f"python{minor}"
        if not binary.is_file() or not (root / "lib" / f"python{minor}" / "LICENSE.txt").is_file():
            raise ValueError("Python archive lacks its interpreter or license")
        for path in root.rglob("*"):
            if path.is_file():
                check_elf(path, readelf, target, baseline)
        environment = {"HOME": str(root), "TMPDIR": str(output), "PATH": "/usr/bin:/bin"}
        prefix = [str(runner), "-L", str(sysroot)] if runner else []
        result = subprocess.run(prefix + [str(binary), "-I", "-S", "-B", "-c", PROBE],
                                check=True, capture_output=True, text=True, env=environment, timeout=30)
        probe = json.loads(result.stdout)
        if (probe["version"] != version or probe["machine"] != TARGETS[target][0]
                or Path(probe["prefix"]) != root or Path(probe["executable"]) != binary):
            raise ValueError("Python runtime is not relocatable or has the wrong identity")
        tags = wheel_tags(minor, target, baseline)
        markers = {"implementation_name": "cpython", "implementation_version": version,
                   "os_name": "posix", "platform_machine": TARGETS[target][0],
                   "platform_python_implementation": "CPython", "platform_release": "",
                   "platform_system": "Linux", "platform_version": "", "python_full_version": version,
                   "python_version": minor, "sys_platform": "linux"}
        descriptor = {"format": "dever-worker-runtime-v1", "ecosystem": "pip", "target": target,
                      "executable": f"bin/python{minor}", "arguments": ["-I", "-S", "-B"],
                      "python_markers": markers, "python_wheel_tags": tags,
                      "python_extension_suffixes": probe["suffixes"]}
        (root / "dever-runtime.json").write_text(json.dumps(descriptor, indent=2) + "\n")
        files = [{"source": str(path.relative_to(output)), "path": str(path.relative_to(root)),
                  "sha256": checksum(path)} for path in sorted(root.rglob("*")) if path.is_file()]
        author = {"name": "cpython", "version": version, "python_markers": markers,
                  "python_wheel_tags": tags, "files": files}
        (output / "python-runtime.json").write_text(json.dumps(author, indent=2) + "\n")
        return {"version": version, "target": target, "glibc_max": ".".join(map(str, baseline)),
                "source_sha256": expected, "files": len(files)}
    except BaseException:
        shutil.rmtree(output)
        raise


def main():
    parser = argparse.ArgumentParser(description="Prepare verified portable Python inputs for native-release")
    parser.add_argument("--archive", type=Path, required=True)
    parser.add_argument("--sha256", required=True)
    parser.add_argument("--version", required=True)
    parser.add_argument("--target", choices=TARGETS, required=True)
    parser.add_argument("--glibc-max", default="2.17")
    parser.add_argument("--readelf", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--runner", type=Path, help="explicit QEMU user-mode executable for cross probing")
    parser.add_argument("--sysroot", type=Path, help="target system libraries for the explicit QEMU runner")
    arguments = parser.parse_args()
    if not re.fullmatch(r"2\.\d+", arguments.glibc_max):
        parser.error("glibc-max must be a 2.x baseline")
    print(json.dumps(prepare(arguments.archive.resolve(), arguments.sha256, arguments.version,
                             arguments.target, tuple(map(int, arguments.glibc_max.split("."))),
                             arguments.readelf, arguments.output.absolute(), arguments.runner,
                             arguments.sysroot), indent=2))


if __name__ == "__main__":
    main()