"""
Valhalla Local Matcher.

Implements the three-phase :class:`BaseMatcher` interface
(trace / snap / enrich) using the pyvalhalla routing engine with
topological node-snapping and nearest-neighbour attribute projection.

Performance characteristics:
    - Per-way BFS pathfinding restricted to small subgraphs.
    - Nearest-node lookup via brute-force on tiny candidate sets.
    - Bulk spatial join in enrich via ``sindex.nearest``.
"""

import json
import logging
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING

import geopandas as gpd
import numpy as np
import pandas as pd
from shapely.geometry import Point
from shapely.strtree import STRtree

from ..base import BaseMatcher
from ._tiles import ensure_tiles

if TYPE_CHECKING:
    from ....core.network import Network

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Pure helpers
# ---------------------------------------------------------------------------

def decode_valhalla(encoded: str) -> List[Tuple[float, float]]:
    """Decode Valhalla's polyline6 encoding into coordinate tuples.

    Valhalla uses a precision-6 variant of Google's encoded polyline
    format.  Each coordinate pair is delta-encoded relative to the
    previous one.

    Args:
        encoded: Polyline6-encoded string from a Valhalla response.

    Returns:
        List of ``(lon, lat)`` tuples in decimal degrees.
    """
    inv = 1.0 / 1e6
    decoded: List[Tuple[float, float]] = []
    previous = [0, 0]
    i = 0
    while i < len(encoded):
        ll = [0, 0]
        for j in (0, 1):
            shift = 0
            byte = 0x20
            while i < len(encoded) and byte >= 0x20:
                byte = ord(encoded[i]) - 63
                i += 1
                ll[j] |= (byte & 0x1F) << shift
                shift += 5
            ll[j] = previous[j] + (~(ll[j] >> 1) if ll[j] & 1 else (ll[j] >> 1))
            previous[j] = ll[j]
        decoded.append((float(f"{ll[1] * inv:.6f}"), float(f"{ll[0] * inv:.6f}")))
    return decoded


def _bfs_path(
    way_graph: Dict[Any, Dict[Any, int]], start_node: Any, end_node: Any
) -> Optional[List[Tuple[Any, Any, int]]]:
    """Find the shortest topological path within a single-way subgraph.

    Performs breadth-first search on *way_graph* (which represents the
    undirected edges of one OSM way) and returns the ordered sequence of
    ``(u, v, edge_id)`` steps from *start_node* to *end_node*.

    Args:
        way_graph: Adjacency dict ``{node: {neighbour: edge_id, ...}, ...}``
            for one OSM way.
        start_node: Source node identifier.
        end_node: Target node identifier.

    Returns:
        List of ``(u, v, edge_id)`` tuples forming the path, an empty
        list when *start_node* equals *end_node*, or ``None`` when no
        path exists.
    """
    if start_node == end_node:
        return []
        
    queue = [(start_node, [])]
    visited = {start_node}
    
    while queue:
        curr, path = queue.pop(0)
        if curr == end_node:
            return path
            
        if curr in way_graph:
            for nxt, eid in way_graph[curr].items():
                if nxt not in visited:
                    visited.add(nxt)
                    queue.append((nxt, path + [(curr, nxt, eid)]))
                    
    return None


def _find_closest_node(
    target_coord: Tuple[float, float],
    nodes_dict: Dict[Any, Tuple[float, float]],
) -> Any:
    """Find the closest node in a small candidate set by squared distance.

    Intended for way-local node sets (typically < 50 nodes) where a
    brute-force scan is faster than building a spatial index.

    Args:
        target_coord: ``(x, y)`` coordinate to snap to.
        nodes_dict: Mapping ``{node_id: (x, y), ...}`` of candidate nodes.

    Returns:
        The *node_id* of the nearest node, or ``None`` if *nodes_dict*
        is empty.
    """
    best_node = None
    best_dist = float('inf')
    tx, ty = target_coord
    for nid, (nx, ny) in nodes_dict.items():
        # Euclidean is perfect here since the candidate set is tiny and local
        dist = (nx - tx)**2 + (ny - ty)**2
        if dist < best_dist:
            best_dist = dist
            best_node = nid
    return best_node

# ---------------------------------------------------------------------------
# ValhallaMatcher
# ---------------------------------------------------------------------------

