"""The ground under the aircraft: map tiles, fetched, cached and dimmed.

A flight path over a black rectangle says how the aircraft moved and nothing
about where it was.  Over a coastline it says which airport it left.  So the
map behind the animation is a real one: standard raster tiles, fetched once,
kept on disk, reprojected onto the picture and dimmed until the aircraft are
the brightest thing on it.

Three things follow from using somebody else's tile server, and all three are
obligations rather than options.  Tiles are **cached** and never fetched
twice.  Every request identifies the program in its User-Agent.  And the
attribution the licence requires is drawn onto the picture, not left to a
readme nobody ships with a GIF.  A drawing is capped at a few dozen tiles: an
aircraft map is a hobby drawing, not a reason to hammer a volunteer-funded
service.

The PNG decoding is here for the same reason the PNG writing is in
:mod:`bandsaunter.images`: a scanner that cannot draw a map because an
imaging library is missing is worse than one that draws it from zlib and
numpy, which is all this needs.
"""

from __future__ import annotations

import json
import math
import os
import struct
import time
import urllib.parse
import urllib.request
import zlib
from pathlib import Path

import numpy as np

from . import __version__

__all__ = ["decode_png", "tile_of", "choose_zoom", "fetch_tile", "mosaic",
           "ground_under", "TILE_URL", "ATTRIBUTION", "MAX_TILES", "MAX_ZOOM",
           "cache_dir", "PNGError", "airports_in", "AIRPORTS_URL",
           "tile_span"]

# The standard OpenStreetMap tiles.  Any {z}/{x}/{y} server can be put here
# instead; nothing below knows anything about this one in particular.
TILE_URL = "https://tile.openstreetmap.org/{z}/{x}/{y}.png"

# Drawn onto every picture that used the tiles.  The licence requires it and
# a GIF travels without its readme.
ATTRIBUTION = "MAP DATA (C) OPENSTREETMAP CONTRIBUTORS"

# The backstop on how many tiles one drawing may fetch.  Reached only by a
# view so wide that no zoom covers it without fetching half a country; the
# usual limit is the size of the picture, which is asked for instead.  A
# screenful at a hundred-mile radius comes to about a hundred and fifty,
# fetched once and kept.
MAX_TILES = 220
MAX_ZOOM = 13

# How much more detail to fetch than the picture holds.  Averaging several
# source pixels into each output one is what makes lettering and coastlines
# come out smooth; taking exactly one leaves them as hard as they were.
OVERSAMPLE = 1.4
MIN_ZOOM = 2
TILE_PIXELS = 256

USER_AGENT = (f"bandsaunter/{__version__} "
              "(+https://github.com/topics/rtl-sdr; aircraft map drawing)")

# Politeness between requests to a volunteer-funded service.  Only paid on a
# tile that was not already on the disk.
FETCH_PAUSE = 0.12


# Where to ask what aerodromes are in a piece of the world.  The same OSM
# data the tiles are drawn from, asked as a question rather than a picture.
AIRPORTS_URL = "https://overpass-api.de/api/interpreter"

# One question covers a whole view and is kept for a month, because runways
# do not move.  The same politeness as the tiles: cache it, say who is
# asking, and do not ask twice for the same thing.
AIRPORT_CACHE_DAYS = 30

# Bumped when the question changes, so that answers to the old one are asked
# again rather than believed.
AIRPORT_CACHE_VERSION = 2

# Enough to name what is under an aircraft and not so many that the map is a
# list of airstrips.
MOST_AIRPORTS = 40


class PNGError(ValueError):
    """A PNG this decoder cannot read."""


# ---------------------------------------------------------------------------
# Reading a PNG
# ---------------------------------------------------------------------------

_CHANNELS = {0: 1, 2: 3, 3: 1, 4: 2, 6: 4}


