#!/usr/bin/env python3
# Sentinel-2 cloudless tiles zoom 0-10 -> MBTiles  (resume-capable)
# Quelle: EOX s2cloudless-2021 CC BY 4.0  https://s2maps.eu
import io, os, sys, sqlite3, time, threading, urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from PIL import Image

MAX_ZOOM = int(os.environ.get("SAT_MAX_ZOOM", "10"))
THREADS  = int(os.environ.get("SAT_THREADS",  "8"))
JPEG_QUALITY = int(os.environ.get("SAT_JPEG_QUALITY", "90"))
JPEG_OPTIMIZE = os.environ.get("SAT_JPEG_OPTIMIZE", "0") == "1"
DB_BATCH = int(os.environ.get("SAT_DB_BATCH", "500"))
CHUNK = int(os.environ.get("SAT_CHUNK", "500"))
# Schreibt in .part-Datei → Martin ignoriert diese während des Downloads
_FINAL   = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/data/satellite.mbtiles")
OUTPUT   = _FINAL.with_suffix(".mbtiles.part")
TILE_URL = ("https://tiles.maps.eox.at/wmts/1.0.0/"
            "s2cloudless-2021_3857/default/GoogleMapsCompatible/{z}/{y}/{x}.jpg")
HEADERS  = {"User-Agent": "rd13-tileserver/1.0 (self-hosted)", "Referer": "https://s2maps.eu"}
_lock    = threading.Lock()
stats    = {"ok": 0, "skip": 0, "err": 0, "conv": 0}
pending  = []
PNG_SIG  = b"\x89PNG\r\n\x1a\n"


def init_db(path):
    path.parent.mkdir(parents=True, exist_ok=True)
    c = sqlite3.connect(str(path), check_same_thread=False)
    c.execute("PRAGMA journal_mode=WAL")
    c.execute("PRAGMA synchronous=NORMAL")
    c.execute("CREATE TABLE IF NOT EXISTS metadata(name TEXT PRIMARY KEY, value TEXT)")
    c.execute("CREATE TABLE IF NOT EXISTS tiles("
              "zoom_level INTEGER, tile_column INTEGER, tile_row INTEGER, tile_data BLOB,"
              "PRIMARY KEY(zoom_level, tile_column, tile_row))")
    c.executemany("INSERT OR REPLACE INTO metadata VALUES(?,?)", [
        ("name",        "satellite"),
        ("format",      "jpg"),
        ("type",        "baselayer"),
        ("description", "Sentinel-2 cloudless 2021 (EOX CC BY 4.0)"),
        ("version",     "1"),
        ("minzoom",     "0"),
        ("maxzoom",     str(MAX_ZOOM)),
        ("bounds",      "-180,-85.051129,180,85.051129"),
        ("center",      "0,0,2"),
    ])
    c.commit()
    return c


def fetch(url):
    req = urllib.request.Request(url, headers=HEADERS)
    for i in range(3):
        try:
            with urllib.request.urlopen(req, timeout=30) as r:
                return r.read() if r.status == 200 else None
        except Exception:
            if i < 2:
                time.sleep(2 ** i)
    return None


def normalize_tile(data):
    if data.startswith(b"\xff\xd8"):
        return data, False
    if not data.startswith(PNG_SIG):
        raise ValueError("unknown tile image format")

    with Image.open(io.BytesIO(data)) as img:
        if img.mode == "RGBA" or "transparency" in img.info:
            rgba = img.convert("RGBA")
            background = Image.new("RGBA", rgba.size, (0, 0, 0, 255))
            background.alpha_composite(rgba)
            rgb = background.convert("RGB")
        else:
            rgb = img.convert("RGB")
        out = io.BytesIO()
        rgb.save(out, format="JPEG", quality=JPEG_QUALITY, optimize=JPEG_OPTIMIZE)
        return out.getvalue(), True


def do_tile(conn, z, x, y):
    y_tms = (1 << z) - 1 - y
    with _lock:
        if conn.execute(
            "SELECT 1 FROM tiles WHERE zoom_level=? AND tile_column=? AND tile_row=?",
            (z, x, y_tms)
        ).fetchone():
            stats["skip"] += 1
            return
    data = fetch(TILE_URL.format(z=z, y=y, x=x))
    converted = False
    if data:
        try:
            data, converted = normalize_tile(data)
        except Exception as e:
            log("[Sat] ERROR normalize z=%d x=%d y=%d: %s" % (z, x, y, e))
            data = None
    with _lock:
        if data:
            pending.append((z, x, y_tms, data))
            if len(pending) >= DB_BATCH:
                conn.executemany("INSERT OR REPLACE INTO tiles VALUES(?,?,?,?)", pending)
                conn.commit()
                pending.clear()
            stats["ok"] += 1
            if converted:
                stats["conv"] += 1
        else:
            stats["err"] += 1


def log(msg):
    print(msg, flush=True)


def tile_coords():
    """Generator: liefert alle (z,x,y)-Koordinaten ohne sie alle im RAM zu halten."""
    for z in range(MAX_ZOOM + 1):
        for x in range(1 << z):
            for y in range(1 << z):
                yield z, x, y


def main():
    total = sum(4 ** z for z in range(MAX_ZOOM + 1))
    log("[Sat] Output:  %s" % OUTPUT)
    log("[Sat] Zoom:    0-%d  |  Threads: %d  |  Tiles: %d" % (MAX_ZOOM, THREADS, total))
    log("[Sat] Quelle:  Sentinel-2 cloudless 2021 (EOX, CC BY 4.0)")
    log("[Sat] Resume:  bereits heruntergeladene Tiles werden übersprungen")
    conn = init_db(OUTPUT)
    t0 = time.time()
    n = 0

    coords = tile_coords()

    try:
        with ThreadPoolExecutor(max_workers=THREADS) as ex:
            while True:
                chunk = []
                for _ in range(CHUNK):
                    coord = next(coords, None)
                    if coord is None:
                        break
                    chunk.append(coord)
                if not chunk:
                    break
                futs = {ex.submit(do_tile, conn, z, x, y): 1 for z, x, y in chunk}
                for f in as_completed(futs):
                    try:
                        f.result()
                    except Exception as e:
                        log("[Sat] ERROR in thread: %s" % e)
                    n += 1
                    if n % 1000 == 0 or n == total:
                        dt  = time.time() - t0 or 0.001
                        eta = (total - n) / n * dt / 3600
                        log("[Sat] %5.1f%%  %d/%d  ok=%d skip=%d err=%d conv=%d  %.0ft/s  ETA=%.1fh" % (
                            n / total * 100, n, total,
                            stats["ok"], stats["skip"], stats["err"], stats["conv"],
                            n / dt, eta))
        with _lock:
            if pending:
                conn.executemany("INSERT OR REPLACE INTO tiles VALUES(?,?,?,?)", pending)
                conn.commit()
        log("[Sat] Erstelle Index...")
        conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx ON tiles(zoom_level, tile_column, tile_row)")
        conn.execute("ANALYZE")
        log("[Sat] Benenne um: %s → %s" % (OUTPUT, _FINAL))
        OUTPUT.replace(_FINAL)
        log("[Sat] Fertig  %.1fh  --  %s" % ((time.time() - t0) / 3600, _FINAL))
    finally:
        conn.close()


if __name__ == "__main__":
    main()