#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

from PIL import Image, ImageDraw
import mapbox_vector_tile
from shapely.geometry import box, shape
from shapely.ops import transform, unary_union
from shapely.prepared import prep


RADIUS = 6378137.0
MAX_MERCATOR_LAT = 85.0511287798066
DEFAULT_BBOX = (129.73110007140482, 32.9464631795233, 131.19176072338317, 34.25430274729531)
COAST_CELL_M = 200.0
FISH_CELL_M = 20.0
HAZARD_CELL_M = 50.0

COAST_SRC = [
    "coastline/C23-06_40_GML.zip",
    "coastline/C23-06_41_GML.zip",
]
FISH_AOI_SRC = "src/pbf/port-overlay/kyushu/port_overlay_kyushu_areas.geojson"
PBF_TILE_ROOT = Path("/home/wwwroot/pbf-delivery-kyushu-reencoded")
PBF_TILE_Z = 12
HAZARD_LAYERS = {
    "navigation_hazard_area",
    "fixed_fishing_gear_area",
    "anchor_caution_hazard_area",
    "navigation_hazard_point",
    "anchor_caution_hazard_point",
    "navigation_marks",
}
LAND_LAYER = "land_area"
POINT_LAYERS = {
    "navigation_hazard_point",
    "anchor_caution_hazard_point",
    "navigation_marks",
}


def mercator_x(lon: float) -> float:
    return RADIUS * math.radians(lon)


def mercator_y(lat: float) -> float:
    lat = max(min(lat, MAX_MERCATOR_LAT), -MAX_MERCATOR_LAT)
    return RADIUS * math.log(math.tan(math.pi / 4.0 + math.radians(lat) / 2.0))


def lon_from_mercator(x: float) -> float:
    return math.degrees(x / RADIUS)


def lat_from_mercator(y: float) -> float:
    return math.degrees(2.0 * math.atan(math.exp(y / RADIUS)) - math.pi / 2.0)


def bbox_intersects(a, b) -> bool:
    return not (a[2] < b[0] or a[0] > b[2] or a[3] < b[1] or a[1] > b[3])


def local_name(tag: str) -> str:
    return tag.split("}", 1)[-1]


def load_xml(zip_path: Path):
    import zipfile
    import xml.etree.ElementTree as ET

    with zipfile.ZipFile(zip_path) as zf:
        xml_name = next(name for name in zf.namelist() if name.endswith(".xml") and "META" not in name)
        return ET.fromstring(zf.read(xml_name))


def lonlat_from_text(text: str) -> tuple[float, float]:
    parts = [p for p in text.replace(",", " ").split() if p]
    if len(parts) < 2:
        raise ValueError(f"invalid coordinate text: {text!r}")
    lat = float(parts[0])
    lon = float(parts[1])
    return lon, lat


def build_point_lookup(root) -> dict[str, tuple[float, float]]:
    lookup: dict[str, tuple[float, float]] = {}
    for point in root.iter():
        if local_name(point.tag) != "GM_Point":
            continue
        point_id = point.attrib.get("id")
        if not point_id:
            continue
        coord_text = None
        for el in point.iter():
            if local_name(el.tag).endswith("DirectPosition.coordinate") and (el.text or "").strip():
                coord_text = el.text.strip()
                break
        if coord_text is None:
            continue
        lookup[point_id] = lonlat_from_text(coord_text)
    return lookup


def extract_curve_coords(curve, point_lookup: dict[str, tuple[float, float]]) -> list[tuple[float, float]]:
    coords: list[tuple[float, float]] = []
    for pos in curve.iter():
        tag = local_name(pos.tag)
        if not tag.startswith("GM_Position"):
            continue
        direct = None
        ref = None
        for child in list(pos):
            child_tag = local_name(child.tag)
            if child_tag.endswith("DirectPosition.coordinate") and (child.text or "").strip():
                direct = child.text.strip()
                break
            if child_tag.endswith("PointRef.point"):
                ref = child.attrib.get("idref")
        if direct is not None:
            coords.append(lonlat_from_text(direct))
        elif ref and ref in point_lookup:
            coords.append(point_lookup[ref])
    return coords


def tile_lonlat_bounds(z: int, x: int, y: int) -> tuple[float, float, float, float]:
    n = 2**z
    lon_left = x / n * 360.0 - 180.0
    lon_right = (x + 1) / n * 360.0 - 180.0
    lat_top = math.degrees(math.atan(math.sinh(math.pi * (1 - 2 * y / n))))
    lat_bottom = math.degrees(math.atan(math.sinh(math.pi * (1 - 2 * (y + 1) / n))))
    return lon_left, lat_bottom, lon_right, lat_top