def decode_png(data: bytes) -> np.ndarray:
    """A PNG's pixels as an ``(h, w, 3)`` array of bytes.

    Eight bits a channel and no interlacing, which is what every tile server
    sends; anything else raises rather than being guessed at.
    """
    if data[:8] != b"\x89PNG\r\n\x1a\x08"[:8] and data[:8] != \
            b"\x89PNG\r\n\x1a\n":
        raise PNGError("not a PNG")
    width = height = depth = colour = interlace = 0
    palette = None
    body = bytearray()
    at = 8
    while at + 8 <= len(data):
        length, tag = struct.unpack(">I4s", data[at:at + 8])
        chunk = data[at + 8:at + 8 + length]
        at += 12 + length
        if tag == b"IHDR":
            (width, height, depth, colour, _compression, _filter,
             interlace) = struct.unpack(">IIBBBBB", chunk)
        elif tag == b"PLTE":
            palette = np.frombuffer(chunk, dtype=np.uint8).reshape(-1, 3)
        elif tag == b"IDAT":
            body += chunk
        elif tag == b"IEND":
            break
    if depth != 8:
        raise PNGError(f"{depth}-bit PNG: only 8 bits a channel is read here")
    if interlace:
        raise PNGError("interlaced PNG")
    if colour not in _CHANNELS:
        raise PNGError(f"colour type {colour}")
    if not width or not height:
        raise PNGError("no image")

    channels = _CHANNELS[colour]
    raw = _unfilter(zlib.decompress(bytes(body)), width, height, channels)
    if colour == 3:
        if palette is None:
            raise PNGError("palette image with no palette")
        return palette[np.clip(raw[:, :, 0], 0, len(palette) - 1)]
    if colour == 0:
        return np.repeat(raw, 3, axis=2)
    if colour == 4:
        return np.repeat(raw[:, :, :1], 3, axis=2)
    return raw[:, :, :3]


def _unfilter(data: bytes, width: int, height: int,
              channels: int) -> np.ndarray:
    """Undo the per-row filters PNG applies before compressing.

    The five of them, from the specification.  None and Up are whole-row
    arithmetic; Sub is a running total along the row, which is a cumulative
    sum once the bytes are grouped by which channel they belong to; Average
    and Paeth each need the byte before them to have been worked out
    already, so those two are the only ones that walk the row.
    """
    stride = width * channels
    if len(data) < height * (stride + 1):
        raise PNGError("truncated image data")
    out = np.zeros((height, stride), dtype=np.uint8)
    previous = np.zeros(stride, dtype=np.uint8)
    at = 0
    for row in range(height):
        kind = data[at]
        line = np.frombuffer(data, dtype=np.uint8, count=stride,
                             offset=at + 1).astype(np.uint16)
        at += stride + 1
        if kind == 0:
            current = line.astype(np.uint8)
        elif kind == 1:
            current = np.empty(stride, dtype=np.uint8)
            for offset in range(channels):
                current[offset::channels] = np.cumsum(
                    line[offset::channels], dtype=np.uint32) % 256
        elif kind == 2:
            current = ((line + previous) % 256).astype(np.uint8)
        elif kind in (3, 4):
            current = _walk_row(bytes(line.astype(np.uint8)), previous,
                                channels, kind)
        else:
            raise PNGError(f"filter type {kind}")
        out[row] = current
        previous = current
    return out.reshape(height, width, channels)


def _walk_row(line: bytes, previous: np.ndarray, channels: int,
              kind: int) -> np.ndarray:
    """Average and Paeth: each byte needs the one ``channels`` back."""
    up = previous.tolist()
    out = [0] * len(line)
    for i, value in enumerate(line):
        left = out[i - channels] if i >= channels else 0
        above = up[i]
        if kind == 3:
            out[i] = (value + ((left + above) >> 1)) & 0xFF
        else:
            upleft = up[i - channels] if i >= channels else 0
            base = left + above - upleft
            da, db, dc = abs(base - left), abs(base - above), abs(base - upleft)
            near = left if (da <= db and da <= dc) else (
                above if db <= dc else upleft)
            out[i] = (value + near) & 0xFF
    return np.array(out, dtype=np.uint8)


# ---------------------------------------------------------------------------
# Which tiles
# ---------------------------------------------------------------------------

def tile_of(lat: float, lon: float, zoom: int) -> tuple[float, float]:
    """Where a coordinate falls in the tile grid, in fractional tiles.

    Web Mercator, which is what every {z}/{x}/{y} tile server serves and is
    not what this program's maps are drawn in -- hence the resampling
    further down rather than a straight paste.
    """
    lat = max(-85.05112878, min(85.05112878, lat))
    n = float(2 ** zoom)
    x = (lon + 180.0) / 360.0 * n
    radians = math.radians(lat)
    y = (1.0 - math.asinh(math.tan(radians)) / math.pi) / 2.0 * n
    return x, y


