"""
OpenStreetMap utilities.

Handles geocoding and two-pass PBF parsing (WayFilterHandler, NodeCoordinateHandler,
write_network_parquet) for building walkingpandas Network graphs from OSM data.
"""

import logging
import os
import shutil
import subprocess
from pathlib import Path
from typing import TYPE_CHECKING, List, Optional, Set, Tuple

import requests

if TYPE_CHECKING:
    from walkingpandas._utils import BoundingBox

logger = logging.getLogger(__name__)


def geocode_place_to_bbox(place: str) -> "BoundingBox":
    """Geocode a place name to a bounding box using Nominatim.

    Args:
        place: Free-text place name (e.g. ``"Zurich, Switzerland"``).

    Returns:
        BoundingBox (north, south, east, west) in decimal degrees.

    Raises:
        ValueError: If Nominatim returns no results for *place*.
    """
    from walkingpandas._utils import BoundingBox

    logger.debug("Geocoding place '%s' via Nominatim...", place)
    r = requests.get(
        "https://nominatim.openstreetmap.org/search",
        params={"q": place, "format": "json", "limit": 1},
        headers={"User-Agent": "walkingpandas-valhalla/1.0"},
        timeout=30,
    )
    r.raise_for_status()
    data = r.json()

    if not data:
        raise ValueError(f"Place not found: {place}")

    b = data[0].get("boundingbox")
    if not b or len(b) < 4:
        lat, lon = float(data[0]["lat"]), float(data[0]["lon"])
        delta = 0.01
        north, south, east, west = lat + delta, lat - delta, lon + delta, lon - delta
    else:
        # Nominatim returns [south, north, west, east] as strings
        south, north, west, east = (
            float(b[0]),
            float(b[1]),
            float(b[2]),
            float(b[3]),
        )
    return BoundingBox.from_nsew(north, south, east, west)


# ---------------------------------------------------------------------------
# Osmium CLI helpers (optional fast-path)
# ---------------------------------------------------------------------------


def _osmium_available() -> bool:
    """Return ``True`` if the ``osmium`` command-line tool is on PATH."""
    return shutil.which("osmium") is not None


def osmium_tags_filter(
    input_pbf: str, output_pbf: str, *expressions: str
) -> bool:
    """Run ``osmium tags-filter`` with the given tag expressions.

    Example::

        osmium_tags_filter("in.osm.pbf", "out.osm.pbf", "w/highway")

    Returns ``True`` on success, ``False`` if the CLI is missing or fails.
    """
    if not _osmium_available():
        return False
    cmd = [
        "osmium", "tags-filter", input_pbf, *expressions,
        "-o", output_pbf, "--overwrite",
    ]
    try:
        logger.debug("Running: %s", " ".join(cmd))
        subprocess.run(cmd, check=True, capture_output=True)
        return True
    except (subprocess.CalledProcessError, OSError) as exc:
        logger.warning("osmium tags-filter failed: %s", exc)
        return False


def osmium_extract_bbox(
    input_pbf: str, output_pbf: str, bbox: "BoundingBox"
) -> bool:
    """Run ``osmium extract -b`` to clip a PBF to a bounding box.

    Returns ``True`` on success, ``False`` if the CLI is missing or fails.
    """
    if not _osmium_available():
        return False
    cmd = [
        "osmium", "extract",
        "-b", bbox.as_osmium(),
        input_pbf,
        "-o", output_pbf, "--overwrite",
    ]
    try:
        logger.debug("Running: %s", " ".join(cmd))
        subprocess.run(cmd, check=True, capture_output=True)
        return True
    except (subprocess.CalledProcessError, OSError) as exc:
        logger.warning("osmium extract failed: %s", exc)
        return False


# ---------------------------------------------------------------------------
# Two-pass PBF parsing for Network.from_osm()
# ---------------------------------------------------------------------------

_ONEWAY_FORWARD = {"yes", "true", "1"}

