"""
Valhalla Local Matcher.

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

import logging
import math
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 LineString, Point

from ..base import BaseMatcher
from ._tiles import ensure_tiles

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

logger = logging.getLogger(__name__)

_MAX_DFS_DEPTH = 50

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


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

    Args:
        encoded: The encoded polyline string from a Valhalla response.

    Returns:
        List of ``(longitude, latitude)`` coordinate pairs.
    """
    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 is_close(
    p1: Tuple[float, ...],
    p2: Tuple[float, ...],
    tol: float = 2e-6,
) -> bool:
    """Check coordinate proximity within *tol* degrees (~0.2 m).

    Args:
        p1: First coordinate tuple.
        p2: Second coordinate tuple.
        tol: Maximum Euclidean distance in degrees to consider "close".

    Returns:
        bool: ``True`` if the points are within *tol* of each other.
    """
    return math.dist(p1[:2], p2[:2]) <= tol


def _find_edge_path(
    chunk: Tuple[str, Tuple[float, float], Tuple[float, float]],
    gdf: gpd.GeoDataFrame,
) -> Optional[List[Tuple[int, int, int]]]:
    """DFS pathfinding over graph edges matching a single Valhalla chunk.

    Pieces together edges whose ``way_id`` / ``osmid`` equals the chunk's
    target way, returning the topological traversal as a list of
    ``(u_node, v_node, edge_id)`` tuples (all ints), or ``None`` if no valid
    path exists.

    Args:
        chunk: ``(way_id, start_coord, end_coord)`` describing the Valhalla
            edge segment to reconstruct.
        gdf: Edges GeoDataFrame with a ``way_id`` or ``osmid`` column.

    Returns:
        List of ``(u_node, v_node, edge_id)`` tuples representing the
        topological traversal, or ``None`` if no valid path is found.
    """
    target_way_id, start_target, end_target = chunk

    if "way_id" in gdf.columns:
        subset = gdf[gdf["way_id"] == target_way_id]
    elif "osmid" in gdf.columns:
        def _matches_way(row: pd.Series) -> bool:
            v = row["osmid"]
            if isinstance(v, list):
                return str(target_way_id) in (str(x) for x in v)
            if isinstance(v, str):
                return str(target_way_id) in [x.strip() for x in v.split(",")]
            return str(v) == str(target_way_id)

        subset = gdf[gdf.apply(_matches_way, axis=1)]
    else:
        raise KeyError("Edges GeoDataFrame must have 'way_id' or 'osmid' column.")

    edges: List[dict] = []
    for idx, row in subset.iterrows():
        u, v, k = idx
        edge_id = int(row["edge_id"])
        coords = list(row.geometry.coords)
        edges.append(
            {
                "u": int(u),
                "v": int(v),
                "edge_id": edge_id,
                "start": coords[0],
                "end": coords[-1],
                "used": False,
            }
        )

    def _dfs(
        current_coord: Tuple[float, float],
        path: List[Tuple[int, int, int]],
        depth: int,
    ) -> Optional[List[Tuple[int, int, int]]]:
        if depth > _MAX_DFS_DEPTH:
            return None
        if path and is_close(current_coord, end_target):
            return path

        for edge in edges:
            if edge["used"]:
                continue

            if is_close(current_coord, edge["start"]):
                edge["used"] = True
                res = _dfs(
                    edge["end"],
                    path + [(edge["u"], edge["v"], edge["edge_id"])],
                    depth + 1,
                )
                if res:
                    return res
                edge["used"] = False

            elif is_close(current_coord, edge["end"]):
                edge["used"] = True
                res = _dfs(
                    edge["start"],
                    path + [(edge["v"], edge["u"], edge["edge_id"])],
                    depth + 1,
                )
                if res:
                    return res
                edge["used"] = False

        return None

    return _dfs(start_target, [], 0)


# ---------------------------------------------------------------------------
# Vectorised osmid filter
# ---------------------------------------------------------------------------