class ValhallaMatcher(BaseMatcher):
    """Map-matching backend powered by the pyvalhalla routing engine.

    Builds a local Valhalla tile extract from OSM data, then exposes the
    three-phase :class:`BaseMatcher` interface:

    1. **trace** — calls ``trace_attributes`` to obtain raw matched shapes.
    2. **snap** — resolves the raw shapes to topological node/edge
       sequences using per-way BFS and nearest-node lookup.
    3. **enrich** — projects original GPS attributes onto the snapped
       sequence via bulk spatial joins.

    Typical usage::

        matcher = ValhallaMatcher(region=(11.5, 48.1, 11.6, 48.2),
                                  osm_pbf="munich.osm.pbf")
        network = matcher.get_network()
        traces  = matcher.trace(gps_df)
        snapped = matcher.snap(traces, network)
        result  = matcher.enrich(snapped, traces, node_attributes=["timestamp"])
    """

    def __init__(
        self,
        region: Optional[Tuple[float, float, float, float]] = None,
        osm_pbf: Optional[str] = None,
        osm_xml: Optional[str] = None,
        costing: str = "pedestrian",
        **kwargs: Any,
    ):
        """Initialise the Valhalla matcher.

        At least one of *region* or *osm_pbf* (or *osm_xml*) must be
        provided for the matcher to set up tiles and load the graph
        automatically.  Otherwise, call :meth:`set_region` later.

        Args:
            region: Bounding box ``(min_lon, min_lat, max_lon, max_lat)``
                used to download OSM data when *osm_pbf* is not given.
            osm_pbf: Path to a local ``.osm.pbf`` file.
            osm_xml: Path to a local ``.osm`` XML file used to build
                the topological graph.  If omitted, is derived from the
                tile-building step.
            costing: Valhalla costing model (default ``"pedestrian"``).
            **kwargs: Extra parameters forwarded verbatim to every
                ``trace_attributes`` request (e.g. ``search_radius``).
        """
        self.region = region
        self.osm_pbf = osm_pbf
        self.osm_xml = osm_xml
        self.costing = costing
        self.trace_config = kwargs

        self.actor = None
        self._nx_graph = None
        self._wp_network: Optional['Network'] = None
        self._nodes_gdf: Optional[gpd.GeoDataFrame] = None
        self._edges_gdf: Optional[gpd.GeoDataFrame] = None
        
        # High-performance routing structures
        self._edge_map: Dict[Tuple[Any, Any], int] = {}
        self._way_graphs: Dict[str, Dict[Any, Dict[Any, int]]] = {}
        self._way_nodes: Dict[str, Dict[Any, Tuple[float, float]]] = {}

        if self.region or self.osm_pbf or self.osm_xml:
            self._setup()

    def set_region(self, bbox: Tuple[float, float, float, float]) -> None:
        """Set the geographic bounding box and trigger full setup.

        Calling this method replaces any previously configured region,
        rebuilds the Valhalla tile extract, and reloads the topological
        graph.

        Args:
            bbox: ``(min_lon, min_lat, max_lon, max_lat)`` in WGS 84.
        """
        self.region = bbox
        self._setup()

    def _setup(self) -> None:
        """Run the full initialisation pipeline.

        Calls :meth:`_ensure_actor` (tiles + Valhalla ``Actor``) then
        :meth:`_load_graph` (topological structures from OSM XML).
        """
        self._ensure_actor()
        self._load_graph()

    def _ensure_actor(self) -> None:
        """Build (or locate) the tile extract and create a Valhalla ``Actor``.

        Uses :func:`ensure_tiles` to obtain the tile tarball.  If the
        tile build produces an ``.osm`` XML file and none was provided
        explicitly, it is stored on ``self.osm_xml`` for later graph
        loading.

        Raises:
            ImportError: If ``pyvalhalla`` is not installed.
        """
        self.tile_extract, downloaded_xml = ensure_tiles(
            bbox=self.region, osm_pbf=self.osm_pbf
        )
        if not self.osm_xml and downloaded_xml:
            self.osm_xml = str(downloaded_xml)

        try:
            from valhalla import Actor, get_config
        except ImportError as e:
            raise ImportError(
                "Install pyvalhalla to use the ValhallaMatcher: pip install walkingpandas[mapmatch]"
            ) from e

        self.config = get_config(
            tile_dir=str(self.tile_extract.parent),
            tile_extract=str(self.tile_extract),
            verbose=False,
        )
        self.actor = Actor(self.config)

    def _load_graph(self) -> None:
        """Load the OSM graph and build high-performance routing structures.

        Reads the graph from ``self.osm_xml`` via ``osmnx``, converts it
        to a :class:`~walkingpandas.core.network.Network`, and populates:

        - ``_way_graphs`` — ``{way_id: {node: {neighbour: edge_id}}}``
        - ``_way_nodes`` — ``{way_id: {node_id: (x, y)}}``
        - ``_edge_map`` — ``{(u, v): edge_id}`` (both directions)

        Also caches ``_nodes_gdf`` and ``_edges_gdf`` as decoded
        :class:`GeoDataFrame` objects.

        Raises:
            ImportError: If ``osmnx`` is not installed.
        """
        if not self.osm_xml:
            return

        try:
            import osmnx as ox
        except ImportError as e:
            raise ImportError("Install osmnx to use the ValhallaMatcher") from e

        logger.info("Loading topological graph from %s ...", self.osm_xml)
        self._nx_graph = ox.graph_from_xml(self.osm_xml, simplify=False)

        # Standardize OSM attributes for DuckDB serialization
        for _, _, _, data in self._nx_graph.edges(keys=True, data=True):
            for k, v in list(data.items()):
                if isinstance(v, list):
                    data[k] = ", ".join(map(str, v))
        for _, data in self._nx_graph.nodes(data=True):
            for k, v in list(data.items()):
                if isinstance(v, list):
                    data[k] = ", ".join(map(str, v))

        # Instantiate Network immediately so it serves as the ultimate source of truth for IDs
        from ....core.network import Network
        self._wp_network = Network.from_networkx(self._nx_graph)
        
        # Extract network with decoded IDs to align perfectly with OSM
        self._nodes_gdf = self._wp_network.get_nodes(as_geopandas=True, decode=True)
        self._edges_gdf = self._wp_network.get_edges(as_geopandas=True, decode=True)
        self._nodes_gdf = self._nodes_gdf.set_index("node_id", drop=False)

        self._way_graphs = {}
        self._way_nodes = {}
        self._edge_map = {}
        
        for _, row in self._edges_gdf.iterrows():
            u = row["source_node"]
            v = row["target_node"]
            eid = row["edge_id"]
            
            # Global edge map for inter-chunk repair
            self._edge_map[(u, v)] = eid
            self._edge_map[(v, u)] = eid
            
            # Extract way_id(s) from JSON attributes
            attrs = {}
            if pd.notna(row.get("attributes")):
                try:
                    attrs = json.loads(row["attributes"])
                except Exception:
                    pass
                    
            osmids = attrs.get("osmid", "")
            wids = [str(x) for x in osmids] if isinstance(osmids, list) else [x.strip() for x in str(osmids).split(",")] if osmids else []
            
            # Build localized maps
            for wid in wids:
                if wid not in self._way_graphs:
                    self._way_graphs[wid] = {}
                    self._way_nodes[wid] = {}
                    
                if u not in self._way_graphs[wid]:
                    self._way_graphs[wid][u] = {}
                if v not in self._way_graphs[wid]:
                    self._way_graphs[wid][v] = {}
                    
                self._way_graphs[wid][u][v] = eid
                self._way_graphs[wid][v][u] = eid  # Undirected equivalent for pedestrian matching
                
                if u not in self._way_nodes[wid]:
                    geom = self._nodes_gdf.at[u, 'geometry']
                    self._way_nodes[wid][u] = (geom.x, geom.y)
                if v not in self._way_nodes[wid]:
                    geom = self._nodes_gdf.at[v, 'geometry']
                    self._way_nodes[wid][v] = (geom.x, geom.y)

    # ------------------------------------------------------------------
    # BaseMatcher interface
    # ------------------------------------------------------------------

    def get_network(self) -> 'Network':
        """Return the :class:`~walkingpandas.core.network.Network` built from OSM.

        Returns:
            The network instance derived from the loaded OSM XML graph.

        Raises:
            ValueError: If the matcher has not been set up with OSM data.
        """
        if self._wp_network is not None:
            return self._wp_network
        raise ValueError("ValhallaMatcher not properly initialised with OSM data.")

    def trace(self, chunk_df: pd.DataFrame) -> pd.DataFrame:
        """Run Valhalla ``trace_attributes`` for each trip (phase 1).

        Groups *chunk_df* by ``trip_id`` and sends each group's GPS
        coordinates to the Valhalla actor.  Failed traces are logged and
        silently skipped.

        Args:
            chunk_df: DataFrame with at least ``trip_id``, ``lat``, and
                ``lon`` columns.

        Returns:
            DataFrame with one row per successfully traced trip,
            containing ``trip_id``, ``raw_shape`` (polyline6),
            ``valhalla_edges`` (list of edge dicts),
            ``matched_points``, and ``original_points``.

        Raises:
            ValueError: If *chunk_df* is missing any of the required
                columns.
        """
        missing = {"trip_id", "lat", "lon"} - set(chunk_df.columns)
        if missing:
            raise ValueError(f"chunk_df is missing required columns: {missing}")

        raw_results: List[dict] = []
        for trip_id, gps_points in chunk_df.groupby("trip_id", sort=False):
            locations = [{"lat": r["lat"], "lon": r["lon"]} for _, r in gps_points.iterrows()]

            query = {
                "shape": locations,
                "costing": self.costing,
                "shape_match": "map_snap",
                "filters": {
                    "attributes": [
                        "shape", "matched_points", "matched.point", 
                        "matched.type", "edge.way_id", 
                        "edge.begin_shape_index", "edge.end_shape_index"
                    ],
                    "action": "include",
                },
            }
            query.update(self.trace_config)

            try:
                response = self.actor.trace_attributes(query)
            except Exception:
                logger.warning("Valhalla trace failed for trip '%s'", trip_id, exc_info=True)
                continue

            if "shape" not in response or "edges" not in response:
                continue

            raw_results.append({
                "trip_id": trip_id,
                "raw_shape": response["shape"],
                "valhalla_edges": response["edges"],
                "matched_points": response.get("matched_points", []),
                "original_points": gps_points,
            })

        return pd.DataFrame(raw_results)

    def snap(self, raw_traces: pd.DataFrame, network: Optional['Network'] = None) -> pd.DataFrame:
        """Snap raw trace geometries to topological node/edge sequences (phase 2).

        For each Valhalla edge in a trace, identifies the nearest start
        and end nodes within the edge's OSM way, then finds the
        topological path between them via :func:`_bfs_path`.  A repair
        step fills in ``exit_edge`` gaps at inter-chunk boundaries using
        ``_edge_map``.

        Args:
            raw_traces: Output of :meth:`trace`, must contain
                ``trip_id``, ``raw_shape``, and ``valhalla_edges``.
            network: Unused; accepted for :class:`BaseMatcher` interface
                compatibility.

        Returns:
            DataFrame with columns ``walk_id``, ``sequence_idx``,
            ``node_id``, and ``exit_edge``.  Empty if *raw_traces* is
            empty.
        """
        if raw_traces.empty:
            return pd.DataFrame()

        missing = {"trip_id", "raw_shape", "valhalla_edges"} - set(raw_traces.columns)
        if missing:
            raise ValueError(f"raw_traces is missing required columns: {missing}")

        walk_rows: List[dict] = []

        for _, trace in raw_traces.iterrows():
            decoded_coords = decode_valhalla(trace["raw_shape"])
            walk_rows_trip: List[dict] = []

            for edge in trace["valhalla_edges"]:
                way_id = str(edge.get("way_id"))
                if way_id == "None":
                    continue

                start_coord = decoded_coords[edge.get("begin_shape_index")]
                end_coord = decoded_coords[edge.get("end_shape_index")]

                way_nodes = self._way_nodes.get(way_id)
                way_graph = self._way_graphs.get(way_id)
                
                if not way_nodes or not way_graph:
                    continue

                # Localized snapping (solves 6-digit precision mismatches elegantly)
                u_node = _find_closest_node(start_coord, way_nodes)
                v_node = _find_closest_node(end_coord, way_nodes)

                if u_node is None or v_node is None:
                    continue

                path = _bfs_path(way_graph, u_node, v_node)
                
                if path is not None:
                    for traversed_u, traversed_v, edge_id in path:
                        if not walk_rows_trip or walk_rows_trip[-1]["node_id"] != traversed_u:
                            walk_rows_trip.append({
                                "walk_id": trace["trip_id"],
                                "node_id": traversed_u,
                                "exit_edge": edge_id
                            })
                        else:
                            walk_rows_trip[-1]["exit_edge"] = edge_id

                        walk_rows_trip.append({
                            "walk_id": trace["trip_id"],
                            "node_id": traversed_v,
                            "exit_edge": None
                        })

            walk_rows.extend(walk_rows_trip)

        df = pd.DataFrame(walk_rows)
        
        if not df.empty and "exit_edge" in df.columns:
            # Phase 2b: Repair inter-chunk gaps
            # Finds instances where one chunk ended, the next began at a different node, but an edge connects them
            next_nodes = df["node_id"].shift(-1)
            same_walk = df["walk_id"] == df["walk_id"].shift(-1)
            needs_repair = df["exit_edge"].isna() & same_walk & next_nodes.notna()
            
            for i in df[needs_repair].index:
                u = df.at[i, "node_id"]
                v = next_nodes.at[i]
                eid = self._edge_map.get((u, v))
                if eid is not None:
                    df.at[i, "exit_edge"] = eid

            # Enforce clean schema formats and re-index sequential flow
            df["node_id"] = df["node_id"].astype(pd.Int64Dtype())
            df["exit_edge"] = df["exit_edge"].astype(pd.Int64Dtype())
            df["sequence_idx"] = df.groupby("walk_id").cumcount().astype(pd.Int64Dtype())

        return df

    def enrich(
        self,
        snapped_df: pd.DataFrame,
        raw_traces: pd.DataFrame,
        node_attributes: Optional[List[str]] = None,
        edge_attributes: Optional[List[str]] = None,
    ) -> pd.DataFrame:
        """Project original GPS attributes onto snapped nodes (phase 3).

        For each walk, builds a spatial index on the matched/original
        GPS points and finds the nearest GPS point for every snapped
        node using ``sindex.nearest(..., return_all=False)``.  A
        monotonically-increasing index constraint ensures strict forward
        progression (DTW-like alignment).

        Args:
            snapped_df: Output of :meth:`snap` with ``walk_id``,
                ``node_id``, and ``exit_edge`` columns.
            raw_traces: Output of :meth:`trace` (needed for
                ``original_points`` and ``matched_points``).
            node_attributes: Column names from the original GPS data
                to transfer to every snapped node row.
            edge_attributes: Column names from the original GPS data
                to transfer only where ``exit_edge`` is present.

        Returns:
            Copy of *snapped_df* augmented with the requested
            attributes.  Returns *snapped_df* unchanged if either input
            is empty.
        """
        if snapped_df.empty or raw_traces.empty:
            return snapped_df

        trace_lookup = raw_traces.set_index("trip_id").to_dict("index")
        results: List[pd.DataFrame] = []

        for walk_id, group in snapped_df.groupby("walk_id", sort=False):
            trace_data = trace_lookup.get(walk_id)
            if trace_data is None:
                results.append(group)
                continue

            orig_pts = trace_data["original_points"].reset_index(drop=True)
            matched_pts = trace_data.get("matched_points", [])

            if matched_pts and len(matched_pts) == len(orig_pts):
                matched_geoms = [Point(pt["lon"], pt["lat"]) for pt in matched_pts]
            else:
                matched_geoms = [Point(r["lon"], r["lat"]) for _, r in orig_pts.iterrows()]
            
            if not matched_geoms:
                results.append(group)
                continue

            matched_gdf = gpd.GeoDataFrame(geometry=matched_geoms, crs="EPSG:4326")
            node_geoms = self._nodes_gdf.loc[group["node_id"]].geometry.values
            nodes_gs = gpd.GeoSeries(node_geoms, crs="EPSG:4326")
            
            # Use return_all=False to ensure EXACTLY ONE match per node
            # This prevents shape mismatch errors when ties exist
            nearest = matched_gdf.sindex.nearest(nodes_gs, return_all=False)
            
            if nearest is None or len(nearest[0]) == 0:
                results.append(group)
                continue

            # Ensure strict alignment with the node sequence using array sorts
            sort_idx = np.argsort(nearest[0])
            nearest_matched_indices = nearest[1][sort_idx]
            
            # Enforce strict forward progression matching logic (Dynamic Time Warping equivalent)
            nearest_matched_indices = np.maximum.accumulate(nearest_matched_indices)

            enriched_group = group.copy()
            
            if node_attributes:
                for attr in node_attributes:
                    if attr in orig_pts.columns:
                        enriched_group[attr] = orig_pts[attr].values[nearest_matched_indices]

            if edge_attributes:
                for attr in edge_attributes:
                    if attr in orig_pts.columns:
                        mask = enriched_group["exit_edge"].notna()
                        if mask.any():
                            # Safely extract boolean mask array values for index alignment
                            enriched_group.loc[mask, attr] = orig_pts[attr].values[nearest_matched_indices[mask.values]]

            results.append(enriched_group)

        return pd.concat(results, ignore_index=True)