#!/usr/bin/env python3
import io
import os
import signal
import sqlite3
import sys
import time
from pathlib import Path

from PIL import Image

DB = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/data/satellite.mbtiles")
BACKUP = Path(os.environ.get("SAT_BACKUP", str(DB) + ".before-png-to-jpg-fix"))
BATCH = int(os.environ.get("SAT_CONVERT_BATCH", "500"))
JPEG_QUALITY = int(os.environ.get("SAT_JPEG_QUALITY", "90"))
SLEEP_SECONDS = float(os.environ.get("SAT_CONVERT_SLEEP", "0"))
REQUIRE_BACKUP = os.environ.get("SAT_REQUIRE_BACKUP", "1") == "1"

PNG_SIG = b"\x89PNG\r\n\x1a\n"
STOP = False


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


def handle_stop(signum, frame):
    global STOP
    STOP = True
    log("[Convert] Stop requested; finishing current batch before exit")


def png_to_jpeg(data):
    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=False)
        return out.getvalue()


def main():
    if REQUIRE_BACKUP and not BACKUP.exists():
        raise SystemExit("[Convert] Backup missing: %s" % BACKUP)

    signal.signal(signal.SIGTERM, handle_stop)
    signal.signal(signal.SIGINT, handle_stop)

    started = time.time()
    converted = 0
    scanned_png = 0
    batch = []

    conn = sqlite3.connect(str(DB))
    conn.execute("PRAGMA busy_timeout=60000")
    journal_mode = conn.execute("PRAGMA journal_mode=DELETE").fetchone()[0]
    conn.execute("PRAGMA synchronous=NORMAL")

    before = conn.execute("SELECT value FROM metadata WHERE name='format'").fetchone()[0]
    conn.execute("UPDATE metadata SET value='jpg' WHERE name='format'")
    conn.commit()
    log("[Convert] DB=%s" % DB)
    log("[Convert] backup=%s" % BACKUP)
    log("[Convert] batch=%d quality=%d sleep=%.3fs journal=%s" % (
        BATCH, JPEG_QUALITY, SLEEP_SECONDS, journal_mode))
    log("[Convert] metadata.format %s -> jpg" % before)

    try:
        for rowid, data in conn.execute("SELECT rowid, tile_data FROM tiles"):
            if STOP:
                break
            if not data.startswith(PNG_SIG):
                continue

            scanned_png += 1
            batch.append((png_to_jpeg(data), rowid))

            if len(batch) >= BATCH:
                conn.executemany("UPDATE tiles SET tile_data=? WHERE rowid=?", batch)
                conn.commit()
                converted += len(batch)
                batch.clear()
                elapsed = time.time() - started
                rate = converted / elapsed if elapsed else 0
                log("[Convert] converted=%d rate=%.1f/s elapsed=%.1fm" % (
                    converted, rate, elapsed / 60))
                if SLEEP_SECONDS:
                    time.sleep(SLEEP_SECONDS)

        if batch:
            conn.executemany("UPDATE tiles SET tile_data=? WHERE rowid=?", batch)
            conn.commit()
            converted += len(batch)

        conn.execute("ANALYZE")
        conn.commit()
    finally:
        conn.close()

    elapsed = time.time() - started
    status = "stopped" if STOP else "done"
    log("[Convert] %s scanned_png=%d converted=%d elapsed=%.1fm" % (
        status, scanned_png, converted, elapsed / 60))


if __name__ == "__main__":
    main()