def tile_mercator_bounds(z: int, x: int, y: int) -> tuple[float, float, float, float]:
    lon_min, lat_min, lon_max, lat_max = tile_lonlat_bounds(z, x, y)
    return mercator_x(lon_min), mercator_y(lat_min), mercator_x(lon_max), mercator_y(lat_max)


def transform_geom_from_tile(geom, bounds_m: tuple[float, float, float, float], extent: int):
    minx, miny, maxx, maxy = bounds_m
    dx = maxx - minx
    dy = maxy - miny

    def coord(x: float, y: float):
        mx = minx + (x / extent) * dx
        my = miny + (y / extent) * dy
        return mx, my

    def walk(obj):
        if isinstance(obj[0], (int, float)):
            return coord(obj[0], obj[1])
        return [walk(item) for item in obj]

    return walk(geom)


def iter_polygons(geom):
    gt = geom.geom_type
    if gt == "Polygon":
        yield geom
    elif gt == "MultiPolygon":
        for g in geom.geoms:
            yield g


def px_from_mercator(x: float, minx: float, px_m: float) -> float:
    return (x - minx) / px_m


def py_from_mercator(y: float, maxy: float, px_m: float) -> float:
    return (maxy - y) / px_m


def cell_to_px_rect(minx: float, maxy: float, px_m: float, x: float, y: float, cell_m: float) -> tuple[float, float, float, float]:
    x0 = px_from_mercator(x, minx, px_m)
    y0 = py_from_mercator(y + cell_m, maxy, px_m)
    x1 = px_from_mercator(x + cell_m, minx, px_m)
    y1 = py_from_mercator(y, maxy, px_m)
    return x0, y0, x1, y1