def tile_span(south: float, west: float, north: float, east: float,
              zoom: int) -> tuple[int, int]:
    """How many tiles wide and tall a box is at one zoom."""
    x0, y0 = tile_of(north, west, zoom)
    x1, y1 = tile_of(south, east, zoom)
    return (int(math.floor(x1)) - int(math.floor(x0)) + 1,
            int(math.floor(y1)) - int(math.floor(y0)) + 1)


def choose_zoom(south: float, west: float, north: float, east: float,
                width: int = 0, max_tiles: int = MAX_TILES,
                most: int = MAX_ZOOM) -> int:
    """The zoom to fetch at: enough for the picture, and no more.

    ``width`` is how many pixels across the picture will be.  Without it
    this fetched whatever the tile budget allowed, which is the wrong
    question in both directions: on a small picture it fetched more than
    could be shown, and on a large one it fetched less and the map was
    blown up to fit -- which is what a low-resolution map looks like.

    So it climbs until the tiles hold at least as many pixels as the
    picture wants, and stops there.  The budget is the backstop, for a view
    so wide that no zoom can cover it without fetching half a country.
    """
    best = MIN_ZOOM
    for zoom in range(MIN_ZOOM, min(most, MAX_ZOOM) + 1):
        wide, tall = tile_span(south, west, north, east, zoom)
        if wide * tall > max_tiles:
            break
        best = zoom
        if width and wide * TILE_PIXELS >= width * OVERSAMPLE:
            break                  # detail enough; more would only cost time
    return best


# ---------------------------------------------------------------------------
# Fetching them, once
# ---------------------------------------------------------------------------

def cache_dir() -> Path:
    root = os.environ.get("XDG_CACHE_HOME") or "~/.cache"
    return Path(root).expanduser() / "bandsaunter" / "tiles"


def fetch_tile(zoom: int, x: int, y: int, url: str = TILE_URL,
               timeout: float = 10.0, cache: Path | None = None) -> bytes | None:
    """One tile, from the disk if it has ever been fetched before.

    Returns the PNG bytes, or None if it could not be had.  A missing tile is
    not an error: the map is drawn with a hole in it, which is better than no
    map and much better than an exception in the middle of an animation.
    """
    where = (cache if cache is not None else cache_dir()) / str(zoom) / str(x)
    path = where / f"{y}.png"
    try:
        return path.read_bytes()
    except OSError:
        pass
    request = urllib.request.Request(url.format(z=zoom, x=x, y=y),
                                     headers={"User-Agent": USER_AGENT})
    try:
        with urllib.request.urlopen(request, timeout=timeout) as answer:
            body = answer.read(2_000_000)
    except Exception:
        return None
    try:
        where.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(".tmp")
        tmp.write_bytes(body)
        tmp.replace(path)
    except OSError:
        pass                      # an unwritable cache is not a reason to stop
    return body


def mosaic(south: float, west: float, north: float, east: float, zoom: int,
           fetch=fetch_tile, pause: float = FETCH_PAUSE, **kw):
    """Every tile the box touches, stitched into one image.

    Returns the pixels and where their top-left corner sits in the world, in
    tile-grid pixels at this zoom, so the resampling below can place them.
    """
    x0, y0 = tile_of(north, west, zoom)
    x1, y1 = tile_of(south, east, zoom)
    left, top = int(math.floor(x0)), int(math.floor(y0))
    right, bottom = int(math.floor(x1)), int(math.floor(y1))
    span = 2 ** zoom
    wide, tall = right - left + 1, bottom - top + 1
    if wide <= 0 or tall <= 0 or wide * tall > MAX_TILES * 4:
        return None, 0, 0
    canvas = np.zeros((tall * TILE_PIXELS, wide * TILE_PIXELS, 3),
                      dtype=np.uint8)
    got = 0
    for row in range(tall):
        for column in range(wide):
            tx, ty = (left + column) % span, top + row
            if not 0 <= ty < span:
                continue
            body = fetch(zoom, tx, ty, **kw)
            if body is None:
                continue
            try:
                tile = decode_png(body)
            except (PNGError, zlib.error, ValueError):
                continue
            if tile.shape[0] != TILE_PIXELS or tile.shape[1] != TILE_PIXELS:
                continue
            canvas[row * TILE_PIXELS:(row + 1) * TILE_PIXELS,
                   column * TILE_PIXELS:(column + 1) * TILE_PIXELS] = tile
            got += 1
            if pause:
                time.sleep(pause)
    if not got:
        return None, 0, 0
    return canvas, left * TILE_PIXELS, top * TILE_PIXELS


