"""Download the Argoverse 2 sensor files we actually use, over plain HTTPS.

The AV2 bucket is public and listable without credentials, so this needs no
s5cmd and no AWS CLI. A full log is several GB of seven cameras, LiDAR and maps;
we read one camera, so per log this fetches four things:

    annotations.feather
    calibration/intrinsics.feather
    calibration/egovehicle_SE3_sensor.feather
    sensors/cameras/<camera>/*.jpg        (~96 MB)

Annotations are downloaded first for every candidate log, because they are 0.29 MB
each and tell us which logs contain cones, barrels or stop signs at all. Images
are then fetched only for the logs worth having.

Usage:
    python scripts/fetch_av2.py --dataroot data/raw/av2 --scan 40 --keep 15
"""

from __future__ import annotations

import argparse
import os
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

BUCKET = "https://s3.amazonaws.com/argoverse"
NAMESPACE = "{http://s3.amazonaws.com/doc/2006-03-01/}"

# Matches AV2_CLASS_MAP in src/data/extract_av2.py.
WANTED_CATEGORIES = ("CONSTRUCTION_CONE", "CONSTRUCTION_BARREL", "STOP_SIGN")


def list_keys(prefix: str, delimiter: str = "") -> tuple[list[str], list[str]]:
    """Page through a public S3 listing. Returns (keys, common prefixes)."""
    keys: list[str] = []
    prefixes: list[str] = []
    token = ""
    while True:
        url = f"{BUCKET}/?list-type=2&prefix={prefix}&max-keys=1000"
        if delimiter:
            url += f"&delimiter={delimiter}"
        if token:
            url += f"&continuation-token={urllib.parse.quote(token, safe='')}"
        root = with_retries(
            lambda: ET.fromstring(urllib.request.urlopen(url, timeout=60).read()),
            f"list {prefix}",
        )

        keys += [e.text for e in root.iter(f"{NAMESPACE}Key")]
        prefixes += [e.findtext(f"{NAMESPACE}Prefix")
                     for e in root.iter(f"{NAMESPACE}CommonPrefixes")]
        if root.findtext(f"{NAMESPACE}IsTruncated") != "true":
            return keys, prefixes
        token = root.findtext(f"{NAMESPACE}NextContinuationToken")


# Thousands of requests over a home connection will hit transient DNS and
# connection failures. Those are routine, not exceptional, so retry rather than
# letting one of them abort a download that is otherwise 95% done.
RETRIES = 4


def with_retries(operation, what: str):
    for attempt in range(RETRIES):
        try:
            return operation()
        except (urllib.error.URLError, TimeoutError, ConnectionError) as error:
            if attempt == RETRIES - 1:
                raise
            delay = 2 ** attempt
            print(f"    retry {attempt + 1}/{RETRIES - 1} in {delay}s ({what}): {error}",
                  flush=True)
            time.sleep(delay)


def download(key: str, dest: Path) -> None:
    if dest.exists():
        return
    dest.parent.mkdir(parents=True, exist_ok=True)
    # Unique per call. A shared "<name>.part" is renamed out from under whichever
    # of two overlapping runs loses the race, which surfaces as a confusing
    # FileNotFoundError on rename rather than as the collision it is.
    temporary = dest.with_suffix(f"{dest.suffix}.{os.getpid()}.{threading.get_ident()}.part")

    def fetch():
        with urllib.request.urlopen(f"{BUCKET}/{key}", timeout=120) as response, \
                open(temporary, "wb") as handle:
            while chunk := response.read(1 << 20):
                handle.write(chunk)

    with_retries(fetch, key.rsplit("/", 1)[-1])
    temporary.replace(dest)


def fetch_metadata(log_prefix: str, log_dir: Path) -> None:
    """The small files: annotations and both calibration tables."""
    for relative in ("annotations.feather",
                     "calibration/intrinsics.feather",
                     "calibration/egovehicle_SE3_sensor.feather"):
        download(log_prefix + relative, log_dir / relative)