DEFAULT_EDGE_ATTRIBUTES: Tuple[str, ...] = ("way_id", "oneway")


def _resolve_direction(
    oneway: str, force_bidirectional: bool
) -> Tuple[bool, bool]:
    """Return ``(add_forward, add_reverse)`` for a way's oneway tag."""
    if force_bidirectional:
        return True, True
    if oneway in _ONEWAY_FORWARD:
        return True, False
    if oneway == "-1":
        return False, True
    return True, True


class WayFilterHandler:
    """Pass 1: collect highway ways and the set of node IDs they reference.

    Args:
        highway_filter: If provided, only ways whose ``highway`` tag value is
            in this set will be kept. ``None`` means accept all highway types.
        edge_attributes: OSM way tags to store per edge.  Defaults to
            :data:`DEFAULT_EDGE_ATTRIBUTES` (``("way_id", "oneway")``).
            ``way_id``, ``highway``, and ``oneway`` are always collected as
            structural fields; any additional names are read from OSM tags.
    """

    def __init__(
        self,
        highway_filter: Optional[Set[str]] = None,
        edge_attributes: Optional[Tuple[str, ...]] = None,
    ):
        try:
            import osmium
        except ImportError as e:
            raise ImportError(
                "The 'osmium' package is required for from_osm(). "
                "Install it via: pip install pyosmium"
            ) from e

        self._handler = osmium.SimpleHandler()
        self._handler.way = self._way
        self.highway_filter = highway_filter
        self.edge_attributes: Tuple[str, ...] = (
            tuple(edge_attributes) if edge_attributes is not None
            else DEFAULT_EDGE_ATTRIBUTES
        )
        self.needed_nodes: Set[int] = set()
        self.ways: List[dict] = []

    def _way(self, w):
        highway = w.tags.get("highway")
        if not highway:
            return
        if w.tags.get("area") == "yes":
            return
        if self.highway_filter is not None and highway not in self.highway_filter:
            return

        node_refs = [n.ref for n in w.nodes]
        self.needed_nodes.update(node_refs)

        _STRUCTURAL = {"way_id", "highway", "oneway", "nodes"}
        extra_tags = {
            attr: w.tags.get(attr, "")
            for attr in self.edge_attributes
            if attr not in _STRUCTURAL
        }

        self.ways.append({
            "way_id": w.id,
            "highway": highway,
            "oneway": w.tags.get("oneway", "no"),
            "nodes": node_refs,
            **extra_tags,
        })

    def apply_file(self, path: str):
        self._handler.apply_file(path)


class NodeCoordinateHandler:
    """Pass 2: collect ``(lon, lat)`` for every node in *needed_nodes*.

    Args:
        needed_nodes: Set of node IDs to collect coordinates for.
        bbox: If provided, only nodes within this bounding box are kept.
            This is the pure-Python fallback for geographic filtering
            when the osmium CLI is not available.
        node_attributes: Optional tuple of OSM node tag names to collect.
            When set, ``self.node_attrs`` maps each node ID to a dict of
            tag key/value pairs.  ``None`` (default) skips tag collection.
    """

    def __init__(
        self,
        needed_nodes: Set[int],
        bbox: Optional["BoundingBox"] = None,
        node_attributes: Optional[Tuple[str, ...]] = None,
    ):
        try:
            import osmium
        except ImportError as e:
            raise ImportError(
                "The 'osmium' package is required for from_osm(). "
                "Install it via: pip install pyosmium"
            ) from e

        self._handler = osmium.SimpleHandler()
        self._handler.node = self._node
        self.needed_nodes = needed_nodes
        self.bbox = bbox
        self.node_attributes: Tuple[str, ...] = (
            tuple(node_attributes) if node_attributes else ()
        )
        self.node_coords: dict = {}
        self.node_attrs: dict = {}

    def _node(self, n):
        if n.id in self.needed_nodes:
            lon, lat = n.location.lon, n.location.lat
            if self.bbox is not None and not self.bbox.contains(lon, lat):
                return
            self.node_coords[n.id] = (lon, lat)
            if self.node_attributes:
                self.node_attrs[n.id] = {
                    k: n.tags.get(k, "") for k in self.node_attributes
                }

    def apply_file(self, path: str):
        self._handler.apply_file(path)


