#!/usr/bin/env python3
"""Generate lightweight osm2streets GeoJSON tiles for Bologna and its outskirts."""
import argparse
import json
import math
import subprocess
import sys
import tempfile
import hashlib
import os
from pathlib import Path
from xml.sax.saxutils import quoteattr

import osmium


DEFAULT_BBOX = (11.18, 44.37, 11.52, 44.62)


class RoadCollector(osmium.SimpleHandler):
    def __init__(self, bounds):
        super().__init__()
        self.bounds = bounds
        self.nodes = {}
        self.node_tags = {}
        self.ways = []

    def node(self, node):
        if not node.tags:
            return
        try:
            lon, lat = node.location.lon, node.location.lat
        except osmium.InvalidLocationError:
            return
        west, south, east, north = self.bounds
        if west-0.01 <= lon <= east+0.01 and south-0.01 <= lat <= north+0.01:
            self.node_tags[node.id] = dict(node.tags)

    def way(self, way):
        tags = dict(way.tags)
        if "highway" not in tags:
            return
        refs, coords = [], []
        try:
            for node in way.nodes:
                refs.append(node.ref)
                coords.append((node.lon, node.lat))
        except osmium.InvalidLocationError:
            return
        if len(coords) < 2:
            return
        west, south, east, north = self.bounds
        xs = [p[0] for p in coords]
        ys = [p[1] for p in coords]
        if max(xs) < west or min(xs) > east or max(ys) < south or min(ys) > north:
            return
        for ref, coord in zip(refs, coords):
            self.nodes[ref] = coord
        self.ways.append((way.id, refs, tags, (min(xs), min(ys), max(xs), max(ys))))


def write_osm(path, nodes, ways, node_tags):
    with path.open("w", encoding="utf-8") as f:
        f.write('<?xml version="1.0" encoding="UTF-8"?>\n<osm version="0.6" generator="aNavigator">\n')
        for node_id in sorted(nodes):
            lon, lat = nodes[node_id]
            tags = node_tags.get(node_id, {})
            if tags:
                f.write(f'<node id="{node_id}" lat="{lat:.7f}" lon="{lon:.7f}">\n')
                for key, value in tags.items():
                    f.write(f'<tag k={quoteattr(key)} v={quoteattr(value)}/>\n')
                f.write('</node>\n')
            else:
                f.write(f'<node id="{node_id}" lat="{lat:.7f}" lon="{lon:.7f}"/>\n')
        for way_id, refs, tags, _ in ways:
            f.write(f'<way id="{way_id}">\n')
            for ref in refs:
                f.write(f'<nd ref="{ref}"/>\n')
            for key, value in tags.items():
                f.write(f'<tag k={quoteattr(key)} v={quoteattr(value)}/>\n')
            f.write('</way>\n')
        f.write('</osm>\n')


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


