#! /usr/bin/python3

from argparse import ArgumentParser
import lxml.etree as ET
from osgeo import gdal, osr
import sys
from zipfile import ZipFile, ZIP_DEFLATED, ZIP_STORED

gdal.UseExceptions()

kml = "http://www.opengis.net/kml/2.2"

def kmlobj(tag, *content, **prop):
    e = ET.Element(ET.QName(kml, tag))
    for (k, v) in prop.items():
        ET.SubElement(e, ET.QName(kml, k)).text = str(v)
    e.extend(content)
    return e

def overlay(href, w, s, e, n):
    o = kmlobj("GroundOverlay",
               kmlobj("LatLonBox", west=w, east=e, north=n, south=s),
               kmlobj("Icon", href=href))
    return o

def pixtoll(ds, x, y):
    # Corner transformation from gdalinfo
    geoxfrm = ds.GetGeoTransform()
    geox = geoxfrm[0] + geoxfrm[1] * x + geoxfrm[2] * y
    geoy = geoxfrm[3] + geoxfrm[4] * x + geoxfrm[5] * y
    proj = osr.SpatialReference(ds.GetProjectionRef())
    wgs84 = osr.SpatialReference(osr.SRS_WKT_WGS84)
    xfrm = osr.CoordinateTransformation(proj, wgs84)
    return xfrm.TransformPoint(geox, geoy)

def kmztile(args, outfile):
    with ZipFile(outfile, "w") as kmz:
        # Force the namespace map because my Garmin eTrex 20 isn't
        # namespace-aware.
        kmldoc = ET.Element(ET.QName(kml, "kml"), nsmap={None: kml})
        doc = kmlobj("Document", name="foo", description="foo")
        kmldoc.append(doc)

        for i, arg in enumerate(args):
            ds = gdal.Open(arg)
            files = ds.GetFileList()
            arcname = "files/%d.jpg" % (i,)
            kmz.write(files[0], arcname)
            w, n, _ = pixtoll(ds, 0, 0)
            e, s, _ = pixtoll(ds, ds.RasterXSize, ds.RasterYSize)
            doc.append(overlay(arcname, w, s, e, n))

        kmz.writestr("%s.kml" % ("doc",), ET.tostring(kmldoc),
                     ZIP_DEFLATED)

parser = ArgumentParser(description="Convert GDAL tiles to KMZ")
parser.add_argument("infile", nargs="+")
parser.add_argument("outfile")
args = parser.parse_args()

kmztile(args.infile, args.outfile)