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

import argparse
from pathlib import Path

from mapbox_vector_tile.Mapbox import vector_tile_pb2


FINAL_RELEASE_PROPERTY_ALLOWLIST = frozenset(
    {
        "canonical_object_type",
        "chart_fill_pattern",
        "chart_fill_style",
        "chart_icon_image",
        "chart_label_position_code",
        "chart_label_subtext",
        "chart_label_text",
        "chart_line_color",
        "chart_line_width",
        "chart_symbol_code",
        "chart_text_color",
        "chart_text_style",
        "depth_value_m",
        "least_depth_m",
        "light_color_code",
        "light_sector_mode",
        "name_ja",
        "place_name_en",
    }
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Create a final release tile set by stripping non-style properties from an existing PBF tree."
    )
    parser.add_argument("--src-root", type=Path, required=True)
    parser.add_argument("--dst-root", type=Path, required=True)
    return parser.parse_args()


def clone_value(dst_value: vector_tile_pb2.tile.value, src_value: vector_tile_pb2.tile.value) -> None:
    dst_value.CopyFrom(src_value)


def strip_tile_properties(raw_tile: bytes, allowlist: set[str] | frozenset[str]) -> bytes:
    src_tile = vector_tile_pb2.tile()
    src_tile.ParseFromString(raw_tile)

    dst_tile = vector_tile_pb2.tile()

    for src_layer in src_tile.layers:
        dst_layer = dst_tile.layers.add()
        dst_layer.version = src_layer.version
        dst_layer.name = src_layer.name
        dst_layer.extent = src_layer.extent

        kept_key_index: dict[int, int] = {}
        kept_value_index: dict[tuple[int, bytes], int] = {}

        for src_feature in src_layer.features:
            dst_feature = dst_layer.features.add()
            if src_feature.HasField("id"):
                dst_feature.id = src_feature.id
            dst_feature.type = src_feature.type
            dst_feature.geometry.extend(src_feature.geometry)

            tags = list(src_feature.tags)
            for pos in range(0, len(tags), 2):
                src_key_idx = tags[pos]
                src_val_idx = tags[pos + 1]
                key = src_layer.keys[src_key_idx]
                if key not in allowlist:
                    continue

                dst_key_idx = kept_key_index.get(src_key_idx)
                if dst_key_idx is None:
                    dst_key_idx = len(dst_layer.keys)
                    dst_layer.keys.append(key)
                    kept_key_index[src_key_idx] = dst_key_idx

                src_value = src_layer.values[src_val_idx]
                value_sig = (src_val_idx, src_value.SerializeToString())
                dst_val_idx = kept_value_index.get(value_sig)
                if dst_val_idx is None:
                    dst_val_idx = len(dst_layer.values)
                    clone_value(dst_layer.values.add(), src_value)
                    kept_value_index[value_sig] = dst_val_idx

                dst_feature.tags.extend([dst_key_idx, dst_val_idx])

    return dst_tile.SerializeToString()


def main() -> None:
    args = parse_args()
    src_root = args.src_root
    dst_root = args.dst_root

    if not src_root.exists():
        raise SystemExit(f"src root not found: {src_root}")

    if dst_root.exists():
        for path in sorted(dst_root.glob("**/*"), reverse=True):
            if path.is_file() or path.is_symlink():
                path.unlink()
            elif path.is_dir():
                try:
                    path.rmdir()
                except OSError:
                    pass
    dst_root.mkdir(parents=True, exist_ok=True)

    written = 0
    for src_path in sorted(src_root.glob("*/*/*.pbf")):
        rel = src_path.relative_to(src_root)
        dst_path = dst_root / rel
        dst_path.parent.mkdir(parents=True, exist_ok=True)
        dst_path.write_bytes(strip_tile_properties(src_path.read_bytes(), FINAL_RELEASE_PROPERTY_ALLOWLIST))
        written += 1

    print(f"minimized release tiles written={written} dst={dst_root}")


if __name__ == "__main__":
    main()