from pathlib import Path
import json
import math
import os
import struct
import zlib

import mercantile
import numpy as np
from shapely import contains_xy
from shapely.geometry import box, shape
from shapely.ops import unary_union
from shapely.strtree import STRtree

from .foundation import (
    DEFAULT_ZOOMS,
    LAND_SEA_TILE_ROOT,
    METADATA_ROOT,
    POLICY_ROOT,
    REGION,
    SOURCE_ASSET,
    TILE_SIZE,
    build_display_policies,
    build_mask_metadata,
)
from .coast_transition import DEFAULT_TRANSITION_RADIUS_BY_ZOOM, build_coast_transition_band


def get_zoom_levels():
    raw_value = os.environ.get("GEO_MASK_ZOOMS", "").strip()
    if not raw_value:
        return DEFAULT_ZOOMS

    zooms = []
    for chunk in raw_value.split(","):
        chunk = chunk.strip()
        if not chunk:
            continue
        zooms.append(int(chunk))
    return sorted(set(zooms))


def load_land_geometries(path=SOURCE_ASSET):
    with path.open(encoding="utf-8") as file_handle:
        payload = json.load(file_handle)

    geometries = [shape(feature["geometry"]) for feature in payload["features"]]
    return geometries


def build_spatial_index(geometries):
    tree = STRtree(geometries)
    return tree, geometries


def png_chunk(chunk_type, data):
    return (
        struct.pack(">I", len(data))
        + chunk_type
        + data
        + struct.pack(">I", zlib.crc32(chunk_type + data) & 0xFFFFFFFF)
    )


def encode_png(rgba):
    height, width, _ = rgba.shape
    raw = b"".join(b"\x00" + rgba[row].tobytes() for row in range(height))
    ihdr = struct.pack(">IIBBBBB", width, height, 8, 6, 0, 0, 0)
    return b"".join(
        [
            b"\x89PNG\r\n\x1a\n",
            png_chunk(b"IHDR", ihdr),
            png_chunk(b"IDAT", zlib.compress(raw, level=6)),
            png_chunk(b"IEND", b""),
        ]
    )


def tile_pixel_lon_lat(tile, tile_size):
    zoom_scale = 2 ** tile.z
    x_pixels = tile.x * tile_size + np.arange(tile_size, dtype=np.float64) + 0.5
    y_pixels = tile.y * tile_size + np.arange(tile_size, dtype=np.float64) + 0.5

    lon = x_pixels / (tile_size * zoom_scale) * 360.0 - 180.0
    mercator_y = math.pi * (1.0 - 2.0 * y_pixels / (tile_size * zoom_scale))
    lat = np.degrees(np.arctan(np.sinh(mercator_y)))
    return np.meshgrid(lon.astype(np.float32), lat.astype(np.float32))


def tile_output_path(tile):
    return LAND_SEA_TILE_ROOT / str(tile.z) / str(tile.x) / f"{tile.y}.png"


def write_json(path, payload):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as file_handle:
        json.dump(payload, file_handle, ensure_ascii=False, indent=2)


def candidate_union(tree, geometries, tile):
    bounds = mercantile.bounds(tile)
    tile_box = box(bounds.west, bounds.south, bounds.east, bounds.north)
    candidates = tree.query(tile_box)
    if len(candidates) == 0:
        return None

    candidate_geometries = [geometries[int(index)] for index in candidates]
    tile_geometries = [geometry for geometry in candidate_geometries if geometry.intersects(tile_box)]
    if not tile_geometries:
        return None
    return unary_union(tile_geometries)


def render_tile(tree, geometries, tile):
    geom = candidate_union(tree, geometries, tile)
    if geom is None:
        rgba = np.zeros((TILE_SIZE, TILE_SIZE, 4), dtype=np.uint8)
        return rgba

    lon_grid, lat_grid = tile_pixel_lon_lat(tile, TILE_SIZE)
    land_mask = contains_xy(geom, lon_grid, lat_grid)
    transition = build_coast_transition_band(
        land_mask,
        DEFAULT_TRANSITION_RADIUS_BY_ZOOM.get(tile.z, 6.0),
    )

    rgba = np.zeros((TILE_SIZE, TILE_SIZE, 4), dtype=np.uint8)
    rgba[:, :, 0][land_mask] = 255
    rgba[:, :, 1] = np.round(transition * 255.0).astype(np.uint8)
    rgba[:, :, 3] = np.where((rgba[:, :, 0] > 0) | (rgba[:, :, 1] > 0), 255, 0).astype(np.uint8)
    return rgba


def generate_tiles():
    if not SOURCE_ASSET.exists():
        raise RuntimeError(f"missing geo mask source asset: {SOURCE_ASSET}")

    geometries = load_land_geometries()
    tree, indexed_geometries = build_spatial_index(geometries)
    tile_count = 0

    for zoom in get_zoom_levels():
        for tile in mercantile.tiles(
            REGION["lon_min"],
            REGION["lat_min"],
            REGION["lon_max"],
            REGION["lat_max"],
            zooms=[zoom],
        ):
            rgba = render_tile(tree, indexed_geometries, tile)
            output_path = tile_output_path(tile)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            output_path.write_bytes(encode_png(rgba))
            tile_count += 1

    write_json(METADATA_ROOT / "land-sea.json", build_mask_metadata())
    write_json(POLICY_ROOT / "display-product-policies.json", build_display_policies())
    print("generated mask tiles:", tile_count)


def main():
    generate_tiles()


if __name__ == "__main__":
    main()
