#!/usr/bin/env python3
"""Turn osmium GeoJSONSeq walking ways into small, versioned graph tiles."""

from __future__ import annotations

import argparse
import gzip
import json
import math
import os
import shutil
import subprocess
import sys
from itertools import groupby
from pathlib import Path

SCHEMA_VERSION = 1
SCALE = 1_000_000
DEFAULT_TILE_DEGREES = 0.05
DEFAULT_SPEED_MPS = 1.35
MAX_EDGE_LENGTH_M = 100
PARTITION_COUNT = 64
BLOCKED_HIGHWAYS = {
    "construction",
    "corridor",
    "motorway",
    "motorway_link",
    "proposed",
    "raceway",
}
BLOCKED_ACCESS = {"no", "private"}
FOOT_OVERRIDE = {"yes", "designated", "permissive", "destination"}
ROUGH_SURFACES = {
    "cobblestone",
    "compacted",
    "dirt",
    "earth",
    "fine_gravel",
    "gravel",
    "ground",
    "mud",
    "pebblestone",
    "sand",
    "unpaved",
}


def tags_for(feature: dict) -> dict:
    properties = feature.get("properties") or {}
    nested = properties.get("tags")
    return {**properties, **(nested if isinstance(nested, dict) else {})}


def walking_profile(tags: dict) -> tuple[float, int] | None:
    highway = str(tags.get("highway") or "").lower()
    route = str(tags.get("route") or "").lower()
    foot = str(tags.get("foot") or "").lower()
    access = str(tags.get("access") or "").lower()
    if str(tags.get("area") or "").lower() == "yes":
        return None
    if foot in BLOCKED_ACCESS:
        return None
    if access in BLOCKED_ACCESS and foot not in FOOT_OVERRIDE:
        return None
    if route == "ferry":
        if foot in BLOCKED_ACCESS:
            return None
        speed = 5.0
    elif not highway or highway in BLOCKED_HIGHWAYS:
        return None
    elif highway == "steps":
        speed = 0.70
    elif highway in {"path", "track", "bridleway"}:
        speed = 1.10
    elif str(tags.get("surface") or "").lower() in ROUGH_SURFACES:
        speed = 1.0
    else:
        speed = DEFAULT_SPEED_MPS
    one_way = str(tags.get("oneway:foot") or "").lower()
    direction = 1 if one_way in {"yes", "1", "true"} else 2 if one_way in {"-1", "reverse"} else 0
    return speed, direction


def blocked_barrier(tags: dict) -> bool:
    barrier = str(tags.get("barrier") or "").lower()
    foot = str(tags.get("foot") or "").lower()
    access = str(tags.get("access") or "").lower()
    return bool(barrier) and (foot in BLOCKED_ACCESS or access in BLOCKED_ACCESS)


def coordinate_lines(geometry: dict):
    geometry_type = geometry.get("type")
    coordinates = geometry.get("coordinates")
    if geometry_type == "LineString" and isinstance(coordinates, list):
        yield coordinates
    elif geometry_type == "MultiLineString" and isinstance(coordinates, list):
        yield from (line for line in coordinates if isinstance(line, list))


def haversine_metres(left: tuple[float, float], right: tuple[float, float]) -> float:
    lon1, lat1 = left
    lon2, lat2 = right
    radians = math.pi / 180
    latitude_delta = (lat2 - lat1) * radians
    longitude_delta = (lon2 - lon1) * radians
    latitude1 = lat1 * radians
    latitude2 = lat2 * radians
    value = math.sin(latitude_delta / 2) ** 2 + math.cos(latitude1) * math.cos(latitude2) * math.sin(longitude_delta / 2) ** 2
    return 2 * 6_371_008.8 * math.asin(min(1, math.sqrt(value)))


def tile_key(latitude: float, longitude: float, tile_degrees: float) -> str:
    latitude_index = math.floor((latitude + 90) / tile_degrees)
    longitude_index = math.floor((longitude + 180) / tile_degrees)
    return f"{latitude_index}_{longitude_index}"


