#!/usr/bin/env python3
"""fgsync.py - fetch FlightGear TerraSync scenery for a region with aria2c.

TerraSync serves one .dirindex per directory listing its files and
subdirectories with SHA1 hashes; aria2c cannot follow that by itself, so this
script walks the indexes, builds a download list restricted to a bounding box,
and hands it to aria2c with resume enabled.  Re-running it only fetches what is
missing or has changed, so an interrupted download continues where it stopped.

  ./fgsync.py --lat 48.11 --lon 16.57 --radius 1 --target ~/.fgfs/TerraSync

Mirrors are discovered from the NAPTR records of terrasync.flightgear.org when
a resolver is available (systemd-resolved, dig or host); otherwise the built-in
list below is probed.  Every mirror that answers is kept: index reads fall
through to the next one when a mirror stops responding, and each download gets
the file's URL on all of them, so aria2c can switch mirrors on its own.
--url overrides the discovery and pins a single mirror.
"""

import argparse
import hashlib
import math
import os
import re
import shutil
import signal
import subprocess
import sys
import urllib.error
import urllib.request

# Fallbacks for when the NAPTR lookup fails - which is the usual reason for
# "no DNS entry found for 'terrasync.flightgear.org'" on mobile connections.
FALLBACK_MIRRORS = [
    "https://terrasync.eti.pg.gda.pl/ws2",
    "https://us1mirror.flightgear.org/terrasync/ws2",
    "https://flightgear.sourceforge.net/scenery",
]

# Top-level directories worth fetching for flying.  Models holds the shared
# object models referenced by the .stg files and is not tile-organised.
DEFAULT_SUBTREES = ["Terrain", "Objects", "Airports", "Models"]

TILE_RE = re.compile(r"^([ew])(\d{3})([ns])(\d{2})$")


def log(msg):
    print(msg, file=sys.stderr, flush=True)


# --------------------------------------------------------------- tile naming

