#!/usr/bin/env python3
"""Prepare and merge non-authoritative pytest-split timing artifacts."""

from __future__ import annotations

import argparse
import io
import json
import math
import os
import re
import sys
import urllib.parse
import urllib.request
import zipfile
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any, cast

ROOT = Path(__file__).resolve().parents[2]
TIMING_SUITES = ("benchmark", "math")
MAX_BYTES = 5 * 1024 * 1024
MAX_ENTRIES = 20_000
MAX_CANDIDATES = 5
HTTP_TIMEOUT_SECONDS = 10
MAX_TIMING_HISTORY_AGE = timedelta(days=30)
"""Maximum age of a successful main-branch timing baseline."""
_SOURCE_SHA = re.compile(r"[0-9a-f]{40}(?:[0-9a-f]{24})?\Z", re.IGNORECASE)


class CrossOriginRedirectHandler(urllib.request.HTTPRedirectHandler):
    """Keep API credentials on GitHub's HTTPS origin during archive redirects."""

    def redirect_request(
        self,
        request: urllib.request.Request,
        file_pointer: Any,
        code: int,
        message: str,
        headers: Any,
        new_url: str,
    ) -> urllib.request.Request | None:
        redirected = super().redirect_request(
            request,
            file_pointer,
            code,
            message,
            headers,
            new_url,
        )
        if redirected is None:
            return None
        original = urllib.parse.urlsplit(request.full_url)
        redirected_url = urllib.parse.urlsplit(new_url)
        if (original.scheme, original.netloc) != (
            redirected_url.scheme,
            redirected_url.netloc,
        ):
            redirected.remove_header("Authorization")
        return redirected


def shard_count(suite: str) -> int:
    if suite not in TIMING_SUITES:
        raise ValueError(f"unsupported timing suite: {suite}")
    return 4


def artifact_name(suite: str) -> str:
    return f"{suite}-test-durations"


def artifact_file(suite: str) -> str:
    return f"{suite}-test-durations.json"


def pytest_split_version() -> str:
    """Return the pinned pytest-split version from pyproject.toml."""

    pyproject = ROOT / "pyproject.toml"
    for line in pyproject.read_text(encoding="utf-8").splitlines():
        stripped = line.strip()
        if stripped.startswith(('"pytest-split==', "'pytest-split==")):
            version = stripped.split("==", 1)[1].rstrip("\",'")
            if version:
                return version
    raise ValueError("pytest-split pin not found in pyproject.toml")


def warning(message: str) -> None:
    print(f"::warning::{message}", file=sys.stderr)


def load_json_bytes(payload: bytes, source: str) -> Any:
    if len(payload) > MAX_BYTES:
        raise ValueError(f"{source} exceeds {MAX_BYTES} bytes")
    return json.loads(payload)


def validate_durations(value: Any, source: str, suite: str) -> dict[str, float]:
    if not isinstance(value, dict):
        raise ValueError(f"{source} must contain a JSON object")
    if len(value) > MAX_ENTRIES:
        raise ValueError(f"{source} exceeds {MAX_ENTRIES} timing entries")

    durations: dict[str, float] = {}
    prefix = "benchmarks/validation/" if suite == "benchmark" else f"tests/{suite}/"
    for nodeid, duration in value.items():
        if not isinstance(nodeid, str) or not nodeid.startswith(prefix):
            raise ValueError(f"{source} contains an invalid test node id: {nodeid!r}")
        if isinstance(duration, bool) or not isinstance(duration, (int, float)):
            raise ValueError(f"{source} contains a non-numeric duration for {nodeid}")
        seconds = float(duration)
        if not math.isfinite(seconds) or seconds < 0:
            raise ValueError(f"{source} contains an invalid duration for {nodeid}")
        durations[nodeid] = seconds
    return durations


def api_json(url: str, token: str) -> Any:
    request = urllib.request.Request(
        url,
        headers={
            "Accept": "application/vnd.github+json",
            "Authorization": f"Bearer {token}",
            "X-GitHub-Api-Version": "2022-11-28",
        },
    )
    with urllib.request.urlopen(request, timeout=HTTP_TIMEOUT_SECONDS) as response:
        return load_json_bytes(response.read(MAX_BYTES + 1), url)


def download(url: str, token: str) -> bytes:
    request = urllib.request.Request(
        url,
        headers={
            "Accept": "application/vnd.github+json",
            "Authorization": f"Bearer {token}",
            "X-GitHub-Api-Version": "2022-11-28",
        },
    )
    opener = urllib.request.build_opener(CrossOriginRedirectHandler())
    with opener.open(request, timeout=HTTP_TIMEOUT_SECONDS) as response:
        return cast(bytes, response.read(MAX_BYTES + 1))