def geometry_bbox(geometry):
    coords = geometry.get("coordinates", [])
    xs, ys = [], []
    def visit(value):
        if value and isinstance(value[0], (int, float)):
            xs.append(value[0]); ys.append(value[1])
        else:
            for child in value:
                visit(child)
    if coords:
        visit(coords)
    return (min(xs), min(ys), max(xs), max(ys)) if xs else None


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input_pbf", type=Path)
    parser.add_argument("output_dir", type=Path)
    parser.add_argument("--builder", type=Path, required=True)
    parser.add_argument("--tile-size", type=float, default=0.025)
    parser.add_argument("--batch-size", type=float, default=0.1)
    parser.add_argument("--bbox", nargs=4, type=float, default=DEFAULT_BBOX)
    args = parser.parse_args()
    bounds = tuple(args.bbox)
    collector = RoadCollector(bounds)
    collector.apply_file(str(args.input_pbf), locations=True, idx="flex_mem")
    print(f"Selected {len(collector.ways)} highway ways and {len(collector.nodes)} nodes")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    tiles = {}
    west, south, east, north = bounds
    cols = math.ceil((east - west) / args.tile_size)
    rows = math.ceil((north - south) / args.tile_size)
    for y in range(rows):
        for x in range(cols):
            tw = west + x * args.tile_size
            ts = south + y * args.tile_size
            tiles[(x, y)] = {"bbox": [tw, ts, min(tw + args.tile_size, east), min(ts + args.tile_size, north)], "features": [], "seen": set()}

    with tempfile.TemporaryDirectory(prefix="anavigator-roads-") as temp_name:
        temp = Path(temp_name)
        sequence = [0]
        failures = []
        def process_batch(core_bbox, depth=0):
            sequence[0] += 1
            pad = 0.003
            batch_bbox = (core_bbox[0]-pad, core_bbox[1]-pad, core_bbox[2]+pad, core_bbox[3]+pad)
            ways = [way for way in collector.ways if intersects(way[3], batch_bbox)]
            if not ways:
                return
            node_ids = {ref for _, refs, _, _ in ways for ref in refs}
            stem = f"batch_{sequence[0]}_{depth}"
            osm_path, geo_path = temp / f"{stem}.osm", temp / f"{stem}.geojson"
            write_osm(osm_path, {ref: collector.nodes[ref] for ref in node_ids}, ways, collector.node_tags)
            env = dict(os.environ, RUST_LOG="off")
            result = subprocess.run([sys.executable, str(args.builder), str(osm_path), str(geo_path)],
                                    env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if result.returncode != 0:
                width, height = core_bbox[2]-core_bbox[0], core_bbox[3]-core_bbox[1]
                if depth >= 3 or max(width, height) <= args.tile_size:
                    failures.append(list(core_bbox))
                    print(f"Fallback for malformed OSM geometry in {core_bbox}")
                    return
                mx, my = (core_bbox[0]+core_bbox[2])/2, (core_bbox[1]+core_bbox[3])/2
                print(f"Subdividing problematic batch {core_bbox}")
                for child in ((core_bbox[0],core_bbox[1],mx,my),(mx,core_bbox[1],core_bbox[2],my),
                              (core_bbox[0],my,mx,core_bbox[3]),(mx,my,core_bbox[2],core_bbox[3])):
                    process_batch(child, depth+1)
                return
            data = json.loads(geo_path.read_text(encoding="utf-8"))
            for feature in data.get("features", []):
                fb = geometry_bbox(feature.get("geometry") or {})
                if not fb:
                    continue
                cx, cy = (fb[0]+fb[2])/2, (fb[1]+fb[3])/2
                if not (core_bbox[0] <= cx <= core_bbox[2] and core_bbox[1] <= cy <= core_bbox[3]):
                    continue
                tx = min(cols-1, max(0, int((cx-west)/args.tile_size)))
                ty = min(rows-1, max(0, int((cy-south)/args.tile_size)))
                tile = tiles.get((tx,ty))
                if tile is not None:
                    encoded = json.dumps(feature, sort_keys=True, separators=(",", ":")).encode()
                    fingerprint = hashlib.sha1(encoded).digest()
                    if fingerprint not in tile["seen"]:
                        tile["seen"].add(fingerprint)
                        tile["features"].append(feature)
            print(f"Processed {core_bbox}: {len(data.get('features', []))} features")

        batch_cols = math.ceil((east-west)/args.batch_size)
        batch_rows = math.ceil((north-south)/args.batch_size)
        for by in range(batch_rows):
            for bx in range(batch_cols):
                bw, bs = west+bx*args.batch_size, south+by*args.batch_size
                process_batch((bw, bs, min(bw+args.batch_size,east), min(bs+args.batch_size,north)))

    manifest_tiles = []
    total_bytes = total_features = 0
    for (x, y), tile in tiles.items():
        if not tile["features"]:
            continue
        name = f"roads_{x}_{y}.geojson"
        payload = {"type": "FeatureCollection", "features": tile["features"]}
        path = args.output_dir / name
        path.write_text(json.dumps(payload, separators=(",", ":")), encoding="utf-8")
        size = path.stat().st_size
        total_bytes += size
        total_features += len(tile["features"])
        manifest_tiles.append({"id": f"{x}_{y}", "file": name, "bbox": tile["bbox"], "bytes": size})
    manifest = {"version": 1, "source": "OpenStreetMap/osm2streets", "coverage": list(bounds),
                "tileSize": args.tile_size, "tiles": manifest_tiles,
                "features": total_features, "bytes": total_bytes, "fallbackAreas": failures}
    (args.output_dir / "manifest.json").write_text(json.dumps(manifest, separators=(",", ":")), encoding="utf-8")
    print(json.dumps({"tiles": len(manifest_tiles), "features": total_features, "bytes": total_bytes}, indent=2))


if __name__ == "__main__":
    main()