#!/usr/bin/env python3
"""
make_kmz_from_geotiff.py

Creates a KMZ (KML + embedded JPEG) from a georeferenced GeoTIFF (preferred) or from an image + manual bounds.

Outputs:
  - <OUT_KMZ> (default: mouza_overlay.kmz)
  - intermediate KML (mouza_overlay.kml) and JPG (annotated_map.jpg)

Usage:
  - Place this script in the folder with your GeoTIFF (mouza_georef.tif)
  - Edit INPUT_IMAGE / GEO_TIFF variables below if different
  - Run: python make_kmz_from_geotiff.py

If GEO_TIFF is provided, bounds are taken automatically. If not, the script will prompt for north/south/east/west.
"""

import os
import zipfile
from PIL import Image
import simplekml
import rasterio
from rasterio.warp import transform_bounds
from pyproj import CRS

# ------------ USER CONFIG ------------
# Input image (will be embedded in KMZ). If your GeoTIFF is also the visual image, set IMAGE = GEO_TIFF.
IMAGE = "img52.jpg"               # path to your raster image (png/jpg/tif)
GEO_TIFF = "mouza_georef.tif"     # georeferenced GeoTIFF (optional but recommended)
OUT_KML = "mouza_overlay.kml"
OUT_KMZ = "mouza_overlay.kmz"
OUT_JPG = "annotated_map.jpg"     # converted image (JPG) to embed in KMZ
# -------------------------------------

def convert_to_jpg(src_path, dst_path):
    im = Image.open(src_path)
    if im.mode in ("RGBA", "LA"):
        bg = Image.new("RGB", im.size, (255,255,255))
        bg.paste(im, mask=im.split()[-1])
        im = bg
    else:
        im = im.convert("RGB")
    im.save(dst_path, "JPEG", quality=90)
    print(f"Saved JPG: {dst_path}")

def get_bounds_from_geotiff(tif_path):
    with rasterio.open(tif_path) as ds:
        src_crs = ds.crs
        bounds = ds.bounds  # left, bottom, right, top in ds.crs
        # we want bounds in EPSG:4326
        if src_crs is None:
            raise RuntimeError("GeoTIFF has no CRS.")
        if CRS.from_user_input(src_crs).to_epsg() == 4326:
            # rasterio bounds: left, bottom, right, top
            west, south, east, north = bounds.left, bounds.bottom, bounds.right, bounds.top
        else:
            # transform bounds to EPSG:4326
            left, bottom, right, top = bounds.left, bounds.bottom, bounds.right, bounds.top
            north, south, east, west = None, None, None, None
            # use rasterio transform_bounds
            b = transform_bounds(src_crs, "EPSG:4326", left, bottom, right, top, densify_pts=21)
            west, south, east, north = b[0], b[1], b[2], b[3]
        print(f"GeoTIFF bounds (lon/lat): north={north}, south={south}, east={east}, west={west}")
        return float(north), float(south), float(east), float(west)

def create_kml(jpg_basename, kml_path, north, south, east, west):
    kml = simplekml.Kml()
    go = kml.newgroundoverlay(name="Mouza overlay")
    go.icon.href = os.path.basename(jpg_basename)
    go.latlonbox.north = float(north)
    go.latlonbox.south = float(south)
    go.latlonbox.east = float(east)
    go.latlonbox.west = float(west)
    # no rotation in this version. If rotation needed then set go.latlonbox.rotation = <deg>
    kml.save(kml_path)
    print(f"Saved KML: {kml_path}")

def package_kmz(kml_path, jpg_path, kmz_path):
    with zipfile.ZipFile(kmz_path, "w", zipfile.ZIP_DEFLATED) as z:
        z.write(kml_path, arcname=os.path.basename(kml_path))
        z.write(jpg_path, arcname=os.path.basename(jpg_path))
    print(f"Saved KMZ: {kmz_path}")

def main():
    if not os.path.exists(IMAGE):
        raise SystemExit(f"Image not found: {IMAGE}")

    # 1) convert image -> JPG (Google Earth likes simple JPG)
    convert_to_jpg(IMAGE, OUT_JPG)

    # 2) determine bounds
    if GEO_TIFF and os.path.exists(GEO_TIFF):
        north, south, east, west = get_bounds_from_geotiff(GEO_TIFF)
    else:
        print("GeoTIFF not found or not set. Please provide bounds in decimal degrees.")
        vals = input("Enter north south east west (space-separated) > ").strip().split()
        if len(vals) != 4:
            raise SystemExit("Invalid bounds provided.")
        north, south, east, west = map(float, vals)

    # 3) fix KML to reference JPG and save
    create_kml(OUT_JPG, OUT_KML, north, south, east, west)

    # 4) package into KMZ
    package_kmz(OUT_KML, OUT_JPG, OUT_KMZ)

    print("\nDone. You can open the resulting KMZ in Google Earth Desktop.")
    print("Files created:", OUT_KML, OUT_JPG, OUT_KMZ)

if __name__ == "__main__":
    main()