def partition_for(tile: str) -> int:
    latitude, longitude = (int(value) for value in tile.split("_", 1))
    return ((latitude * 73_856_093) ^ (longitude * 19_349_663)) % PARTITION_COUNT


def write_tile(output: Path, version: str, tile: str, rows) -> tuple[int, int]:
    nodes: list[list[int]] = []
    node_indexes: dict[tuple[int, int], int] = {}
    edges: list[list[int]] = []

    def node_index(latitude: int, longitude: int) -> int:
        key = (latitude, longitude)
        found = node_indexes.get(key)
        if found is not None:
            return found
        found = len(nodes)
        node_indexes[key] = found
        nodes.append([latitude, longitude])
        return found

    for row in rows:
        parts = row.rstrip("\n").split("\t")
        if len(parts) != 9:
            continue
        _, lat1, lon1, lat2, lon2, distance, duration, direction, _source = parts
        start = node_index(int(lat1), int(lon1))
        end = node_index(int(lat2), int(lon2))
        if start == end:
            continue
        edges.append([start, end, int(distance), int(duration), int(direction)])

    payload = {
        "schemaVersion": SCHEMA_VERSION,
        "version": version,
        "tile": tile,
        "scale": SCALE,
        "nodes": nodes,
        "edges": edges,
    }
    filename = output / "tiles" / f"{tile}.json.gz"
    filename.parent.mkdir(parents=True, exist_ok=True)
    with gzip.open(filename, "wt", encoding="utf-8", compresslevel=9) as handle:
        json.dump(payload, handle, ensure_ascii=False, separators=(",", ":"))
    return len(nodes), len(edges)


