#!/usr/bin/env python3
"""Install and smoke-test a web-scraping stack on Python 3.15, with pip and uv.

Run inside a python:3.15-slim container, once as shipped (no C compiler) and once after
`apt-get install -y --no-install-recommends gcc libc6-dev`:

    python3 lab315.py --label no-compiler --out results-no-compiler.json
    python3 lab315.py --label compiler    --out results-compiler.json

Every case gets a fresh virtual environment and no package cache. For each installed
distribution the script records the wheel tag it was installed from: a manylinux or
py3-none-any tag means a wheel from PyPI, a plain linux_x86_64 tag means pip or uv
compiled it from the source archive on this machine. Smoke tests run against a local
http.server on 127.0.0.1; nothing is fetched from the internet except packages.
"""
import argparse
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
SPECS = [
    "scrapy", "requests[socks]", "httpx[socks]", "httpx[brotli]", "httpx[zstd]", "aiohttp",
    "aiohttp[speedups]", "aiohttp-socks", "urllib3[brotli]", "curl_cffi", "playwright", "lxml",
    "beautifulsoup4", "parsel", "selectolax", "pandas", "pydantic", "orjson", "tenacity",
]
# Extra install commands: (case id, tool, extra args, specs, smoke test or None).
EXTRA = [
    ("only-binary-all", "pip", ["--only-binary", ":all:"], ["httpx[brotli]", "aiohttp[speedups]"], None),
    ("no-build", "uv", ["--no-build"], ["httpx[brotli]", "aiohttp[speedups]"], None),
    ("only-binary-brotli", "pip", ["--only-binary", "brotli"], ["scrapy"], "scrapy"),
    ("only-binary-brotli", "uv", ["--only-binary", "brotli"], ["scrapy"], "scrapy"),
    ("only-binary-brotli-pinned", "pip", ["--only-binary", "brotli"], ["scrapy==2.19.0"], None),
    ("only-binary-brotli-pinned", "uv", ["--only-binary", "brotli"], ["scrapy==2.19.0"], None),
    # brotlicffi ships an abi3 wheel; HTTPX, urllib3 and aiohttp accept it in place of brotli.
    ("httpx+brotlicffi", "pip", [], ["httpx", "brotlicffi"], "httpx[brotli]"),
    ("httpx+brotlicffi", "uv", [], ["httpx", "brotlicffi"], "httpx[brotli]"),
    ("urllib3+brotlicffi", "pip", [], ["urllib3", "brotlicffi"], "urllib3[brotli]"),
    ("aiohttp+aiodns+brotlicffi", "pip", [], ["aiohttp", "aiodns", "brotlicffi"], "aiohttp[speedups]"),
    # On 3.14+ urllib3 and aiohttp decode zstd with the standard library (compression.zstd).
    ("urllib3 (zstd)", "pip", [], ["urllib3"], "urllib3-zstd"),
    ("aiohttp (zstd)", "pip", [], ["aiohttp"], "aiohttp-zstd"),
    ("zope.interface", "pip", [], ["zope.interface"], "zope.interface"),
    ("zope.interface", "uv", [], ["zope.interface"], "zope.interface"),
]
ERROR_LINE = re.compile(r"error|failed|×|╰─>|warning|does not provide|no matching|cannot|unable", re.I)


def run(cmd, timeout=1800, env=None):
    t = time.monotonic()
    p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=env)
    return p.returncode, p.stdout, p.stderr, round(time.monotonic() - t, 1)


def venv(tool, path, python):
    if path.exists():
        shutil.rmtree(path)
    if tool == "uv":
        run(["uv", "venv", "--quiet", "--python", python, str(path)])
    else:
        run([python, "-m", "venv", str(path)])
    return str(path / "bin" / "python")


def site_packages(py):
    return Path(run([py, "-c", "import sysconfig; print(sysconfig.get_paths()['purelib'])"])[1].strip())


def inventory(py, report):
    """Every installed distribution with its wheel tag; pip's --report adds the downloaded file name."""
    downloaded = {}
    if report and report.exists():
        for item in json.loads(report.read_text()).get("install", []):
            name = item["metadata"]["name"].lower().replace("_", "-")
            downloaded[name] = item["download_info"]["url"].rsplit("/", 1)[-1]
    dists = []
    for info in sorted(site_packages(py).glob("*.dist-info")):
        meta = (info / "METADATA").read_text(errors="replace")
        name = re.search(r"^Name: (.+)$", meta, re.M).group(1).strip()
        version = re.search(r"^Version: (.+)$", meta, re.M).group(1).strip()
        wheel = (info / "WHEEL").read_text(errors="replace") if (info / "WHEEL").exists() else ""
        tags = re.findall(r"^Tag: (.+)$", wheel, re.M)
        key = name.lower().replace("_", "-").replace(".", "-")
        if key in ("pip", "setuptools"):
            continue
        if any("manylinux" in t or "musllinux" in t for t in tags):
            origin = "binary wheel"
        elif tags and all(t.endswith("-none-any") for t in tags):
            origin = "pure-Python wheel"
        elif tags:
            origin = "built from source here"
        else:
            origin = "unknown"
        file = downloaded.get(name.lower().replace("_", "-")) or downloaded.get(key)
        if file and file.endswith((".tar.gz", ".zip")):
            origin = "built from source here"
        dists.append({"name": name, "version": version, "tags": tags, "origin": origin, "file": file})
    return dists