def _get_way_attr_value(way: dict, attr: str) -> object:
    """Resolve an attribute value from a way dict.

    Structural fields (``way_id``, ``highway``, ``oneway``) live at the
    top level.  Everything else was stored as a top-level key by
    :class:`WayFilterHandler`.
    """
    return way.get(attr, "")


def write_network_parquet(
    ways: List[dict],
    node_coords: dict,
    nodes_path: str,
    edges_path: str,
    force_bidirectional: bool = True,
    chunk_size: int = 1_000_000,
    edge_attributes: Optional[Tuple[str, ...]] = None,
    node_attrs: Optional[dict] = None,
) -> Tuple[int, int]:
    """Stream nodes and edges to Parquet files in fixed-size chunks.

    This avoids materialising full GeoDataFrames in RAM.  The output
    Parquet files use a ``geometry`` column containing WKB bytes so that
    :meth:`Network._set_nodes_from_parquet` /
    :meth:`Network._set_edges_from_parquet` can ingest them directly.

    Args:
        ways: List of way dicts from :class:`WayFilterHandler`.
        node_coords: ``{node_id: (lon, lat)}`` from
            :class:`NodeCoordinateHandler`.
        nodes_path: Destination path for the nodes Parquet file.
        edges_path: Destination path for the edges Parquet file.
        force_bidirectional: Passed through to the directionality logic.
        chunk_size: Rows per Parquet row-group / flush cycle.
        edge_attributes: Tuple of attribute names to write as edge
            columns.  Defaults to :data:`DEFAULT_EDGE_ATTRIBUTES`.
        node_attrs: Optional ``{node_id: {tag: value}}`` mapping from
            :class:`NodeCoordinateHandler`.  When provided, each tag
            becomes a separate column in the nodes Parquet file.

    Returns:
        ``(total_nodes, total_edges)`` counts.
    """
    import numpy as np
    import pyarrow as pa
    import pyarrow.parquet as pq
    import shapely
    from tqdm import tqdm

    if edge_attributes is None:
        edge_attributes = DEFAULT_EDGE_ATTRIBUTES

    # ---- Nodes ----
    node_fields = [
        ("node_id", pa.int64()),
        ("geometry", pa.binary()),
    ]
    has_node_attrs = bool(node_attrs)
    node_attr_keys: list[str] = []
    if has_node_attrs:
        all_keys: set[str] = set()
        for attrs in node_attrs.values():
            all_keys.update(attrs.keys())
        node_attr_keys = sorted(all_keys)
        for key in node_attr_keys:
            node_fields.append((key, pa.string()))
    nodes_schema = pa.schema(node_fields)

    node_items = list(node_coords.items())
    total_nodes = len(node_items)
    logger.debug("Streaming %d nodes to %s ...", total_nodes, nodes_path)

    writer = pq.ParquetWriter(nodes_path, nodes_schema)
    try:
        for start in range(0, total_nodes, chunk_size):
            chunk = node_items[start : start + chunk_size]
            ids = [nid for nid, _ in chunk]
            coords = np.array([c for _, c in chunk])
            wkb = shapely.to_wkb(shapely.points(coords))
            col_data: dict = {
                "node_id": pa.array(ids, type=pa.int64()),
                "geometry": pa.array(
                    [bytes(w) for w in wkb], type=pa.binary()
                ),
            }
            if has_node_attrs:
                for key in node_attr_keys:
                    col_data[key] = pa.array(
                        [str(node_attrs.get(nid, {}).get(key, "")) for nid in ids],
                        type=pa.string(),
                    )
            batch = pa.table(col_data, schema=nodes_schema)
            writer.write_table(batch)
    finally:
        writer.close()

    # ---- Edges: dynamic attribute columns ----
    edge_schema_fields = [
        ("edge_id", pa.int64()),
        ("source_node", pa.int64()),
        ("target_node", pa.int64()),
        ("geometry", pa.binary()),
    ]
    for attr in edge_attributes:
        pa_type = pa.int64() if attr == "way_id" else pa.string()
        edge_schema_fields.append((attr, pa_type))
    edges_schema = pa.schema(edge_schema_fields)

    logger.debug("Streaming edges to %s ...", edges_path)
    writer = pq.ParquetWriter(edges_path, edges_schema)

    edge_id_counter = 0
    buf_source: list = []
    buf_target: list = []
    buf_src_coords: list = []
    buf_dst_coords: list = []
    attr_bufs: dict[str, list] = {a: [] for a in edge_attributes}
    skipped = 0

    def _flush() -> None:
        nonlocal edge_id_counter
        if not buf_source:
            return
        n = len(buf_source)
        src_arr = np.array(buf_src_coords)
        dst_arr = np.array(buf_dst_coords)
        line_coords = np.stack([src_arr, dst_arr], axis=1)
        wkb = shapely.to_wkb(shapely.linestrings(line_coords))

        ids = list(range(edge_id_counter, edge_id_counter + n))
        edge_id_counter += n

        col_data: dict = {
            "edge_id": pa.array(ids, type=pa.int64()),
            "source_node": pa.array(buf_source, type=pa.int64()),
            "target_node": pa.array(buf_target, type=pa.int64()),
            "geometry": pa.array(
                [bytes(w) for w in wkb], type=pa.binary()
            ),
        }
        for attr in edge_attributes:
            pa_type = pa.int64() if attr == "way_id" else pa.string()
            col_data[attr] = pa.array(attr_bufs[attr], type=pa_type)

        batch = pa.table(col_data, schema=edges_schema)
        writer.write_table(batch)

        buf_source.clear()
        buf_target.clear()
        buf_src_coords.clear()
        buf_dst_coords.clear()
        for buf in attr_bufs.values():
            buf.clear()

    try:
        for way in tqdm(ways, desc="Building edges", unit="way"):
            way_nodes = way["nodes"]
            oneway = way.get("oneway", "no")

            add_forward, add_reverse = _resolve_direction(oneway, force_bidirectional)

            attr_vals = {a: _get_way_attr_value(way, a) for a in edge_attributes}

            for i in range(len(way_nodes) - 1):
                u_id = way_nodes[i]
                v_id = way_nodes[i + 1]

                if u_id not in node_coords or v_id not in node_coords:
                    skipped += 1
                    continue

                u_coord = node_coords[u_id]
                v_coord = node_coords[v_id]

                if add_forward:
                    buf_source.append(u_id)
                    buf_target.append(v_id)
                    buf_src_coords.append(u_coord)
                    buf_dst_coords.append(v_coord)
                    for a in edge_attributes:
                        attr_bufs[a].append(attr_vals[a])

                if add_reverse:
                    buf_source.append(v_id)
                    buf_target.append(u_id)
                    buf_src_coords.append(v_coord)
                    buf_dst_coords.append(u_coord)
                    for a in edge_attributes:
                        attr_bufs[a].append(attr_vals[a])

                if len(buf_source) >= chunk_size:
                    _flush()

        _flush()
    finally:
        writer.close()

    total_edges = edge_id_counter

    if skipped:
        logger.debug(
            "Skipped %d edge segments due to missing node coordinates",
            skipped,
        )

    logger.debug(
        "Wrote %d nodes and %d edges to Parquet.", total_nodes, total_edges
    )
    return total_nodes, total_edges