def bucket_name(lon, lat):
    """10x10 degree bucket directory, e.g. e010n40."""
    blon = int((lon // 10) * 10)
    blat = int((lat // 10) * 10)
    return "%s%03d%s%02d" % ("w" if blon < 0 else "e", abs(blon),
                             "s" if blat < 0 else "n", abs(blat))


def tile_name(lon, lat):
    """1x1 degree tile directory, e.g. e016n48."""
    tlon = int(lon // 1)
    tlat = int(lat // 1)
    return "%s%03d%s%02d" % ("w" if tlon < 0 else "e", abs(tlon),
                             "s" if tlat < 0 else "n", abs(tlat))


# ----------------------------------------------------------- finished regions
#
# A run that read every index it asked for and downloaded everything it
# wanted appends its bounding box here.  fgfs-run asks with --check before a
# start whether the airport is covered, and only then goes to the mirrors -
# so a start with the scenery already there costs no network round trip
# and works offline.  A run that lost a mirror halfway (walk() logs and goes
# on) or whose aria2c failed records nothing, and is repeated next time.

REGIONS_FILE = ".fgsync-regions"


def load_regions(target):
    """[(left, bottom, right, top, set(subtrees)), ...] recorded so far."""
    out = []
    try:
        with open(os.path.join(target, REGIONS_FILE)) as fh:
            for line in fh:
                parts = line.split()
                if len(parts) < 5:
                    continue
                try:
                    left, bottom, right, top = (float(v) for v in parts[:4])
                except ValueError:
                    continue
                out.append((left, bottom, right, top, set(parts[4].split(","))))
    except OSError:
        pass
    return out


def record_region(target, left, bottom, right, top, subtrees):
    """Four decimals, each edge rounded outward.  Rounded to nearest, an
    edge could land inside the box that was actually fetched, and a --check
    with the same radius - which is what fgfs-run asks - then failed
    immediately after a successful download."""
    with open(os.path.join(target, REGIONS_FILE), "a") as fh:
        fh.write("%.4f %.4f %.4f %.4f %s\n"
                 % (math.floor(left * 10000) / 10000.0,
                    math.floor(bottom * 10000) / 10000.0,
                    math.ceil(right * 10000) / 10000.0,
                    math.ceil(top * 10000) / 10000.0,
                    ",".join(sorted(subtrees))))


def tiles_on_disk(target, left, bottom, right, top):
    """True when every Terrain tile of the box is present.  The record file
    alone is not proof: it sits next to the data, and someone freeing space
    deletes the tiles, not the dotfile."""
    buckets, tiles = wanted_dirs(left, bottom, right, top)
    terrain = os.path.join(target, "Terrain")
    if not os.path.isdir(terrain):
        return False
    for tile in tiles:
        # the bucket the tile lives in, back from its name
        lon = (1 if tile[0] == "e" else -1) * int(tile[1:4])
        lat = (1 if tile[4] == "n" else -1) * int(tile[5:7])
        bucket = bucket_name(lon, lat)
        if not os.path.isdir(os.path.join(terrain, bucket, tile)):
            return False
    return True


def is_covered(regions, left, bottom, right, top, subtrees):
    """True when one recorded region contains the whole box with all the
    subtrees asked for.  Boxes straddling the date line are not merged;
    a region fetched there is simply asked for again."""
    for rl, rb, rr, rt, rs in regions:
        if rl <= left and rb <= bottom and rr >= right and rt >= top \
                and set(subtrees) <= rs:
            return True
    return False


def wanted_dirs(left, bottom, right, top):
    """The bucket and tile directory names covering the bounding box.

    Stepping by 1.0 from the left edge lost the eastern column whenever the
    accumulated float landed one ulp above the right edge (it happens for
    longitudes just below a power of two, e.g. -0.45 - 1), and the region
    was recorded as fetched all the same.  Integers cannot do that.
    Longitudes are wrapped, so a box across the date line asks for the
    tiles that exist rather than for w181."""
    buckets, tiles = set(), set()
    for ilat in range(int(math.floor(bottom)), int(math.floor(top)) + 1):
        lat = max(-90.0, min(89.0, float(ilat)))
        for ilon in range(int(math.floor(left)), int(math.floor(right)) + 1):
            lon = ((ilon + 180.0) % 360.0) - 180.0
            buckets.add(bucket_name(lon, lat))
            tiles.add(tile_name(lon, lat))
    return buckets, tiles


# ------------------------------------------------------------------ mirrors

def naptr_mirrors():
    """Ask the system resolver for the TerraSync mirror list."""
    probes = [
        ["resolvectl", "query", "--type=NAPTR", "terrasync.flightgear.org"],
        ["systemd-resolve", "--type=NAPTR", "terrasync.flightgear.org"],
        ["dig", "+short", "NAPTR", "terrasync.flightgear.org"],
        ["host", "-t", "NAPTR", "terrasync.flightgear.org"],
    ]
    found = []
    for cmd in probes:
        if not shutil.which(cmd[0]):
            continue
        try:
            out = subprocess.run(cmd, capture_output=True, text=True,
                                 timeout=20).stdout
        except Exception:
            continue
        for url in re.findall(r'https?://[^\s"\']+', out):
            url = url.rstrip('".').rstrip("/")
            if url not in found:
                found.append(url)
        if found:
            log("NAPTR-Spiegel gefunden: %d" % len(found))
            return found
    log("keine NAPTR-Antwort - benutze die eingebaute Spiegelliste")
    return []


def fetch(url, timeout=30):
    req = urllib.request.Request(url, headers={"User-Agent": "fgsync.py"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.read()


def pick_mirrors(candidates):
    """Every mirror that serves a usable root .dirindex, in the order tried."""
    good = []
    for base in candidates:
        if base in good:
            continue
        try:
            data = fetch(base + "/.dirindex", timeout=20)
        except Exception as exc:
            log("  %s -> %s" % (base, exc))
            continue
        if b"version:" in data or b"d:" in data:
            log("  %s -> ok" % base)
            good.append(base)
        else:
            log("  %s -> unerwarteter Inhalt" % base)
    return good


def fetch_any(mirrors, rel, timeout=30):
    """Read rel from the first mirror that answers.  A failing mirror moves to
    the back of the list, so the next call does not wait on it again."""
    last = None
    for _ in range(len(mirrors)):
        base = mirrors[0]
        try:
            return fetch("%s/%s" % (base, rel), timeout=timeout)
        except urllib.error.HTTPError as exc:
            # Not necessarily absent: a mirror that lags behind serves 404 for
            # a directory the others have.  Only when every mirror says 404
            # is it really gone - and walk() counts that as a failure, so an
            # index that was asked for but never read cannot end up recorded
            # as fetched.
            last = exc
        except Exception as exc:
            last = exc
        if len(mirrors) > 1:
            log("  Spiegel %s antwortet nicht (%s), wechsle" % (base, last))
            mirrors.append(mirrors.pop(0))
    raise last if last is not None else RuntimeError("kein Spiegel erreichbar")


# ----------------------------------------------------------------- dirindex

def parse_dirindex(data):
    """Return (dirs, files) from a .dirindex.

    Lines are colon separated: 'd:<name>:<sha1>' for subdirectories,
    'f:<name>:<sha1>[:<size>]' for files; 'version:', 'path:' and 't:' are
    metadata.  Unknown line types are ignored so a format extension does not
    break the walk.
    """
    dirs, files = [], []
    for raw in data.decode("utf-8", "replace").splitlines():
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        parts = line.split(":")
        if parts[0] == "d" and len(parts) >= 3:
            dirs.append((parts[1], parts[2]))
        elif parts[0] == "f" and len(parts) >= 3:
            size = int(parts[3]) if len(parts) > 3 and parts[3].isdigit() else None
            files.append((parts[1], parts[2], size))
    return dirs, files


def sha1_of(path):
    h = hashlib.sha1()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def keep_dir(name, depth, subtree, buckets, tiles):
    """Restrict the walk to the requested region.

    Only Terrain/Objects are tile organised: depth 1 is the 10x10 bucket,
    depth 2 the 1x1 tile.  Airports and Models are not filtered by position.
    """
    if subtree not in ("Terrain", "Objects"):
        return True
    if depth == 1:
        return name in buckets
    if depth == 2:
        return name in tiles
    return True


def walk(mirrors, target, subtree, buckets, tiles, want, stats, path=None, depth=0):
    """Recursively read .dirindex files and collect what needs downloading."""
    rel = subtree if path is None else path
    local_dir = os.path.join(target, *rel.split("/"))
    try:
        data = fetch_any(mirrors, "%s/.dirindex" % rel)
    except Exception as exc:
        # Every directory walked here is either a subtree root or one its
        # parent index listed, so not reading it means the region was not
        # covered - whatever the reason.
        log("  %s: %s" % (rel, exc))
        stats["failed"] += 1
        return

    stats["indexes"] += 1
    if stats["indexes"] % 25 == 0:
        log("  %d Indizes gelesen, %d Dateien vorgemerkt" % (stats["indexes"], len(want)))

    os.makedirs(local_dir, exist_ok=True)
    index_path = os.path.join(local_dir, ".dirindex")
    with open(index_path, "wb") as fh:
        fh.write(data)

    dirs, files = parse_dirindex(data)
    for name, sha, size in files:
        dest = os.path.join(local_dir, name)
        if os.path.exists(dest):
            try:
                if sha1_of(dest) == sha:
                    stats["current"] += 1
                    continue
            except OSError:
                pass
        want.append(([("%s/%s/%s" % (m, rel, name)) for m in mirrors], local_dir, name))

    for name, sha in dirs:
        if not keep_dir(name, depth + 1, subtree, buckets, tiles):
            continue
        walk(mirrors, target, subtree, buckets, tiles, want, stats,
             "%s/%s" % (rel, name), depth + 1)


# --------------------------------------------------------------------- main

def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--target", default=os.path.expanduser("~/.fgfs/TerraSync"),
                    help="scenery directory (default: ~/.fgfs/TerraSync)")
    ap.add_argument("--lat", type=float, help="centre latitude")
    ap.add_argument("--lon", type=float, help="centre longitude")
    ap.add_argument("--radius", type=float, default=1.0,
                    help="degrees around the centre (default: 1)")
    ap.add_argument("--left", type=float, help="bounding box west edge")
    ap.add_argument("--right", type=float, help="bounding box east edge")
    ap.add_argument("--bottom", type=float, help="bounding box south edge")
    ap.add_argument("--top", type=float, help="bounding box north edge")
    ap.add_argument("--url", help="scenery root URL, skips mirror discovery")
    ap.add_argument("--subtrees", default=",".join(DEFAULT_SUBTREES),
                    help="comma separated top level directories")
    ap.add_argument("--connections", type=int, default=4,
                    help="aria2c connections per server (default: 4)")
    ap.add_argument("--concurrent", type=int, default=8,
                    help="aria2c parallel downloads (default: 8)")
    ap.add_argument("--list-only", action="store_true",
                    help="only write the aria2c input file, do not download")
    ap.add_argument("--check", action="store_true",
                    help="do not download; exit 0 if a finished run covers "
                         "--margin degrees around --lat/--lon, 2 if the tiles "
                         "are on disk but no finished run is on record, else 1")
    ap.add_argument("--margin", type=float, default=0.5,
                    help="degrees around the centre that --check requires (default: 0.5)")
    args = ap.parse_args()

    subtrees = [s.strip() for s in args.subtrees.split(",") if s.strip()]

    if args.check:
        if args.lat is None or args.lon is None:
            ap.error("--check needs --lat and --lon")
        m = min(args.margin, args.radius)
        regions = load_regions(args.target)
        box = (args.lon - m, args.lat - m, args.lon + m, args.lat + m)
        on_disk = tiles_on_disk(args.target, box[0], box[1], box[2], box[3])
        if is_covered(regions, box[0], box[1], box[2], box[3], subtrees):
            if not on_disk:
                log("Szenerie um %.2f,%.2f ist vermerkt, aber die Kacheln fehlen"
                    % (args.lat, args.lon))
                return 1
            log("Szenerie um %.2f,%.2f vorhanden (%d fertige Regionen)"
                % (args.lat, args.lon, len(regions)))
            return 0
        if on_disk:
            # Fetched before this record existed (or the record was lost):
            # the tiles are there, so flying is possible right away; whether
            # they are complete only a run against the servers can tell.
            log("Szenerie um %.2f,%.2f auf der Platte, aber nicht als fertig vermerkt"
                % (args.lat, args.lon))
            return 2
        log("Szenerie um %.2f,%.2f fehlt (%d fertige Regionen)"
            % (args.lat, args.lon, len(regions)))
        return 1

    if args.lat is not None and args.lon is not None:
        left, right = args.lon - args.radius, args.lon + args.radius
        bottom, top = args.lat - args.radius, args.lat + args.radius
    elif None not in (args.left, args.right, args.bottom, args.top):
        left, right, bottom, top = args.left, args.right, args.bottom, args.top
    else:
        ap.error("either --lat/--lon or --left/--right/--bottom/--top")

    buckets, tiles = wanted_dirs(left, bottom, right, top)
    log("Region %.2f..%.2f x %.2f..%.2f -> %d Kacheln in %d Buckets"
        % (left, right, bottom, top, len(tiles), len(buckets)))
    log("Kacheln: %s" % " ".join(sorted(tiles)))

    if args.url:
        mirrors = [args.url.rstrip("/")]
    else:
        mirrors = pick_mirrors(naptr_mirrors() + FALLBACK_MIRRORS)
    if not mirrors:
        log("kein erreichbarer Spiegel - mit --url einen angeben")
        return 2
    log("Spiegel (%d): %s" % (len(mirrors), ", ".join(mirrors)))

    os.makedirs(args.target, exist_ok=True)
    want = []
    stats = {"indexes": 0, "current": 0, "failed": 0}
    for subtree in subtrees:
        log("Durchsuche %s ..." % subtree)
        before = stats["indexes"]
        walk(mirrors, args.target, subtree, buckets, tiles, want, stats)
        if stats["indexes"] == before:
            log("  %s: kein einziger Index gelesen" % subtree)
            stats["failed"] += 1

    log("%d Indizes, %d Dateien aktuell, %d zu laden, %d Indizes nicht lesbar"
        % (stats["indexes"], stats["current"], len(want), stats["failed"]))
    if stats["failed"]:
        log("unvollstaendig gelesen - die Region wird nicht als fertig vermerkt")
    if not want:
        log("nichts zu tun")
        if not stats["failed"] and not args.list_only:
            record_region(args.target, left, bottom, right, top, subtrees)
        return 0

    list_path = os.path.join(args.target, "fgsync-aria2.txt")
    with open(list_path, "w") as fh:
        for urls, local_dir, name in want:
            # aria2 takes several URIs for one file on a single tab separated
            # line and falls through to the next when one of them fails.
            fh.write("%s\n  dir=%s\n  out=%s\n" % ("\t".join(urls), local_dir, name))
    log("Liste: %s" % list_path)

    if args.list_only:
        return 0
    if not shutil.which("aria2c"):
        log("aria2c nicht gefunden - Liste steht bereit, mit -i abarbeiten")
        return 3

    cmd = ["aria2c", "-i", list_path,
           "--continue=true",                  # resume partial files
           "--auto-file-renaming=false",       # overwrite, never write .1 copies
           "--allow-overwrite=true",
           "--max-connection-per-server=%d" % args.connections,
           "--max-concurrent-downloads=%d" % args.concurrent,
           "--split=%d" % args.connections,
           "--retry-wait=5", "--max-tries=5", "--timeout=30",
           "--connect-timeout=15",
           # spread the connections for one file over the mirrors we found
           "--min-split-size=1M",
           "--summary-interval=10", "--console-log-level=warn"]
    log("aria2c startet ...")
    # The application stops a start by terminating this process; aria2c is a
    # child of it and would otherwise keep downloading unattached.
    child = subprocess.Popen(cmd)

    def stop(signum, frame):
        try:
            child.terminate()
            child.wait(10)
        except Exception:
            try:
                child.kill()
            except Exception:
                pass
        log("abgebrochen")
        sys.exit(130)

    for sig in (signal.SIGTERM, signal.SIGINT):
        try:
            signal.signal(sig, stop)
        except (ValueError, OSError):
            pass
    rc = child.wait()
    if rc == 0 and not stats["failed"]:
        record_region(args.target, left, bottom, right, top, subtrees)
        log("Region vermerkt")
    return rc


if __name__ == "__main__":
    sys.exit(main())
