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

import argparse
import os
from pathlib import Path

import mapbox_vector_tile


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Normalize vector-tile layer extents by rescaling geometry coordinates."
    )
    parser.add_argument("--src-root", type=Path, required=True)
    parser.add_argument("--dst-root", type=Path, required=True)
    parser.add_argument("--target-extent", type=int, default=4096)
    parser.add_argument(
        "--rewrite-all",
        action="store_true",
        help="rewrite every tile even when all layers already use the target extent",
    )
    return parser.parse_args()


def scale_coordinates(value, factor: float):
    if isinstance(value, (list, tuple)):
        if len(value) >= 2 and isinstance(value[0], (int, float)) and isinstance(value[1], (int, float)):
            return [round(value[0] * factor), round(value[1] * factor)]
        return [scale_coordinates(item, factor) for item in value]
    return value


def hardlink_or_copy(src: Path, dst: Path) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    if dst.exists():
        dst.unlink()
    try:
        os.link(src, dst)
    except OSError:
        dst.write_bytes(src.read_bytes())


def normalize_tile(src: Path, dst: Path, target_extent: int, rewrite_all: bool) -> bool:
    decoded = mapbox_vector_tile.decode(src.read_bytes())
    extents = {
        layer_name: int(layer.get("extent") or target_extent)
        for layer_name, layer in decoded.items()
    }
    if not rewrite_all and all(extent == target_extent for extent in extents.values()):
        hardlink_or_copy(src, dst)
        return False

    layers = []
    per_layer_options: dict[str, dict[str, int]] = {}
    for layer_name, layer in decoded.items():
        source_extent = extents[layer_name]
        factor = target_extent / source_extent if source_extent else 1.0
        features = []
        for feature in layer.get("features", []):
            output_feature = {
                "geometry": {
                    "type": feature["geometry"]["type"],
                    "coordinates": scale_coordinates(feature["geometry"]["coordinates"], factor),
                },
                "properties": dict(feature.get("properties") or {}),
            }
            if feature.get("id") is not None:
                output_feature["id"] = feature["id"]
            features.append(output_feature)
        layers.append({"name": layer_name, "features": features})
        per_layer_options[layer_name] = {"extents": target_extent}

    dst.parent.mkdir(parents=True, exist_ok=True)
    if dst.exists():
        dst.unlink()
    dst.write_bytes(mapbox_vector_tile.encode(layers, per_layer_options=per_layer_options))
    return True


def main() -> None:
    args = parse_args()
    if not args.src_root.exists():
        raise SystemExit(f"src root not found: {args.src_root}")

    written = 0
    linked = 0
    for src_path in sorted(args.src_root.glob("*/*/*.pbf")):
        dst_path = args.dst_root / src_path.relative_to(args.src_root)
        changed = normalize_tile(src_path, dst_path, args.target_extent, args.rewrite_all)
        if changed:
            written += 1
        else:
            linked += 1

    print(
        f"normalize complete: src={args.src_root} dst={args.dst_root} "
        f"rewritten={written} reused={linked} target_extent={args.target_extent}"
    )


if __name__ == "__main__":
    main()