# SPDX-License-Identifier: GPL-3.0-only
"""Download an approved Atlas collection page-by-page without loading it all in RAM.

python3 atlas_download.py --credentials credentials.json --collection trees --output trees.geojson
Uses only the Python standard library. Credentials must be kept private.
"""

import argparse
import base64
import json
import os
import sys
import tempfile
import time
from pathlib import Path
from urllib.error import HTTPError
from urllib.parse import urlencode, urljoin, urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener


class NoRedirects(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None  # Never forward credentials through redirects.


class Atlas:
    def __init__(self, credentials):
        self.endpoint = credentials["token_endpoint"]
        parsed = urlsplit(self.endpoint)
        if parsed.scheme != "https" and not (
            parsed.scheme == "http"
            and parsed.hostname in ("localhost", "127.0.0.1", "::1")
        ):
            raise ValueError("HTTPS is required except for local development")
        if parsed.username or parsed.password or parsed.query or parsed.fragment:
            raise ValueError("Invalid token endpoint")
        if parsed.path != "/partner/v1/oauth/token":
            raise ValueError("Unexpected token endpoint path")
        self.origin = f"{parsed.scheme}://{parsed.netloc}"
        self.credentials = credentials
        self.token = None
        self.expires = 0
        self.opener = build_opener(NoRedirects())

    def json_request(self, url, headers, data=None):
        for _ in range(8):
            try:
                with self.opener.open(
                    Request(url, data=data, headers=headers), timeout=45
                ) as response:
                    payload = response.read(32 * 1024 * 1024 + 1)
                    if len(payload) > 32 * 1024 * 1024:
                        raise ValueError("Response too large; request a smaller page")
                    return json.loads(payload)
            except HTTPError as exc:
                if exc.code != 429:
                    raise RuntimeError(
                        f"Atlas returned HTTP {exc.code}; check access policy"
                    ) from None
                wait = min(86400, max(1, int(exc.headers.get("Retry-After", "60"))))
                print(f"Quota reached; waiting {wait}s", file=sys.stderr)
                time.sleep(wait)
        raise RuntimeError("Repeated quota responses; try again later")

    def refresh(self):
        value = f"{self.credentials['client_id']}:{self.credentials['client_secret']}"
        basic = base64.b64encode(value.encode()).decode()
        result = self.json_request(
            self.endpoint,
            {
                "Authorization": f"Basic {basic}",
                "Content-Type": "application/x-www-form-urlencoded",
            },
            b"grant_type=client_credentials",
        )
        self.token = result["access_token"]
        self.expires = time.monotonic() + result["expires_in"] - 30

    def page(self, path, collection):
        url = urljoin(self.origin, path)
        parsed = urlsplit(url)
        if (
            f"{parsed.scheme}://{parsed.netloc}" != self.origin
            or parsed.path != f"/partner/v1/collections/{collection}/items"
        ):
            raise ValueError(
                "Refusing to send a token outside the approved collection endpoint"
            )
        for attempt in range(2):
            if not self.token or time.monotonic() >= self.expires:
                self.refresh()
            try:
                return self.json_request(url, {"Authorization": f"Bearer {self.token}"})
            except RuntimeError as exc:
                if "HTTP 401" not in str(exc) or attempt:
                    raise
                self.token = None
        raise RuntimeError("Unable to authenticate")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--credentials", required=True, type=Path)
    parser.add_argument("--collection", choices=["trees", "green-areas"], required=True)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--limit", type=int, default=100)
    parser.add_argument("--location-id")
    parser.add_argument("--bbox")
    args = parser.parse_args()
    if not 1 <= args.limit <= 1000:
        parser.error("--limit must be between 1 and 1000")
    atlas = Atlas(json.loads(args.credentials.read_text()))
    params = {"limit": args.limit}
    if args.location_id:
        params["location_id"] = args.location_id
    if args.bbox:
        params["bbox"] = args.bbox
    path = f"/partner/v1/collections/{args.collection}/items?{urlencode(params)}"
    count = 0
    args.output.parent.mkdir(parents=True, exist_ok=True)
    if args.output.exists():
        parser.error("Output already exists; choose another path")
    descriptor, temp = tempfile.mkstemp(
        prefix=".atlas-download-", dir=args.output.parent
    )
    try:
        with os.fdopen(descriptor, "w") as stream:
            stream.write('{"type":"FeatureCollection","features":[')
            first = True
            while path:
                page = atlas.page(path, args.collection)
                for feature in page["features"]:
                    if not first:
                        stream.write(",")
                    json.dump(
                        feature, stream, ensure_ascii=False, separators=(",", ":")
                    )
                    first = False
                    count += 1
                path = next(
                    (
                        link["href"]
                        for link in page.get("links", [])
                        if link.get("rel") == "next"
                    ),
                    None,
                )
                print(f"Downloaded {count} features", file=sys.stderr)
            stream.write("]}")
        os.link(
            temp, args.output
        )  # Fail rather than overwrite a concurrently created file.
        os.unlink(temp)
    finally:
        if os.path.exists(temp):
            os.unlink(temp)
    print(f"Saved {count} features to {args.output}", file=sys.stderr)


if __name__ == "__main__":
    main()