def write_durations(path: Path, durations: dict[str, float]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(durations, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )


def _timestamp(value: Any, label: str) -> datetime:
    if not isinstance(value, str):
        raise ValueError(f"{label} must be an RFC3339 timestamp")
    try:
        parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError as exc:
        raise ValueError(f"{label} must be an RFC3339 timestamp") from exc
    if parsed.tzinfo is None:
        raise ValueError(f"{label} must include a timezone")
    return parsed.astimezone(UTC)


def _require_recent_timestamp(value: Any, label: str, *, now: datetime) -> None:
    timestamp = _timestamp(value, label)
    if timestamp > now:
        raise ValueError(f"{label} is in the future")
    if now - timestamp > MAX_TIMING_HISTORY_AGE:
        raise ValueError(
            f"{label} exceeds the {MAX_TIMING_HISTORY_AGE.days}-day timing-history age"
        )


def artifact_durations(
    artifact: dict[str, Any],
    *,
    api: str,
    repository: str,
    token: str,
    suite: str,
    now: datetime | None = None,
) -> dict[str, float] | None:
    now = now or datetime.now(UTC)
    run = api_json(
        f"{api}/repos/{repository}/actions/runs/{artifact['workflow_run']['id']}",
        token,
    )
    if run.get("head_branch") != "main" or run.get("conclusion") != "success":
        return None
    _require_recent_timestamp(
        artifact.get("created_at"), "timing artifact creation", now=now
    )

    archive = download(artifact["archive_download_url"], token)
    if len(archive) > MAX_BYTES:
        raise ValueError("timing artifact archive is too large")
    filename = artifact_file(suite)
    with zipfile.ZipFile(io.BytesIO(archive)) as bundle:
        member = bundle.getinfo(filename)
        if member.file_size > MAX_BYTES:
            raise ValueError("timing artifact payload is too large")
        payload = load_json_bytes(bundle.read(member), filename)
    if not isinstance(payload, dict):
        raise ValueError("timing artifact payload must be a JSON object")
    if (
        payload.get("version") != 1
        or payload.get("suite") != suite
        or payload.get("shard_count") != shard_count(suite)
    ):
        raise ValueError("timing artifact metadata is incompatible")
    _require_recent_timestamp(
        payload.get("generated_at"), "timing artifact generation", now=now
    )
    source_sha = payload.get("source_sha")
    if not isinstance(source_sha, str) or _SOURCE_SHA.fullmatch(source_sha) is None:
        raise ValueError("timing artifact source SHA is invalid")
    if source_sha != run.get("head_sha"):
        raise ValueError("timing artifact source SHA does not match its workflow run")
    return validate_durations(payload.get("durations"), filename, suite)


def prepare(output: Path, suite: str) -> None:
    try:
        repository = os.environ["GITHUB_REPOSITORY"]
        token = os.environ["GH_TOKEN"]
        api = os.environ.get("GITHUB_API_URL", "https://api.github.com")
        listing = api_json(
            f"{api}/repos/{repository}/actions/artifacts"
            f"?name={artifact_name(suite)}&per_page={MAX_CANDIDATES}",
            token,
        )
        artifacts = sorted(
            [artifact for artifact in listing["artifacts"] if not artifact["expired"]],
            key=lambda candidate: candidate["created_at"],
            reverse=True,
        )[:MAX_CANDIDATES]
        durations = None
        for artifact in artifacts:
            try:
                durations = artifact_durations(
                    artifact,
                    api=api,
                    repository=repository,
                    token=token,
                    suite=suite,
                )
            except (
                IndexError,
                KeyError,
                OSError,
                TypeError,
                ValueError,
                zipfile.BadZipFile,
            ) as exc:
                artifact_id = artifact.get("id", "unknown")
                warning(
                    f"rejected {suite} timing artifact "
                    f"{artifact_id}: {type(exc).__name__}"
                )
                continue
            if durations is not None:
                break
        if durations is None:
            raise ValueError(
                f"no valid successful main-branch {suite} timing artifact is available"
            )
    except (
        IndexError,
        KeyError,
        OSError,
        TypeError,
        ValueError,
        zipfile.BadZipFile,
    ) as exc:
        warning(f"{exc}; {suite} shards will use equal weighting")
        durations = {}
    write_durations(output, durations)


def merge(
    inputs: list[Path],
    output: Path,
    source_sha: str,
    python_version: str,
    pytest_split_version: str,
    suite: str,
) -> None:
    count = shard_count(suite)
    if len(inputs) != count:
        raise ValueError(f"expected {count} shard timing files, received {len(inputs)}")
    merged: dict[str, float] = {}
    for path in inputs:
        payload = load_json_bytes(path.read_bytes(), str(path))
        durations = validate_durations(payload, str(path), suite)
        duplicates = merged.keys() & durations.keys()
        for nodeid in sorted(duplicates):
            previous = merged[nodeid]
            incoming = durations[nodeid]
            if incoming > previous:
                merged[nodeid] = incoming
            warning(
                f"duplicate timing entry across shards for {nodeid}; "
                f"keeping max({previous}, {incoming})"
            )
        merged.update({k: v for k, v in durations.items() if k not in duplicates})

    envelope = {
        "version": 1,
        "suite": suite,
        "source_sha": source_sha,
        "generated_at": datetime.now(UTC).isoformat(),
        "python_version": python_version,
        "shard_count": count,
        "pytest_split_version": pytest_split_version,
        "durations": dict(sorted(merged.items())),
    }
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(envelope, indent=2) + "\n", encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    prepare_parser = subparsers.add_parser("prepare")
    prepare_parser.add_argument("--output", type=Path, required=True)
    prepare_parser.add_argument("--suite", choices=TIMING_SUITES, required=True)

    merge_parser = subparsers.add_parser("merge")
    merge_parser.add_argument("--input", action="append", type=Path, required=True)
    merge_parser.add_argument("--output", type=Path, required=True)
    merge_parser.add_argument("--source-sha", required=True)
    merge_parser.add_argument("--python-version", required=True)
    merge_parser.add_argument("--suite", choices=TIMING_SUITES, required=True)
    merge_parser.add_argument(
        "--pytest-split-version",
        default=None,
        help="defaults to the pytest-split pin in pyproject.toml",
    )

    args = parser.parse_args()
    if args.command == "prepare":
        prepare(args.output, args.suite)
    else:
        split_version = args.pytest_split_version or pytest_split_version()
        merge(
            args.input,
            args.output,
            args.source_sha,
            args.python_version,
            split_version,
            args.suite,
        )


if __name__ == "__main__":
    main()
