#!/usr/bin/env python3
"""Resume a public immutable artifact with independent HTTP byte ranges."""

from __future__ import annotations

import argparse
import concurrent.futures
import hashlib
import json
import os
import shutil
import time
import urllib.request
from pathlib import Path


def digest(path: Path, algorithm: str) -> str:
    value = hashlib.new(algorithm)
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            value.update(chunk)
    return value.hexdigest()


def fetch_part(url: str, path: Path, start: int, end: int, retries: int) -> dict[str, object]:
    expected = end - start + 1
    path.parent.mkdir(parents=True, exist_ok=True)
    for attempt in range(1, retries + 1):
        current = path.stat().st_size if path.exists() else 0
        if current == expected:
            return {"path": str(path), "bytes": current, "reused": True, "attempts": attempt - 1}
        if current > expected:
            raise RuntimeError(f"oversized range file {path}: {current} > {expected}")
        request_start = start + current
        request = urllib.request.Request(
            url,
            headers={
                "Range": f"bytes={request_start}-{end}",
                "User-Agent": "mcu-ir-research/1.0",
            },
        )
        try:
            with urllib.request.urlopen(request, timeout=120) as response:
                status = getattr(response, "status", None)
                content_range = response.headers.get("Content-Range", "")
                if status != 206 or not content_range.startswith(f"bytes {request_start}-{end}/"):
                    raise RuntimeError(f"unexpected range response: status={status} content-range={content_range!r}")
                with path.open("ab") as output:
                    shutil.copyfileobj(response, output, length=1024 * 1024)
        except Exception:
            if attempt == retries:
                raise
            time.sleep(min(2**attempt, 15))
    raise AssertionError("unreachable")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", required=True)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--expected-bytes", required=True, type=int)
    parser.add_argument("--expected-md5", required=True)
    parser.add_argument("--parts", type=int, default=16)
    parser.add_argument("--retries", type=int, default=5)
    args = parser.parse_args()
    output = args.output.resolve()
    if output.exists():
        actual_md5 = digest(output, "md5")
        if output.stat().st_size == args.expected_bytes and actual_md5 == args.expected_md5:
            print(json.dumps({"status": "PASS", "reused": True, "output": str(output), "bytes": output.stat().st_size, "md5": actual_md5, "sha256": digest(output, "sha256")}, sort_keys=True))
            return 0
        raise SystemExit(f"refusing to overwrite invalid existing output: {output}")

    part_dir = output.parent / f".{output.name}.parts"
    part_dir.mkdir(parents=True, exist_ok=True)
    base, remainder = divmod(args.expected_bytes, args.parts)
    ranges: list[tuple[int, Path, int, int]] = []
    offset = 0
    for index in range(args.parts):
        size = base + (1 if index < remainder else 0)
        ranges.append((index, part_dir / f"part-{index:03d}", offset, offset + size - 1))
        offset += size
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.parts) as executor:
        futures = {
            executor.submit(fetch_part, args.url, path, start, end, args.retries): index
            for index, path, start, end in ranges
        }
        for future in concurrent.futures.as_completed(futures):
            result = future.result()
            print(json.dumps({"part": futures[future], **result}, sort_keys=True), flush=True)

    temporary = output.with_suffix(output.suffix + ".assembling")
    md5 = hashlib.md5(usedforsecurity=False)
    sha256 = hashlib.sha256()
    with temporary.open("wb") as destination:
        for _, path, start, end in ranges:
            expected = end - start + 1
            if path.stat().st_size != expected:
                raise RuntimeError(f"range size mismatch for {path}")
            with path.open("rb") as source:
                for chunk in iter(lambda: source.read(1024 * 1024), b""):
                    destination.write(chunk)
                    md5.update(chunk)
                    sha256.update(chunk)
        destination.flush()
        os.fsync(destination.fileno())
    if temporary.stat().st_size != args.expected_bytes or md5.hexdigest() != args.expected_md5:
        raise RuntimeError(
            f"assembled artifact mismatch: bytes={temporary.stat().st_size} md5={md5.hexdigest()}"
        )
    os.replace(temporary, output)
    print(json.dumps({"status": "PASS", "reused": False, "output": str(output), "bytes": output.stat().st_size, "md5": md5.hexdigest(), "sha256": sha256.hexdigest()}, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())