def count_wanted(log_dir: Path) -> dict[str, int]:
    import pandas as pd

    annotations = pd.read_feather(log_dir / "annotations.feather")
    counts = annotations["category"].value_counts()
    return {c: int(counts.get(c, 0)) for c in WANTED_CATEGORIES}


def fetch_images(log_prefix: str, log_dir: Path, camera: str, workers: int) -> int:
    prefix = f"{log_prefix}sensors/cameras/{camera}/"
    keys, _ = list_keys(prefix)

    def fetch_one(key: str) -> bool:
        try:
            download(key, log_dir / key[len(log_prefix):])
            return True
        except Exception as error:               # noqa: BLE001 -- report, do not abort
            print(f"    skipped {key.rsplit('/', 1)[-1]}: {error}", flush=True)
            return False

    with ThreadPoolExecutor(max_workers=workers) as pool:
        return sum(pool.map(fetch_one, keys))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--dataroot", type=Path, default=Path("data/raw/av2"))
    parser.add_argument("--split", default="val", choices=["train", "val", "test"])
    parser.add_argument("--camera", default="ring_front_center")
    parser.add_argument("--scan", type=int, default=40,
                        help="how many logs to read annotations for")
    parser.add_argument("--keep", type=int, default=15,
                        help="how many of those logs to fetch images for")
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--prioritise", default=None,
                        help="rank logs by this category alone rather than by "
                             "the total, e.g. STOP_SIGN")
    parser.add_argument("--min-count", type=int, default=1,
                        help="skip logs with fewer than this many of the "
                             "prioritised category")
    args = parser.parse_args()

    split_prefix = f"datasets/av2/sensor/{args.split}/"
    _, log_prefixes = list_keys(split_prefix, delimiter="/")
    log_prefixes = sorted(log_prefixes)[:args.scan]
    print(f"scanning {len(log_prefixes)} logs in {args.split}")

    scored: list[tuple[int, str, Path, dict]] = []
    for index, log_prefix in enumerate(log_prefixes, start=1):
        log_dir = args.dataroot / log_prefix[len(split_prefix):].rstrip("/")
        fetch_metadata(log_prefix, log_dir)
        counts = count_wanted(log_dir)
        score = counts[args.prioritise] if args.prioritise else sum(counts.values())
        scored.append((score, log_prefix, log_dir, counts))
        print(f"  [{index}/{len(log_prefixes)}] {log_dir.name[:8]}  "
              f"cone={counts['CONSTRUCTION_CONE']:5d} "
              f"barrel={counts['CONSTRUCTION_BARREL']:5d} "
              f"stop={counts['STOP_SIGN']:5d}", flush=True)

    scored.sort(reverse=True, key=lambda row: row[0])
    chosen = [row for row in scored if row[0] >= args.min_count][:args.keep]
    print(f"\n{len(chosen)} logs contain our classes; fetching images for them")

    total_images = 0
    for index, (score, log_prefix, log_dir, counts) in enumerate(chosen, start=1):
        count = fetch_images(log_prefix, log_dir, args.camera, args.workers)
        total_images += count
        print(f"  [{index}/{len(chosen)}] {log_dir.name[:8]}  {count} images  "
              f"({score} objects)", flush=True)

    # Logs with none of our classes keep only their metadata; drop them so the
    # extractor does not walk directories that have no images.
    for score, _, log_dir, _ in scored:
        if score == 0:
            for path in sorted(log_dir.rglob("*"), reverse=True):
                path.unlink() if path.is_file() else path.rmdir()
            log_dir.rmdir()

    size = sum(p.stat().st_size for p in args.dataroot.rglob("*") if p.is_file())
    print(f"\n{len(chosen)} logs, {total_images} images, {size / 1e6:.0f} MB "
          f"in {args.dataroot}")


if __name__ == "__main__":
    main()