from __future__ import annotations

import re
from ast import literal_eval
from collections import defaultdict

import httpx
from utils import ROOT

URL = "https://raw.githubusercontent.com/ofek/pyapp/master/build.rs"
OUTPUT_FILE = ROOT / "src" / "hatch" / "python" / "distributions.py"
ARCHES = {("linux", "x86"): "i686", ("windows", "x86_64"): "amd64", ("windows", "x86"): "i386"}

# system, architecture, ABI, CPU variant, GIL variant
MAX_IDENTIFIER_COMPONENTS = 5


def parse_distributions(contents: str, constant: str):
    match = re.search(f"^const {constant}.+?^];$", contents, flags=re.DOTALL | re.MULTILINE)
    if not match:
        message = f"Could not find {constant} in {URL}"
        raise ValueError(message)

    block = match.group(0).replace('",\n', '",')
    for raw_line in block.splitlines()[1:-1]:
        line = raw_line.strip()
        if not line or line.startswith("//"):
            continue

        identifier, *data, source = literal_eval(line[:-1])
        os, arch = data[:2]
        if arch == "powerpc64":
            arch = "ppc64le"
        elif os == "macos" and arch == "aarch64":
            arch = "arm64"

        # Force everything to have the proper number of variants to maintain structure
        if len(data) != MAX_IDENTIFIER_COMPONENTS:
            data.extend(("", ""))

        data[1] = ARCHES.get((os, arch), arch)
        yield identifier, tuple(data), source


def main():
    response = httpx.get(URL)
    response.raise_for_status()

    contents = response.text
    distributions = defaultdict(list)
    ordering_data = defaultdict(dict)

    for i, distribution_type in enumerate(("DEFAULT_CPYTHON_DISTRIBUTIONS", "DEFAULT_PYPY_DISTRIBUTIONS")):
        for identifier, data, source in parse_distributions(contents, distribution_type):
            ordering_data[i][identifier] = None
            distributions[identifier].append((data, source))

    ordered = [identifier for identifiers in ordering_data.values() for identifier in reversed(identifiers)]
    output = [
        "from __future__ import annotations",
        "",
        "# fmt: off",
        "ORDERED_DISTRIBUTIONS: tuple[str, ...] = (",
    ]
    output.extend(f"    {identifier!r}," for identifier in ordered)
    output.extend((")", "DISTRIBUTIONS: dict[str, dict[tuple[str, ...], str]] = {"))

    for identifier, data in distributions.items():
        output.append(f"    {identifier!r}: {{")
        for d, source in data:
            output.extend((f"        {d!r}:", f"            {source!r},"))
        output.append("    },")

    output.extend(("}", ""))
    output = "\n".join(output)

    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        f.write(output)


if __name__ == "__main__":
    main()