def key_lines(text, limit=12):
    lines = [l.rstrip() for l in text.splitlines() if ERROR_LINE.search(l)]
    return lines[-limit:]


def install(tool, py, args, specs, workdir):
    report = workdir / "report.json"
    if tool == "pip":
        cmd = [py, "-m", "pip", "install", "--no-cache-dir", "--disable-pip-version-check",
               "--report", str(report), *args, *specs]
    else:
        cmd = ["uv", "pip", "install", "--no-cache", "--python", py, *args, *specs]
    rc, out, err, secs = run(cmd)
    built = sorted(set(re.findall(r"Built (\S+==\S+)", err))) if tool == "uv" else []
    shown = " ".join(cmd) if tool == "uv" else "pip " + " ".join(cmd[3:])
    return {"command": shown,
            "exit": rc, "seconds": secs, "key_lines": key_lines(out + "\n" + err),
            "uv_built": built, "output_tail": (out + err).strip().splitlines()[-40:]}, report


def smoke(py, spec):
    rc, out, err, secs = run([py, str(HERE / "smoke315.py"), spec], timeout=300)
    try:
        result = json.loads(out.strip().splitlines()[-1])
    except (IndexError, json.JSONDecodeError):
        result = {"ok": False, "error": (err or out).strip().splitlines()[-5:]}
    result["seconds"] = secs
    return result


def environment(label):
    env = {"label": label, "python": sys.version, "platform": platform.platform(),
           "compilers": {c: shutil.which(c) for c in ("cc", "gcc", "x86_64-linux-gnu-gcc")}}
    env["pip"] = run([sys.executable, "-m", "pip", "--version"])[1].strip()
    env["uv"] = run(["uv", "--version"])[1].strip() if shutil.which("uv") else None
    if Path("/etc/os-release").exists():
        env["os_release"] = re.search(r'^PRETTY_NAME="(.+)"', Path("/etc/os-release").read_text(), re.M).group(1)
    # PEP 686: UTF-8 mode on by default; check what open() picks in the C locale.
    probe = ("import sys, locale, tempfile; f = tempfile.TemporaryFile('w'); "
             "print(sys.flags.utf8_mode, locale.getpreferredencoding(False), f.encoding)")
    clean = {"PATH": os.environ.get("PATH", ""), "LANG": "C", "LC_ALL": "C"}
    env["utf8_mode_c_locale"] = run([sys.executable, "-c", probe], env=clean)[1].strip()
    env["utf8_mode_opt_out"] = run([sys.executable, "-X", "utf8=0", "-c", probe], env=clean)[1].strip()
    # The one-line check from the post: is there a wheel for this interpreter yet?
    env["wheel_check"] = {}
    for name in ("brotli", "zstandard", "zope.interface", "brotlicffi"):
        rc, out, err, _ = run([sys.executable, "-m", "pip", "download", "--disable-pip-version-check", "--no-cache-dir",
                               "--only-binary", ":all:", "--no-deps", "-d", "/tmp/wheel-check", name])
        files = sorted(os.listdir("/tmp/wheel-check")) if os.path.isdir("/tmp/wheel-check") else []
        env["wheel_check"][name] = {
            "command": f"python3 -m pip download --only-binary :all: --no-deps -d /tmp/wheel-check {name}",
            "exit": rc, "key_lines": key_lines(out + err, 3), "files": files}
        shutil.rmtree("/tmp/wheel-check", ignore_errors=True)
    return env


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--label", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--tools", default="pip,uv")
    ap.add_argument("--only", nargs="*")
    ap.add_argument("--work", default="/tmp/lab315")
    a = ap.parse_args()
    work, python = Path(a.work), sys.executable
    results = {"environment": environment(a.label), "cases": []}
    print(json.dumps(results["environment"], indent=1), flush=True)
    cases = [(spec, tool, [], [spec], spec) for tool in a.tools.split(",") for spec in SPECS]
    cases += [c for c in EXTRA if c[1] in a.tools.split(",")]
    for case_id, tool, args, specs, smoke_spec in cases:
        if a.only and case_id not in a.only:
            continue
        d = work / tool / re.sub(r"[^a-z0-9]+", "-", case_id.lower())
        py = venv(tool, d / "venv", python)
        inst, report = install(tool, py, args, specs, d)
        row = {"case": case_id, "tool": tool, "specs": specs, "install": inst}
        if inst["exit"] == 0:
            row["installed"] = inventory(py, report if tool == "pip" else None)
            if smoke_spec:
                row["smoke"] = smoke(py, smoke_spec)
        results["cases"].append(row)
        shutil.rmtree(d, ignore_errors=True)
        mark = "ok" if inst["exit"] == 0 else "FAIL"
        print(f"{a.label:12} {tool:3} {case_id:28} install={mark:4} {inst['seconds']:6}s "
              f"smoke={row.get('smoke', {}).get('ok')}", flush=True)
        Path(a.out).write_text(json.dumps(results, indent=1))


if __name__ == "__main__":
    main()
