"""
Extract tram routes from OSM and dump them into GeoJson.
"""

import json
from pathlib import Path

import duckdb


ROOT_DIR = Path(__file__).parent


def main():
    pbf = ROOT_DIR / "croatia-latest.osm.pbf"

    duckdb.install_extension("spatial")
    duckdb.load_extension("spatial")

    print("Loading relations...")
    relations = duckdb.query(f"""
        SELECT tags, refs, ref_roles, ref_types
        FROM ST_ReadOSM('{pbf}')
        WHERE kind='relation'
          AND tags.route = 'tram'
          AND tags.operator = 'ZET'
    """).fetchall()

    print("Loading ways...")
    way_ids = flatten(get_way_ids(r) for r in relations)
    ways = duckdb.query(f"""
        SELECT id, refs
        FROM ST_ReadOSM('{pbf}')
        WHERE kind='way'
          AND id IN {way_ids}
    """).fetchall()
    ways_map = {id: refs for id, refs in ways}

    print("Loading nodes...")
    node_ids = flatten([w[1] for w in ways])
    nodes = duckdb.query(f"""
        SELECT id, lat, lon
        FROM ST_ReadOSM('{pbf}')
        WHERE kind='node'
          AND id IN {node_ids}
    """).fetchall()
    nodes_map = {id: (lat, lon) for id, lat, lon in nodes}

    for relation in relations:
        collection = {"type": "FeatureCollection", "features": []}
        filename = f"{relation[0]['route']}_{relation[0]['ref'].rjust(2, '0')}.geojson"
        print(f"Writing {filename}")

        way_ids = get_way_ids(relation)
        for way_id in way_ids:
            coordinates = []
            node_ids = ways_map[way_id]
            for node_id in node_ids:
                lat, lon = nodes_map[node_id]
                coordinates.append([lon, lat])
            collection["features"].append(
                {
                    "type": "Feature",
                    "geometry": {
                        "type": "LineString",
                        "coordinates": coordinates,
                    },
                }
            )
            with open(ROOT_DIR / filename, "w") as f:
                json.dump(collection, f, indent=4)


def get_way_ids(relation):
    _, refs, roles, types = relation
    for ref, role, type in zip(refs, roles, types):
        if role is None and type == "way":
            yield ref


def flatten(xss):
    return [x for xs in xss for x in xs]


if __name__ == "__main__":
    main()