# ---------------------------------------------------------------------------
# What is on the ground
# ---------------------------------------------------------------------------

def _resample(values: np.ndarray, edges: np.ndarray, axis: int) -> np.ndarray:
    """Average each output cell over the source pixels that fall in it.

    A box filter, done as the difference of a running total so the whole
    axis is two passes rather than a loop.  Where the source is coarser than
    the output -- a map zoomed in further than the tiles go -- a cell covers
    less than one source pixel, and it takes that one.
    """
    length = values.shape[axis]
    starts = np.clip(np.floor(edges[:-1]).astype(np.int64), 0, length - 1)
    ends = np.clip(np.ceil(edges[1:]).astype(np.int64), 1, length)
    ends = np.maximum(ends, starts + 1)
    running = np.cumsum(values, axis=axis, dtype=np.float64)
    pad = np.zeros_like(np.take(running, [0], axis=axis))
    running = np.concatenate([pad, running], axis=axis)
    total = (np.take(running, ends, axis=axis)
             - np.take(running, starts, axis=axis))
    counts = (ends - starts).astype(np.float64)
    shape = [1] * values.ndim
    shape[axis] = counts.size
    return (total / counts.reshape(shape)).astype(np.float32)


def _airport_cache(box) -> Path:
    root = os.environ.get("XDG_CACHE_HOME") or "~/.cache"
    name = "_".join(f"{round(v, 1):+06.1f}" for v in box)
    return Path(root).expanduser() / "bandsaunter" / "airports" / f"{name}.json"


def airports_in(south: float, west: float, north: float, east: float,
                url: str = AIRPORTS_URL, timeout: float = 45.0,
                cache: Path | None = None, ask=None) -> list[dict]:
    """Every aerodrome in a piece of the world, with a code and a position.

    Asked of the same OpenStreetMap data the tiles are drawn from, as a
    question rather than a picture, and kept on disk afterwards: a runway
    does not move, so one question covers a view for a month.

    An aerodrome with no code is left out.  A map wants to say *which*
    airport an aircraft is over, and there are a great many landing strips
    with a name and nothing else; the ones worth marking have a code.
    """
    box = (round(south, 1), round(west, 1), round(north, 1), round(east, 1))
    path = cache if cache is not None else _airport_cache(box)
    try:
        body = json.loads(path.read_text())
        if int(body.get("version") or 0) >= AIRPORT_CACHE_VERSION and \
                time.time() - float(body.get("fetched_at") or 0) < \
                AIRPORT_CACHE_DAYS * 86_400:
            return body.get("airports") or []
    except (OSError, ValueError):
        pass

    found: list[dict] = []
    try:
        raw = (ask or _ask_overpass)(box, url, timeout)
        for element in (raw or {}).get("elements", []):
            tags = element.get("tags") or {}
            code = _airport_code(tags)
            lat = element.get("lat") or (element.get("center") or {}).get("lat")
            lon = element.get("lon") or (element.get("center") or {}).get("lon")
            if not code or lat is None or lon is None:
                continue
            found.append({"code": code,
                          "name": (tags.get("name") or "").strip(),
                          "icao": bool(tags.get("icao")),
                          "latitude": float(lat), "longitude": float(lon)})
    except Exception:
        return []                  # a map with no airports on it, not a crash

    # The ones with a real ICAO code first, since those are the ones an
    # aircraft is likely to be flying to or from.
    found.sort(key=lambda a: (0 if a.get("icao") else 1, a["code"]))
    # One airport is often tagged twice -- a point for the terminal and an
    # outline for the field -- and marking it twice writes its name over
    # itself.
    once: dict[str, dict] = {}
    for one in found:
        once.setdefault(one["code"], one)
    found = list(once.values())[:MOST_AIRPORTS]
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps({"fetched_at": time.time(),
                                    "version": AIRPORT_CACHE_VERSION,
                                    "airports": found}))
    except OSError:
        pass
    return found