def _filter_edges_by_way_ids(
    edges_gdf: gpd.GeoDataFrame,
    osm_way_ids: set[str],
) -> gpd.GeoDataFrame:
    """Return the subset of *edges_gdf* whose ``osmid`` matches *osm_way_ids*.

    Uses a fast ``isin`` path for scalar osmid columns and falls back to
    row-wise ``apply`` only when the column contains lists or comma-
    separated strings.

    Args:
        edges_gdf: Edges GeoDataFrame with an ``osmid`` column.
        osm_way_ids: Set of OSM way ID strings to match against.

    Returns:
        GeoDataFrame: Filtered copy containing only matching edges.
    """
    col = edges_gdf["osmid"]
    sample = col.dropna().iloc[:100] if len(col) > 0 else col

    has_complex = sample.apply(lambda v: isinstance(v, (list, str)) and (isinstance(v, list) or "," in v)).any()

    if not has_complex:
        return edges_gdf[col.astype(str).isin(osm_way_ids)].copy()

    def _is_in(osmid: object) -> bool:
        if isinstance(osmid, list):
            return any(str(oid) in osm_way_ids for oid in osmid)
        if isinstance(osmid, str):
            return any(oid.strip() in osm_way_ids for oid in osmid.split(","))
        return str(osmid) in osm_way_ids

    return edges_gdf[col.apply(_is_in)].copy()


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

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

    Args:
        region: ``(north, south, east, west)`` bounding box for on-the-fly
            tile building.
        osm_pbf: Path to an existing ``.osm.pbf`` file.
        osm_xml: Path to an existing ``.osm`` XML file (used to build the
            topological graph via OSMnx).
        costing: Valhalla costing model (default ``"pedestrian"``).
        **kwargs: Extra keys forwarded to every ``trace_attributes`` call.
    """

    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,
    ):
        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
        self._edge_map: Dict[Tuple[int, int], int] = {}

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

    # ------------------------------------------------------------------
    # Setup
    # ------------------------------------------------------------------

    def set_region(self, bbox: Tuple[float, float, float, float]) -> None:
        """Set the geographic region and rebuild tiles / graph.

        Args:
            bbox: ``(north, south, east, west)`` bounding box.
        """
        self.region = bbox
        self._setup()

    def _setup(self) -> None:
        self._ensure_actor()
        self._load_graph()

    def _ensure_actor(self) -> None:
        """Build (or fetch cached) Valhalla tiles and create the Actor."""
        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 OSMnx graph and prepare cached GeoDataFrames."""
        if not self.osm_xml:
            return

        try:
            import osmnx as ox
        except ImportError as e:
            raise ImportError(
                "Install osmnx to use the ValhallaMatcher: "
                "pip install walkingpandas[mapmatch]"
            ) from e

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

        for u, v, k, data in self._nx_graph.edges(keys=True, data=True):
            for attr_key, attr_val in list(data.items()):
                if isinstance(attr_val, list):
                    data[attr_key] = ", ".join(map(str, attr_val))

        for _, data in self._nx_graph.nodes(data=True):
            for attr_key, attr_val in list(data.items()):
                if isinstance(attr_val, list):
                    data[attr_key] = ", ".join(map(str, attr_val))

        self._nodes_gdf, self._edges_gdf = ox.graph_to_gdfs(self._nx_graph)

        self._nodes_gdf.index = self._nodes_gdf.index.astype(np.int64)
        self._nodes_gdf.crs = "EPSG:4326"

        if isinstance(self._edges_gdf.index, pd.MultiIndex):
            new_levels = [
                self._edges_gdf.index.levels[0].astype(np.int64),
                self._edges_gdf.index.levels[1].astype(np.int64),
                self._edges_gdf.index.levels[2],
            ]
            self._edges_gdf.index = self._edges_gdf.index.set_levels(new_levels)

        self._edges_gdf["edge_id"] = np.arange(len(self._edges_gdf), dtype=np.int64)
        self._edges_gdf.crs = "EPSG:4326"

        self._edge_map = {}
        for idx, row in self._edges_gdf.iterrows():
            u, v, _k = idx
            eid = int(row["edge_id"])
            self._edge_map[(int(u), int(v))] = eid
            self._edge_map[(int(v), int(u))] = eid

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

    def get_network(self) -> 'Network':
        """Return the walkingpandas Network that backs this matcher.

        Returns:
            Network: The topological network derived from the loaded OSM
            graph.

        Raises:
            ValueError: If the matcher has not been initialised with OSM
                data.
        """
        if self._wp_network is not None:
            return self._wp_network

        if self._nx_graph is None:
            raise ValueError(
                "ValhallaMatcher not properly initialised with OSM data."
            )

        from ....core.network import Network

        self._wp_network = Network.from_networkx(self._nx_graph)
        return self._wp_network

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

        Args:
            chunk_df: GPS observations with ``trip_id``, ``lat``, ``lon``.

        Returns:
            DataFrame with columns ``trip_id``, ``raw_shape``,
            ``valhalla_edges``, ``matched_points``, ``original_points``.
        """
        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_attributes failed for trip '%s'", trip_id,
                    exc_info=True,
                )
                continue

            if "shape" not in response or "edges" not in response:
                logger.debug(
                    "Incomplete response for trip '%s' (keys: %s)",
                    trip_id,
                    list(response.keys()),
                )
                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:
        """Phase 2 -- topological alignment.

        Prunes overhangs using spatial metric distance and reconstructs the
        discrete node path from Valhalla's edge sequence.

        Args:
            raw_traces: Output of :meth:`trace`.
            network: Optional fallback network (prefers cached GeoDataFrames).

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

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

        nodes_gdf = (
            self._nodes_gdf
            if self._nodes_gdf is not None
            else network.get_nodes(as_geopandas=True)
        )
        edges_gdf = (
            self._edges_gdf
            if self._edges_gdf is not None
            else network.get_edges(as_geopandas=True)
        )

        walk_rows: List[dict] = []

        for _, trace in raw_traces.iterrows():
            decoded_coords = decode_valhalla(trace["raw_shape"])
            path_geom = LineString(decoded_coords)

            osm_way_ids: set[str] = {
                str(edge.get("way_id"))
                for edge in trace["valhalla_edges"]
                if edge.get("way_id")
            }

            candidate_edges = _filter_edges_by_way_ids(edges_gdf, osm_way_ids)
            if candidate_edges.empty:
                continue

            path_gs = gpd.GeoSeries([path_geom], crs="EPSG:4326")
            local_metric_crs = nodes_gdf.estimate_utm_crs()

            u_node_ids = candidate_edges.index.get_level_values("u")
            v_node_ids = candidate_edges.index.get_level_values("v")

            u_points = gpd.GeoSeries(
                nodes_gdf.loc[u_node_ids].geometry.values,
                index=candidate_edges.index,
                crs="EPSG:4326",
            ).to_crs(local_metric_crs)
            v_points = gpd.GeoSeries(
                nodes_gdf.loc[v_node_ids].geometry.values,
                index=candidate_edges.index,
                crs="EPSG:4326",
            ).to_crs(local_metric_crs)
            path_geom_proj = path_gs.to_crs(local_metric_crs).iloc[0]

            distance_threshold_m = 0.2
            valid_mask = u_points.dwithin(
                path_geom_proj, distance_threshold_m
            ) & v_points.dwithin(path_geom_proj, distance_threshold_m)
            valid_edges = candidate_edges[valid_mask].copy()

            if valid_edges.empty:
                continue

            sequence: List[Tuple[str, Tuple, Tuple]] = []
            for edge in trace["valhalla_edges"]:
                way_id = str(edge.get("way_id"))
                if way_id == "None":
                    continue
                start = decoded_coords[edge.get("begin_shape_index")]
                end = decoded_coords[edge.get("end_shape_index")]
                sequence.append((way_id, start, end))

            walk_rows_trip: List[dict] = []
            seq_idx = 0

            for i, chunk in enumerate(sequence):
                path = _find_edge_path(chunk, valid_edges)

                if path:
                    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"],
                                    "sequence_idx": seq_idx,
                                    "node_id": int(traversed_u),
                                    "exit_edge": int(edge_id),
                                }
                            )
                            seq_idx += 1
                        else:
                            walk_rows_trip[-1]["exit_edge"] = int(edge_id)

                        walk_rows_trip.append(
                            {
                                "walk_id": trace["trip_id"],
                                "sequence_idx": seq_idx,
                                "node_id": int(traversed_v),
                                "exit_edge": None,
                            }
                        )
                        seq_idx += 1
                else:
                    if 0 < i < len(sequence) - 1:
                        logger.warning(
                            "Could not connect chunk %d (way %s) between %s and %s",
                            i,
                            chunk[0],
                            chunk[1],
                            chunk[2],
                        )

            walk_rows.extend(walk_rows_trip)

        df = pd.DataFrame(walk_rows)

        if not df.empty and "exit_edge" in df.columns:
            edge_map = self._edge_map
            if not edge_map:
                edge_map = {}
                for eidx, erow in valid_edges.iterrows():
                    eu, ev, _ek = eidx
                    eid = int(erow["edge_id"])
                    edge_map[(int(eu), int(ev))] = eid
                    edge_map[(int(ev), int(eu))] = eid

            for i, r in df[df["exit_edge"].isnull()].iterrows():
                nid = int(r["node_id"])
                nid_next = int(df.iloc[i + 1]["node_id"]) if i < len(df) - 1 else None
                if nid_next is not None:
                    eid = edge_map.get((nid, nid_next))
                    if eid is not None:
                        df.at[i, "exit_edge"] = eid

            df["node_id"] = df["node_id"].astype(pd.Int64Dtype())
            df["exit_edge"] = df["exit_edge"].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:
        """Phase 3 -- attribute projection.

        Maps attributes from original GPS observations onto the topological
        walk structure using matched-point nearest-neighbour assignment that
        strictly advances forward in sequence.

        Args:
            snapped_df: Output of :meth:`snap`.
            raw_traces: Output of :meth:`trace`.
            node_attributes: Attributes to project onto each visited node.
            edge_attributes: Attributes to project onto each traversed edge.

        Returns:
            Enriched walk DataFrame.
        """
        if snapped_df.empty or raw_traces.empty:
            return snapped_df

        trace_lookup = raw_traces.set_index("trip_id").to_dict("index")
        nodes_gdf = self._nodes_gdf

        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"]
            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()
                ]

            orig_gdf = orig_pts.copy().reset_index(drop=True)
            matched_gdf = gpd.GeoDataFrame(
                orig_gdf, geometry=matched_geoms, crs="EPSG:4326"
            )

            enriched_group = group.copy()
            current_matched_idx = 0

            for idx, row in enriched_group.iterrows():
                if current_matched_idx >= len(matched_gdf):
                    break

                nid = int(row["node_id"])
                node_geom = nodes_gdf.loc[nid].geometry

                remaining_matched = matched_gdf.iloc[current_matched_idx:]

                nearest_pos = remaining_matched.sindex.nearest(node_geom)[1][0]
                best_idx = remaining_matched.index[nearest_pos]

                nearest_gps_row = matched_gdf.loc[best_idx]

                if node_attributes:
                    for attr in node_attributes:
                        enriched_group.at[idx, attr] = nearest_gps_row.get(attr)

                if edge_attributes and row["exit_edge"]:
                    for attr in edge_attributes:
                        enriched_group.at[idx, attr] = nearest_gps_row.get(attr)

                current_matched_idx = best_idx + 1

            results.append(enriched_group)

        return pd.concat(results, ignore_index=True)