def main() -> None:
    parser = argparse.ArgumentParser(description="Render Fukuoka/Saga three-layer navigation preview")
    parser.add_argument("--output-png", default="src/pbf/coastline-only/fukuoka_saga_three_layer/fukuoka_saga_three_layer.png")
    parser.add_argument("--bbox", nargs=4, type=float, default=list(DEFAULT_BBOX))
    parser.add_argument("--pixel-size-m", type=float, default=20.0)
    parser.add_argument("--coast-buffer-m", type=float, default=12.0)
    parser.add_argument("--point-buffer-m", type=float, default=25.0)
    args = parser.parse_args()

    project_root = Path(__file__).resolve().parent.parent
    out_path = Path(args.output_png)
    if not out_path.is_absolute():
        out_path = project_root / out_path
    out_path.parent.mkdir(parents=True, exist_ok=True)

    bbox = tuple(args.bbox)
    px_m = args.pixel_size_m

    # 1) 渔港 AOI
    aoi_path = project_root / FISH_AOI_SRC
    data = json.loads(aoi_path.read_text(encoding="utf-8"))
    fish_polygons = []
    fish_bounds_list = []
    for feature in data.get("features", []):
        props = feature.get("properties") or {}
        if props.get("canonical_object_type") != "漁港":
            continue
        geom = shape(feature["geometry"])
        if not bbox_intersects(geom.bounds, bbox):
            continue
        mg = transform(lambda x, y, z=None: (mercator_x(x), mercator_y(y)), geom)
        fish_polygons.append(mg)
        fish_bounds_list.append(mg.bounds)

    if not fish_polygons:
        raise SystemExit("no fish-port AOI polygons selected")

    fish_geom = unary_union(fish_polygons)
    fish_prepared = prep(fish_geom)

    # 1b) 海岸线底座，用于限制渔港 20m 小框只出现在重叠区域
    coast_lines = []
    for src in COAST_SRC:
        root = load_xml(project_root / src)
        point_lookup = build_point_lookup(root)
        for curve in root.findall(".//{http://www.opengis.net/gml/3.2}Curve"):
            coords: list[tuple[float, float]] = []
            for pos_list in curve.findall(".//{http://www.opengis.net/gml/3.2}posList"):
                if not pos_list.text:
                    continue
                values = [float(part) for part in pos_list.text.split() if part]
                if len(values) < 4 or len(values) % 2 != 0:
                    continue
                for i in range(0, len(values), 2):
                    lat = values[i]
                    lon = values[i + 1]
                    coords.append((mercator_x(lon), mercator_y(lat)))
            if len(coords) < 2:
                coords = [(mercator_x(lon), mercator_y(lat)) for lon, lat in extract_curve_coords(curve, point_lookup)]
            if len(coords) < 2:
                continue
            coast_lines.append(shape({"type": "LineString", "coordinates": coords}))

    if not coast_lines:
        raise SystemExit("no coastline curves were parsed")

    coast_geom = unary_union(coast_lines)
    coast_buffer = coast_geom.buffer(args.coast_buffer_m, cap_style=2, join_style=2)
    coast_prepared = prep(coast_buffer)

    # 2) 黄色层：PBF 危险对象
    z = PBF_TILE_Z
    n = 2**z
    lon_min, lat_min, lon_max, lat_max = bbox
    x0 = max(0, int((lon_min + 180.0) / 360.0 * n))
    x1 = min(n - 1, int((lon_max + 180.0) / 360.0 * n))
    def ytile(lat: float) -> int:
        return int((1.0 - math.asinh(math.tan(math.radians(lat))) / math.pi) / 2.0 * n)
    y0 = max(0, ytile(lat_max))
    y1 = min(n - 1, ytile(lat_min))

    hazard_cells_50m: set[tuple[int, int]] = set()
    fish_reef_points: list[tuple[float, float]] = []
    hazard_source_counts: dict[str, int] = {k: 0 for k in HAZARD_LAYERS}
    hazard_features = 0
    tile_count = 0

    for tx in range(x0, x1 + 1):
        for ty in range(y0, y1 + 1):
            p = PBF_TILE_ROOT / f"{z}/{tx}/{ty}.pbf"
            if not p.exists():
                continue
            try:
                tile = mapbox_vector_tile.decode(p.read_bytes())
            except Exception:
                continue
            tile_count += 1
            bounds_m = tile_mercator_bounds(z, tx, ty)
            for layer_name in HAZARD_LAYERS:
                layer = tile.get(layer_name)
                if not layer:
                    continue
                extent = int(layer.get("extent") or 1048576)
                for feature in layer.get("features", []):
                    geom = feature.get("geometry") or {}
                    if not geom:
                        continue
                    props = feature.get("properties") or {}
                    is_fish_reef = (
                        props.get("canonical_object_type") == "fish_reef"
                        or props.get("chart_symbol_code") == "fish_reef"
                        or props.get("class_name") == "魚礁"
                    )
                    gtype = geom.get("type")
                    coords = transform_geom_from_tile(geom.get("coordinates"), bounds_m, extent)
                    shp = shape({"type": gtype, "coordinates": coords})
                    if shp.is_empty:
                        continue
                    if is_fish_reef:
                        fish_reef_points.append((shp.centroid.x, shp.centroid.y))
                        hazard_features += 1
                        hazard_source_counts[layer_name] += 1
                        continue
                    if layer_name in POINT_LAYERS:
                        shp = shp.buffer(args.point_buffer_m)
                    hazard_features += 1
                    hazard_source_counts[layer_name] += 1
                    minx, miny, maxx, maxy = shp.bounds
                    start_x = math.floor(minx / HAZARD_CELL_M) * HAZARD_CELL_M
                    start_y = math.floor(miny / HAZARD_CELL_M) * HAZARD_CELL_M
                    end_x = math.ceil(maxx / HAZARD_CELL_M) * HAZARD_CELL_M
                    end_y = math.ceil(maxy / HAZARD_CELL_M) * HAZARD_CELL_M
                    prep_geom = prep(shp)
                    cy = start_y
                    while cy < end_y:
                        cx = start_x
                        while cx < end_x:
                            cell = box(cx, cy, cx + HAZARD_CELL_M, cy + HAZARD_CELL_M)
                            if prep_geom.intersects(cell):
                                hazard_cells_50m.add((int(round(cx / HAZARD_CELL_M)), int(round(cy / HAZARD_CELL_M))))
                            cx += HAZARD_CELL_M
                        cy += HAZARD_CELL_M

    # 2b) 黄色层单独导出：50m 障碍格 + 鱼礁黄点
    hazard_minx = mercator_x(bbox[0])
    hazard_miny = mercator_y(bbox[1])
    hazard_maxx = mercator_x(bbox[2])
    hazard_maxy = mercator_y(bbox[3])
    hazard_width = max(1, math.ceil((hazard_maxx - hazard_minx) / px_m))
    hazard_height = max(1, math.ceil((hazard_maxy - hazard_miny) / px_m))
    hazard_image = Image.new("RGBA", (hazard_width, hazard_height), (0, 0, 0, 0))
    hazard_draw = ImageDraw.Draw(hazard_image, "RGBA")

    fish_reef_count = 0
    for ix, iy in hazard_cells_50m:
        x = ix * HAZARD_CELL_M
        y = iy * HAZARD_CELL_M
        x0, y0, x1, y1 = cell_to_px_rect(hazard_minx, hazard_maxy, px_m, x, y, HAZARD_CELL_M)
        hazard_draw.rectangle((x0, y0, x1, y1), fill=(242, 194, 48, 150), outline=(235, 180, 20, 160))

    for x, y in fish_reef_points:
        px = px_from_mercator(x, hazard_minx, px_m)
        py = py_from_mercator(y, hazard_maxy, px_m)
        r = 4.0
        hazard_draw.ellipse((px - r, py - r, px + r, py + r), fill=(242, 194, 48, 235), outline=(141, 103, 0, 230), width=1)
        fish_reef_count += 1

    hazard_overlay_path = out_path.parent / "fukuoka_saga_hazard_50m_overlay.png"
    hazard_image.save(hazard_overlay_path)
    hazard_manifest = {
        "bbox_lonlat": list(bbox),
        "pixel_size_m": px_m,
        "hazard_cell_m": HAZARD_CELL_M,
        "hazard_cells_50m": len(hazard_cells_50m),
        "fish_reef_points": fish_reef_count,
        "hazard_features": hazard_features,
        "hazard_source_counts": hazard_source_counts,
        "pbf_tiles_scanned": tile_count,
        "image_width": hazard_width,
        "image_height": hazard_height,
        "png_file": hazard_overlay_path.name,
    }
    hazard_manifest_path = out_path.parent / "fukuoka_saga_hazard_50m_overlay_manifest.json"
    hazard_manifest_path.write_text(json.dumps(hazard_manifest, ensure_ascii=False, indent=2), encoding="utf-8")

    # 4) 20m 渔港不可航行格：只对渔港 AOI 内着色
    minx_all, miny_all = fish_geom.bounds[0], fish_geom.bounds[1]
    maxx_all, maxy_all = fish_geom.bounds[2], fish_geom.bounds[3]
    width = max(1, math.ceil((maxx_all - minx_all) / px_m))
    height = max(1, math.ceil((maxy_all - miny_all) / px_m))
    image = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    draw = ImageDraw.Draw(image, "RGBA")

    red_cells = 0
    coast_cells = 0
    yellow_cells = 0

    for poly in iter_polygons(fish_geom):
        pb = poly.bounds
        start_x = math.floor(pb[0] / FISH_CELL_M) * FISH_CELL_M
        start_y = math.floor(pb[1] / FISH_CELL_M) * FISH_CELL_M
        end_x = math.ceil(pb[2] / FISH_CELL_M) * FISH_CELL_M
        end_y = math.ceil(pb[3] / FISH_CELL_M) * FISH_CELL_M
        prep_poly = prep(poly)
        y = start_y
        while y < end_y:
            x = start_x
            while x < end_x:
                cell = box(x, y, x + FISH_CELL_M, y + FISH_CELL_M)
                if prep_poly.intersects(cell) and coast_prepared.intersects(cell):
                    x0, y0, x1, y1 = cell_to_px_rect(minx_all, maxy_all, px_m, x, y, FISH_CELL_M)
                    draw.rectangle((x0, y0, x1, y1), fill=(222, 92, 92, 170))
                    red_cells += 1
                x += FISH_CELL_M
            y += FISH_CELL_M

    # 5) 黄色 50m 格
    for ix, iy in hazard_cells_50m:
        x = ix * HAZARD_CELL_M
        y = iy * HAZARD_CELL_M
        x0, y0, x1, y1 = cell_to_px_rect(minx_all, maxy_all, px_m, x, y, HAZARD_CELL_M)
        draw.rectangle((x0, y0, x1, y1), fill=(242, 194, 48, 155))
        yellow_cells += 1

    image.save(out_path)

    manifest = {
        "bbox_lonlat": list(bbox),
        "pixel_size_m": px_m,
        "coast_cell_m": COAST_CELL_M,
        "fish_cell_m": FISH_CELL_M,
        "hazard_cell_m": HAZARD_CELL_M,
        "fish_polygon_count": len(fish_polygons),
        "fish_green_cells": 0,
        "fish_red_cells": red_cells,
        "hazard_cells_50m": yellow_cells,
        "hazard_features": hazard_features,
        "hazard_source_counts": hazard_source_counts,
        "pbf_tiles_scanned": tile_count,
        "image_width": width,
        "image_height": height,
        "png_file": out_path.name,
    }
    manifest_path = out_path.parent / "fukuoka_saga_three_layer_manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    print("已导出:")
    print(" -", out_path)
    print(" -", manifest_path)
    print(" -", hazard_overlay_path)
    print(" -", hazard_manifest_path)


if __name__ == "__main__":
    main()