#!/usr/bin/env python3
"""Build lightweight, confidence-aware aNavigator road geometry with osm2streets."""
import argparse
import json
from pathlib import Path
import xml.etree.ElementTree as ET

import osm2streets_python


def compact_geometry(geometry):
    """Keep about 10 cm precision while avoiding verbose Rust float output."""
    def compact(value):
        if isinstance(value, float):
            return round(value, 6)
        if isinstance(value, list):
            return [compact(child) for child in value]
        return value
    return {"type": geometry.get("type"), "coordinates": compact(geometry.get("coordinates", []))}


def compact_feature(geometry, properties):
    return {"type": "Feature", "geometry": compact_geometry(geometry), "properties": properties}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input_osm", type=Path)
    parser.add_argument("output_geojson", type=Path)
    args = parser.parse_args()

    osm_bytes = args.input_osm.read_bytes()
    root = ET.fromstring(osm_bytes)
    way_tags = {}
    explicit_lanes = set()
    explicit_width = set()
    explicit_markings = set()
    for way in root.findall("way"):
        way_id = int(way.attrib["id"])
        tags = {t.attrib["k"]: t.attrib["v"] for t in way.findall("tag")}
        way_tags[way_id] = tags
        if tags.get("lanes") or tags.get("lanes:forward") or tags.get("lanes:backward"):
            explicit_lanes.add(way_id)
        if tags.get("width") or tags.get("width:carriageway"):
            explicit_width.add(way_id)
        if tags.get("lane_markings") is not None or tags.get("divider") is not None:
            explicit_markings.add(way_id)

    options = json.dumps({
        "debug_each_step": False,
        "dual_carriageway_experiment": False,
        "sidepath_zipping_experiment": False,
        "inferred_sidewalks": False,
        "inferred_kerbs": False,
        "date_time": None,
        "override_driving_side": "Right",
    })
    network = osm2streets_python.PyStreetNetwork(osm_bytes, "", options)
    roads = json.loads(network.to_geojson_plain())
    markings = json.loads(network.to_lane_markings_geojson())
    intersections = json.loads(network.to_intersection_markings_geojson())

    output = []
    road_confidence = {}
    for feature in roads.get("features", []):
        props = feature.get("properties") or {}
        if props.get("type") == "road":
            ids = [int(x) for x in props.get("osm_way_ids", [])]
            if ids and all(x in explicit_width for x in ids):
                confidence = "verified_width"
            elif ids and all(x in explicit_lanes for x in ids):
                confidence = "verified_lanes"
            else:
                confidence = "inferred_width"
            road_confidence[int(props.get("id", -1))] = confidence
            output.append(compact_feature(feature["geometry"], {
                "layer": "asphalt", "confidence": confidence,
                "osm_way_ids": ids,
            }))
        elif props.get("type") == "intersection":
            output.append(compact_feature(feature["geometry"], {
                "layer": "intersection", "confidence": "derived_osm_geometry"
            }))

    # Group thousands of dash polygons by road/type into MultiPolygon features.
    groups = {}
    skipped_inferred = 0
    for feature in markings.get("features", []):
        props = feature.get("properties") or {}
        marking_type = props.get("type")
        if marking_type not in {"center line", "lane separator", "vehicle stop line", "bike stop line"}:
            continue
        ids = tuple(int(x) for x in props.get("osm_way_ids", []))
        road_id = int(props.get("road", -1))
        if marking_type in {"center line", "lane separator"}:
            if not ids or not all(x in explicit_lanes for x in ids):
                skipped_inferred += 1
                continue
            confidence = "verified_marking" if all(x in explicit_markings for x in ids) else "verified_lanes_standard_marking"
        else:
            confidence = "verified_osm_control"
        geometry = feature.get("geometry") or {}
        polygons = geometry.get("coordinates", [])
        if geometry.get("type") == "Polygon":
            polygons = [polygons]
        elif geometry.get("type") != "MultiPolygon":
            continue
        key = (road_id, marking_type, confidence, props.get("highway", ""), ids)
        groups.setdefault(key, []).extend(polygons)

    for (road_id, marking_type, confidence, highway, ids), polygons in groups.items():
        output.append(compact_feature(
            {"type": "MultiPolygon", "coordinates": polygons},
            {"layer": "road_marking", "marking": marking_type,
             "confidence": confidence, "road": road_id,
             "highway": highway, "osm_way_ids": list(ids)}
        ))

    for feature in intersections.get("features", []):
        props = feature.get("properties") or {}
        if props.get("type") != "marked crossing line":
            continue
        output.append(compact_feature(feature["geometry"], {
            "layer": "crosswalk", "confidence": "derived_from_osm_crossing"
        }))

    metadata = {
        "generator": "osm2streets/aNavigator",
        "source": "OpenStreetMap",
        "explicit_lane_ways": len(explicit_lanes),
        "explicit_marking_ways": len(explicit_markings),
        "skipped_inferred_marking_polygons": skipped_inferred,
        "features": len(output),
    }
    result = {"type": "FeatureCollection", "anavigator": metadata, "features": output}
    args.output_geojson.parent.mkdir(parents=True, exist_ok=True)
    args.output_geojson.write_text(json.dumps(result, separators=(",", ":")), encoding="utf-8")
    print(json.dumps(metadata, indent=2))
    print(f"output_bytes={args.output_geojson.stat().st_size}")


if __name__ == "__main__":
    main()