def build(args: argparse.Namespace) -> dict:
    output = Path(args.output)
    if output.exists():
        shutil.rmtree(output)
    output.mkdir(parents=True)
    partitions = output / ".partitions"
    partitions.mkdir()
    handles = [open(partitions / f"{index:02d}.tsv", "w", encoding="utf-8") for index in range(PARTITION_COUNT)]
    source_ways = 0
    source_segments = 0
    blocked_points: set[tuple[int, int]] = set()
    try:
        for raw_line in sys.stdin:
            line = raw_line.lstrip("\x1e").strip()
            if not line:
                continue
            try:
                feature = json.loads(line)
            except json.JSONDecodeError:
                continue
            tags = tags_for(feature)
            geometry = feature.get("geometry") or {}
            if geometry.get("type") == "Point" and blocked_barrier(tags):
                coordinates = geometry.get("coordinates") or []
                if isinstance(coordinates, list) and len(coordinates) >= 2:
                    try:
                        blocked_points.add((round(float(coordinates[1]) * SCALE), round(float(coordinates[0]) * SCALE)))
                    except (TypeError, ValueError):
                        pass
                continue
            profile = walking_profile(tags)
            if not profile:
                continue
            speed, direction = profile
            accepted = False
            for coordinates in coordinate_lines(geometry):
                for start, end in zip(coordinates, coordinates[1:]):
                    if not (
                        isinstance(start, list)
                        and isinstance(end, list)
                        and len(start) >= 2
                        and len(end) >= 2
                    ):
                        continue
                    try:
                        left = (float(start[0]), float(start[1]))
                        right = (float(end[0]), float(end[1]))
                    except (TypeError, ValueError):
                        continue
                    distance = haversine_metres(left, right)
                    if not math.isfinite(distance) or distance < 0.5:
                        continue
                    part_count = max(1, math.ceil(distance / MAX_EDGE_LENGTH_M))
                    source = str(feature.get("id") or (feature.get("properties") or {}).get("@id") or "")[:40]
                    for part in range(part_count):
                        start_fraction = part / part_count
                        end_fraction = (part + 1) / part_count
                        part_left = (
                            left[0] + (right[0] - left[0]) * start_fraction,
                            left[1] + (right[1] - left[1]) * start_fraction,
                        )
                        part_right = (
                            left[0] + (right[0] - left[0]) * end_fraction,
                            left[1] + (right[1] - left[1]) * end_fraction,
                        )
                        left_key = (round(part_left[1] * SCALE), round(part_left[0] * SCALE))
                        right_key = (round(part_right[1] * SCALE), round(part_right[0] * SCALE))
                        if left_key in blocked_points or right_key in blocked_points:
                            continue
                        part_distance = distance / part_count
                        tile = tile_key(
                            (part_left[1] + part_right[1]) / 2,
                            (part_left[0] + part_right[0]) / 2,
                            args.tile_degrees,
                        )
                        handle = handles[partition_for(tile)]
                        handle.write(
                            f"{tile}\t{left_key[0]}\t{left_key[1]}\t{right_key[0]}\t{right_key[1]}\t"
                            f"{max(1, round(part_distance))}\t{max(1, math.ceil(part_distance / speed))}\t"
                            f"{direction}\t{source}\n"
                        )
                        source_segments += 1
                        accepted = True
            if accepted:
                source_ways += 1
    finally:
        for handle in handles:
            handle.close()

    tiles: list[str] = []
    node_total = 0
    edge_total = 0
    for partition in sorted(partitions.glob("*.tsv")):
        if partition.stat().st_size == 0:
            partition.unlink()
            continue
        sorted_partition = partition.with_suffix(".sorted.tsv")
        environment = {**os.environ, "LC_ALL": "C"}
        with open(sorted_partition, "w", encoding="utf-8") as sorted_handle:
            subprocess.run(
                ["sort", "--buffer-size=64M", "-t", "\t", "-k1,1", str(partition)],
                check=True,
                stdout=sorted_handle,
                env=environment,
            )
        with open(sorted_partition, encoding="utf-8") as rows:
            for tile, grouped in groupby(rows, key=lambda row: row.split("\t", 1)[0]):
                node_count, edge_count = write_tile(output, args.version, tile, grouped)
                tiles.append(tile)
                node_total += node_count
                edge_total += edge_count
        partition.unlink()
        sorted_partition.unlink()
    partitions.rmdir()

    tiles.sort()
    manifest = {
        "schemaVersion": SCHEMA_VERSION,
        "version": args.version,
        "generatedAt": args.generated_at,
        "tileDegrees": args.tile_degrees,
        "tileUrlTemplate": "/api/walking-atlas/tiles/{version}/{tile}.json.gz",
        "maxRouteDistanceM": args.max_route_distance,
        "maxSnapDistanceM": args.max_snap_distance,
        "walkingSpeedMps": DEFAULT_SPEED_MPS,
        "tiles": tiles,
        "tileCount": len(tiles),
        "nodeReferences": node_total,
        "edgeCount": edge_total,
        "sourceWayCount": source_ways,
        "sourceSegmentCount": source_segments,
        "coverage": "United Kingdom",
        "sourceName": "OpenStreetMap contributors",
        "sourceUrl": "https://www.openstreetmap.org/copyright",
        "sourceLicence": "Open Database Licence",
    }
    with open(output / "manifest.json", "w", encoding="utf-8") as handle:
        json.dump(manifest, handle, ensure_ascii=False, separators=(",", ":"))
    return manifest


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", required=True)
    parser.add_argument("--version", required=True)
    parser.add_argument("--generated-at", required=True)
    parser.add_argument("--tile-degrees", type=float, default=DEFAULT_TILE_DEGREES)
    parser.add_argument("--max-route-distance", type=int, default=3_000)
    parser.add_argument("--max-snap-distance", type=int, default=250)
    return parser.parse_args()


if __name__ == "__main__":
    result = build(parse_args())
    print(
        f"Walking atlas built: {result['tileCount']} tiles, "
        f"{result['edgeCount']} edges, version {result['version']}",
        file=sys.stderr,
    )