#!/usr/bin/env python3
"""Convert the coverage outline of a DJI FlightHub 2 / Terra map export to a shapefile zip.

Usage: flighthub_to_shape.py <export.zip | export-folder> [output.zip]

Reads report/map_report.json (outline in pixels + pixel grid) and AT/sfm_geo_desc.json
(local ENU origin), writes a WGS84 polygon shapefile zip and prints the zskarte layer
settings (name, zip, format, OpenLayers flat style) for copy-paste. Standard library only.
"""
import json
import math
import struct
import sys
import zipfile
from pathlib import Path

STYLE = {
    "fill-color": "rgba(52, 152, 219, 0.25)",
    "stroke-color": "#1f618d",
    "stroke-width": 2,
}
PRJ_WGS84 = ('GEOGCS["GCS_WGS_1984",DATUM["D_WGS_1984",SPHEROID["WGS_1984",6378137.0,298.257223563]],'
             'PRIMEM["Greenwich",0.0],UNIT["Degree",0.0174532925199433]]')


def read_json(source, rel_path):
    if source.is_dir():
        return json.loads((source / rel_path).read_text())
    with zipfile.ZipFile(source) as z:
        # Tolerate a top-level folder inside the zip.
        name = next(n for n in z.namelist() if n == rel_path or n.endswith("/" + rel_path))
        return json.loads(z.read(name))


def outline_lonlat(source):
    report = read_json(source, "report/map_report.json")
    geo = read_json(source, "AT/sfm_geo_desc.json")

    poly = next(p for p in report["valid_range_polygons"] if p["type"] == "RGB")
    tfw = next(d for d in report["datasets_tfw"] if d["type"] == "RGB")["tfw"]
    ring = [(tfw["corner_tl_x"] + pt["x"] * tfw["scale_x"],
             tfw["corner_tl_y"] + pt["y"] * tfw["scale_y"]) for pt in poly["valid_range_polygon"][0]]

    # Local east/north metres -> WGS84 around the reference point (exact enough for a few km).
    lat0, lon0 = geo["ref_GPS"]["latitude"], geo["ref_GPS"]["longitude"]
    a, e2 = 6378137.0, 6.69437999014e-3
    s = math.sin(math.radians(lat0))
    m_radius = a * (1 - e2) / (1 - e2 * s * s) ** 1.5
    n_radius = a / math.sqrt(1 - e2 * s * s)
    lonlat = [(lon0 + math.degrees(e / (n_radius * math.cos(math.radians(lat0)))),
               lat0 + math.degrees(n / m_radius)) for e, n in ring]

    area = abs(sum(x0 * y1 - x1 * y0 for (x0, y0), (x1, y1) in zip(ring, ring[1:] + ring[:1]))) / 2
    return lonlat, area


def write_shapefile(points, area, name):
    # Shapefile outer rings are clockwise and closed.
    if sum(x0 * y1 - x1 * y0 for (x0, y0), (x1, y1) in zip(points, points[1:] + points[:1])) > 0:
        points = points[::-1]
    if points[0] != points[-1]:
        points = points + [points[0]]
    xs, ys = [p[0] for p in points], [p[1] for p in points]
    bbox = (min(xs), min(ys), max(xs), max(ys))

    content = (struct.pack("<i4dii", 5, *bbox, 1, len(points)) + struct.pack("<i", 0)
               + b"".join(struct.pack("<2d", *p) for p in points))
    record = struct.pack(">ii", 1, len(content) // 2) + content

    def header(file_len):
        return (struct.pack(">i5ii", 9994, 0, 0, 0, 0, 0, file_len // 2)
                + struct.pack("<ii4d4d", 1000, 5, *bbox, 0, 0, 0, 0))

    shp = header(100 + len(record)) + record
    shx = header(108) + struct.pack(">ii", 50, len(content) // 2)

    fields = [(b"id", b"N", 10, 0), (b"name", b"C", 80, 0), (b"area_m2", b"N", 12, 1)]
    rec_len = 1 + sum(f[2] for f in fields)
    dbf = struct.pack("<B3BIHH20x", 3, 126, 1, 1, 1, 32 + 32 * len(fields) + 1, rec_len)
    for fname, ftype, flen, fdec in fields:
        dbf += struct.pack("<11sc4xBB14x", fname, ftype, flen, fdec)
    dbf += b"\r " + b"1".rjust(10) + name.encode()[:80].ljust(80) + f"{area:.1f}".encode().rjust(12) + b"\x1a"

    return {"shp": shp, "shx": shx, "dbf": dbf, "prj": PRJ_WGS84.encode(), "cpg": b"UTF-8"}


def main():
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    source = Path(sys.argv[1])
    name = source.stem if source.suffix.lower() == ".zip" else source.name
    out = Path(sys.argv[2]) if len(sys.argv) > 2 else source.parent / f"{name}_shape.zip"

    points, area = outline_lonlat(source)
    files = write_shapefile(points, area, name)
    with zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as z:
        for ext, data in files.items():
            z.writestr(f"{name}.{ext}", data)

    print(f"Wrote {out} ({len(points)} points, {area:.0f} m2)\n")
    print(f"Name of Export:\n{name}\n")
    print(f"Name of Zip:\n{out.name}\n")
    print("Format:\nOpenLayers Flat\n")
    print(f"Style Config:\n{json.dumps(STYLE, indent=2)}")


if __name__ == "__main__":
    main()