def _airport_code(tags: dict) -> str:
    """What to call an aerodrome, or nothing if it has no code at all.

    The ICAO code where it has one.  Otherwise a reference, and only when it
    is four letters: a great many landing strips carry a local identifier
    like "14AZ" or "MX-0492", which names nothing anybody would recognise
    and turns a map into a list of airstrips.
    """
    icao = (tags.get("icao") or "").strip().upper()
    if len(icao) == 4 and icao.isalpha():
        return icao
    ref = (tags.get("ref") or "").strip().upper()
    if len(ref) == 4 and ref.isalpha():
        return ref
    return ""


def _ask_overpass(box, url: str, timeout: float) -> dict:
    """The one question this program asks Overpass."""
    south, west, north, east = box
    where = f"({south},{west},{north},{east})"
    # Relations as well as nodes and ways: a big airport is a relation more
    # often than not -- Tucson International and Davis-Monthan both are --
    # so asking only for the other two finds every airstrip in the county
    # and misses the two the county is known for.
    query = ("[out:json][timeout:25];"
             f'(node["aeroway"="aerodrome"]{where};'
             f' way["aeroway"="aerodrome"]{where};'
             f' relation["aeroway"="aerodrome"]{where};);'
             "out center tags;")
    request = urllib.request.Request(
        url, data=urllib.parse.urlencode({"data": query}).encode(),
        headers={"User-Agent": USER_AGENT})
    with urllib.request.urlopen(request, timeout=timeout) as answer:
        return json.loads(answer.read(4_000_000).decode("utf8", "replace"))


# ---------------------------------------------------------------------------
# Putting it under the picture
# ---------------------------------------------------------------------------

def ground_under(south: float, west: float, north: float, east: float,
                 width: int, height: int, shades: int = 32,
                 fetch=fetch_tile, zoom: int | None = None,
                 **kw) -> np.ndarray | None:
    """The map for one picture, as ``shades`` levels of brightness.

    The tiles are Web Mercator and the picture is not, so every output pixel
    asks the mosaic where its own latitude and longitude landed rather than
    the mosaic being pasted in.  Over the couple of hundred miles a receiver
    hears, the difference is a few pixels of drift at the top of the frame --
    which is a few pixels an aircraft would be drawn wrong by, and the whole
    point of putting a coastline under it is that the coastline is where the
    aircraft was.

    Returns None when nothing could be fetched, which the caller draws as the
    plain grid it drew before.
    """
    if width < 1 or height < 1 or north <= south or east <= west:
        return None
    if zoom is None:
        zoom = choose_zoom(south, west, north, east, width=width)
    tiles, origin_x, origin_y = mosaic(south, west, north, east, zoom,
                                       fetch=fetch, **kw)
    if tiles is None:
        return None

    # The edges of each output pixel rather than its middle, so that what
    # lands in it can be averaged.  Taking the nearest source pixel instead
    # throws away most of a tile -- the mosaic is commonly half again the
    # size of the picture -- and what survives is the aliasing: hard, broken
    # lettering and roads that come and go along their length.
    lons = west + (east - west) * np.arange(width + 1) / width
    lats = north - (north - south) * np.arange(height + 1) / height
    span = float(2 ** zoom) * TILE_PIXELS
    xs = (lons + 180.0) / 360.0 * span - origin_x
    clipped = np.clip(lats, -85.05112878, 85.05112878)
    ys = (1.0 - np.arcsinh(np.tan(np.radians(clipped))) / math.pi) / 2.0 \
        * span - origin_y

    # Brightness only, inverted, and dimmed.  Inverted because a printed map
    # is ink on white paper and this picture is the other way round: the
    # things drawn on the map -- coastlines, roads, the names of towns --
    # are the dark parts of a tile, and they are what should show against a
    # night background.  Dimmed because the map is the ground under the
    # aircraft rather than the subject: anything drawn on top has to stay the
    # brightest thing on the picture.
    whole = (0.299 * tiles[:, :, 0] + 0.587 * tiles[:, :, 1]
             + 0.114 * tiles[:, :, 2]).astype(np.float32)
    luma = _resample(_resample(whole, ys, axis=0), xs, axis=1)
    low, high = float(luma.min()), float(luma.max())
    if high - low < 1.0:
        levels = np.zeros_like(luma)
    else:
        levels = 1.0 - (luma - low) / (high - low)
    return np.clip((levels * (shades - 1)).round(), 0,
                   shades - 1).astype(np.uint8)