"""
Valhalla Local Matcher.

An integrated single-pass map matcher using the pyvalhalla routing engine.
Transforms raw GPS coordinates into snapped, topological network paths and 
interpolates dynamic temporal/spatial attributes onto the graph edges.
"""

import logging
import uuid
from typing import Any, List, Tuple, Optional, Union, TYPE_CHECKING
from scipy.spatial import KDTree
import geopandas as gpd
import numpy as np
import pandas as pd
import networkx as nx

from ..base import BaseMatcher
from ._tiles import ensure_tiles

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

logger = logging.getLogger(__name__)

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

def patch_network_holes(result_df: pd.DataFrame, edges_df: pd.DataFrame) -> pd.DataFrame:
    """
    Fills NaN holes in a route DataFrame using the shortest path from a network of valid edges.
    """
    # 1. Build the graph using the ultra-fast zip method
    G = nx.DiGraph()
    edges_to_add = zip(
        edges_df['source_node'], 
        edges_df['target_node'], 
        edges_df['edge_id'], 
        edges_df.get('way_id', [None] * len(edges_df))
    )
    
    for src, tgt, eid, wid in edges_to_add:
        G.add_edge(src, tgt, edge_id=eid, way_id=wid)
        
    # 2. Identify contiguous blocks of valid rows and NaN rows
    is_nan = result_df.isna().any(axis=1)
    block_ids = (is_nan != is_nan.shift()).cumsum()
    blocks = [group for _, group in result_df.groupby(block_ids)]
    
    patched_chunks = []
    
    # 3. Iterate through the blocks and patch the holes
    for i, block in enumerate(blocks):
        # If it's a valid block (no NaNs), keep it as is
        if not is_nan.loc[block.index[0]]:
            patched_chunks.append(block)
            continue
            
        # Handle edge cases where a hole is at the very beginning or end
        if i == 0 or i == len(blocks) - 1:
            logger.debug("Hole at the bounds of data. Cannot patch. Leaving as NaN.")
            patched_chunks.append(block)
            continue
            
        # Get the adjacent valid blocks to find our start and end points
        prev_block = blocks[i - 1]
        next_block = blocks[i + 1]
        
        start_node = prev_block.iloc[-1]['target_node']
        end_node = next_block.iloc[0]['source_node']
        
        try:
            path_nodes = nx.shortest_path(G, source=start_node, target=end_node)
            path_edges = []
            
            for j in range(len(path_nodes) - 1):
                u = path_nodes[j]
                v = path_nodes[j+1]
                edge_data = G[u][v]
                
                path_edges.append({
                    'source_node': u,
                    'target_node': v,
                    'edge_id': edge_data['edge_id'],
                    'way_id': edge_data['way_id']
                })
                
            patched_chunks.append(pd.DataFrame(path_edges))
            
        except nx.NetworkXNoPath:
            logger.debug(f"No path found between {start_node} and {end_node}. Leaving NaN block.")
            patched_chunks.append(block)
        except nx.NodeNotFound as e:
            logger.debug(f"Node missing in network graph - {e}. Leaving NaN block.")
            patched_chunks.append(block)

    return pd.concat(patched_chunks, ignore_index=True).astype('Int64')


def decode_valhalla(encoded: str) -> List[Tuple[float, float]]:
    """Decode Valhalla's polyline6 encoding into coordinate tuples."""
    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 spatial_enrich(
    snapped_path_gdf: gpd.GeoDataFrame, 
    raw_gps_gdf: gpd.GeoDataFrame, 
    columns: List[str]
) -> gpd.GeoDataFrame:
    """Enrich snapped nodes with nearest attributes from the raw GPS trajectory.

    Uses ``np.maximum.accumulate`` on nearest-neighbour indices to enforce
    forward-only (monotonic) matching along the path.  This prevents a
    later snapped point from "reaching back" to an earlier GPS observation.
    The trade-off is that if the raw GPS trace doubles back on itself, the
    matched attributes may skip observations rather than assign the
    geographically closest one.
    """
    if not columns:
        return snapped_path_gdf
        
    nearest = raw_gps_gdf.sindex.nearest(snapped_path_gdf.geometry, return_all=False)
    if nearest is not None and len(nearest[0]) > 0:
        sort_idx = np.argsort(nearest[0])
        nearest_matched_indices = nearest[1][sort_idx]
        
        # Force forward-progression matching equivalent to Dynamic Time Warping
        nearest_matched_indices = np.maximum.accumulate(nearest_matched_indices)
        
        for attr in columns:
            if attr in raw_gps_gdf.columns:
                snapped_path_gdf[attr] = raw_gps_gdf[attr].values[nearest_matched_indices]
                
    return snapped_path_gdf


def interpolate_along_path(
    snapped_path_gdf: gpd.GeoDataFrame, 
    raw_gps_gdf: gpd.GeoDataFrame, 
    columns: List[str],
    keep_original: bool = False
) -> gpd.GeoDataFrame:
    """Interpolate numeric or datetime columns along the topological path distance.

    Distances are computed by reprojecting to EPSG:3857 (Web Mercator).
    This is accurate enough for moderate latitudes, but distortion grows
    significantly above ~60 degrees N/S.  For polar regions an equal-area
    or local UTM projection would be more appropriate.
    """
    if not columns:
        return snapped_path_gdf
        
    # Reproject to metric bounds for distance calcs
    raw_meters = raw_gps_gdf.to_crs("EPSG:3857")
    path_meters = snapped_path_gdf.to_crs("EPSG:3857")
    
    gps_cum_dist = raw_meters.geometry.distance(raw_meters.geometry.shift()).fillna(0).cumsum()
    path_cum_dist = path_meters.geometry.distance(path_meters.geometry.shift()).fillna(0).cumsum()

    gps_total = gps_cum_dist.iloc[-1] if gps_cum_dist.iloc[-1] > 0 else 1
    path_total = path_cum_dist.iloc[-1] if path_cum_dist.iloc[-1] > 0 else 1
    
    gps_progression = gps_cum_dist / gps_total
    path_progression = path_cum_dist / path_total
    
    for attr in columns:
        if attr not in raw_gps_gdf.columns:
            continue
            
        target_col = f"{attr}_interp" if keep_original else attr
        
        # Handle datetime columns (convert to seconds, interpolate, and back)
        if pd.api.types.is_datetime64_any_dtype(raw_gps_gdf[attr]):
            start_time = raw_gps_gdf[attr].iloc[0]
            elapsed_seconds = (raw_gps_gdf[attr] - start_time).dt.total_seconds()
            interp_elapsed = np.interp(path_progression, gps_progression, elapsed_seconds)
            snapped_path_gdf[target_col] = start_time + pd.to_timedelta(np.round(interp_elapsed), unit='s')
            
        # Handle numeric columns directly
        elif pd.api.types.is_numeric_dtype(raw_gps_gdf[attr]):
            vals = raw_gps_gdf[attr].values
            interp_vals = np.interp(path_progression, gps_progression, vals)
            snapped_path_gdf[target_col] = interp_vals
            
        else:
            logger.warning(f"Skipping interpolation for non-numeric/datetime column: {attr}")

    return snapped_path_gdf


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

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

    def __init__(
        self,
        network: 'Network',
        osm_pbf: str,
        costing: str = "pedestrian",
        enrich_columns: Union[str, List[str], None] = "all",
        interpolate_columns: Union[str, List[str], None] = None,
        keep_interpolated_original: bool = False,
        **kwargs: Any,
    ):
        """Initialise the Valhalla matcher.

        Args:
            network: Built walkingpandas Network.
            osm_pbf: Path to local OSM PBF file for Valhalla tiles.
            costing: Routing model (default ``'pedestrian'``).
            enrich_columns: Columns to snap to the route via nearest-neighbour
                matching.  Use ``'all'`` for dynamic resolution of every
                non-geometry column, or ``None`` to skip enrichment entirely.
            interpolate_columns: Numeric/datetime columns to interpolate
                evenly along the path distance.  Pass ``'default'`` for the
                built-in preset ``['timestamp']``, an explicit list, or
                ``None`` to skip interpolation.
            keep_interpolated_original: If ``True``, interpolation adds a new
                ``<col>_interp`` column instead of overwriting the original.
        """
        self._network = network
        self.osm_pbf = osm_pbf
        self.costing = costing
        
        self.enrich_columns = enrich_columns
        self.interpolate_columns = ["timestamp"] if interpolate_columns == "default" else interpolate_columns
        self.keep_interpolated_original = keep_interpolated_original
        self.kwargs = kwargs

        self.actor = None
        self._ensure_actor()
        
    @property
    def network(self) -> 'Network':
        """Return the underlying walkingpandas Network."""
        return self._network

    def _ensure_actor(self) -> None:
        """Build (or locate) the tile extract and create a Valhalla Actor."""
        self.tile_extract = ensure_tiles(osm_pbf=self.osm_pbf)
        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 match_trajectory(
        self, trajectory: pd.DataFrame, walk_id: Optional[str] = None
    ) -> pd.DataFrame:
        """Process a single trajectory into a validated WalkFrame format."""

        if trajectory.empty or "lat" not in trajectory.columns or "lon" not in trajectory.columns:
            return pd.DataFrame()

        if walk_id is None:
            if "walk_id" in trajectory.columns:
                walk_id = str(trajectory["walk_id"].iloc[0])
            else:
                walk_id = uuid.uuid4().hex[:12]
        locations = trajectory[["lat", "lon"]].to_dict("records")

        query = {
            "shape": locations,
            "costing": self.costing,
            "shape_match": "map_snap",
            "filters": {
                "attributes": ["shape", "matched_points", "matched.point", "edge.way_id"],
                "action": "include",
            },
        }
        
        try:
            response = self.actor.trace_attributes(query)
        except Exception as e:
            logger.debug(f"Valhalla engine failed to trace trajectory '{walk_id}': {e}")
            return pd.DataFrame()
            
        encoded_shape = response.get("shape")
        if not encoded_shape:
            return pd.DataFrame()
            
        decoded_shape_coords = decode_valhalla(encoded_shape)

        # 1. Fetch Candidate Topology from DuckDB
        way_ids = [edge.get("way_id") for edge in response.get("edges", []) if edge.get("way_id") is not None]
        if not way_ids:
            return pd.DataFrame()

        conn = self._network.connection
        conn.execute(
            "CREATE OR REPLACE TEMP TABLE _tmp_way_ids AS SELECT UNNEST(?) AS way_id",
            [list(set(way_ids))],
        )
        candidate_edges_df = conn.execute("""
            SELECT e.source_node, e.target_node, e.edge_id, e.way_id
            FROM wp_edges e SEMI JOIN _tmp_way_ids t ON e.way_id = t.way_id
        """).df()
        if candidate_edges_df.empty:
            return pd.DataFrame()

        node_ids = list(
            set(candidate_edges_df['source_node']).union(set(candidate_edges_df['target_node']))
        )
        conn.execute(
            "CREATE OR REPLACE TEMP TABLE _tmp_node_ids AS SELECT UNNEST(?) AS node_id",
            [node_ids],
        )
        candidate_nodes_df = conn.execute("""
            SELECT n.node_id, ST_X(n.geom) AS lon, ST_Y(n.geom) AS lat
            FROM wp_nodes n SEMI JOIN _tmp_node_ids t ON n.node_id = t.node_id
        """).df()

        # 2. Node Spatial Resolution (KDTree on Candidates)
        node_mapping = {
            (round(lon, 6), round(lat, 6)): nid 
            for nid, lon, lat in zip(candidate_nodes_df["node_id"], candidate_nodes_df["lon"], candidate_nodes_df["lat"])
        }

        node_keys = list(node_mapping.keys())
        tree = KDTree(node_keys)
        resolved_node_ids = []
        
        # 1.05e-6 degrees is approx 0.1 meters (strict exact-match tolerance)
        for idx, coord in enumerate(decoded_shape_coords):
            node_id = node_mapping.get(coord)
            if node_id is not None:
                resolved_node_ids.append(node_id)
                continue
                
            dist, match_idx = tree.query(coord, distance_upper_bound=1.05e-6)
            if dist <= 1.05e-6:
                resolved_node_ids.append(node_mapping[node_keys[match_idx]])
            else:
                # Silently ignore the coordinate drop unless it's a critical inner node
                if idx not in (0, len(decoded_shape_coords) - 1):
                    logger.debug(f"Coordinate desync at index {idx} in trace {walk_id}: {coord}")

        if not resolved_node_ids:
            return pd.DataFrame()

        # 3. Reconstruct the Edges 
        edges_list = list(zip(resolved_node_ids, resolved_node_ids[1:]))
        order_df = pd.DataFrame(edges_list, columns=['source_node', 'target_node'])
        result_df = order_df.merge(candidate_edges_df, on=['source_node', 'target_node'], how='left')

        # 4. Patch Topology Holes
        if result_df['edge_id'].isna().any():
            result_df = patch_network_holes(result_df, candidate_edges_df)

        if result_df.empty:
            return pd.DataFrame()

        final_destination = result_df['target_node'].iloc[-1]
        path_df = result_df[['source_node', 'edge_id']].rename(columns={
            'source_node': 'node_id',
            'edge_id': 'exit_edge'
        }).astype('Int64')

        final_row = pd.DataFrame({
            'node_id': [final_destination],
            'exit_edge': [pd.NA]
        }, dtype='Int64')

        formatted_path = pd.concat([path_df, final_row], ignore_index=True)
        
        # 5. Determine Reference Geometry for Enrichment/Interpolation
        matched_pts_df = pd.DataFrame(response.get("matched_points", []))
        if matched_pts_df.empty or len(matched_pts_df) != len(trajectory):
            # Fallback to raw coords if valhalla pruned points unexpectedly
            raw_gps_gdf = gpd.GeoDataFrame(
                trajectory,
                geometry=gpd.points_from_xy(trajectory['lon'], trajectory['lat']),
                crs="EPSG:4326"
            )
        else:
            raw_gps_gdf = gpd.GeoDataFrame(
                trajectory,
                geometry=gpd.points_from_xy(matched_pts_df['lon'], matched_pts_df['lat']),
                crs="EPSG:4326"
            )

        merged_path = formatted_path.merge(candidate_nodes_df, on='node_id', how='left')
        snapped_path_gdf = gpd.GeoDataFrame(
            merged_path,
            geometry=gpd.points_from_xy(merged_path['lon'], merged_path['lat']),
            crs="EPSG:4326"
        ).drop(columns=['lon', 'lat'])

        # 6. Atomic Enrichment
        if self.enrich_columns:
            if self.enrich_columns == "all":
                skip_cols = ['lat', 'lon', 'geometry'] + (self.interpolate_columns or [])
                cols_to_enrich = [c for c in trajectory.columns if c not in skip_cols]
            else:
                cols_to_enrich = self.enrich_columns
                
            snapped_path_gdf = spatial_enrich(snapped_path_gdf, raw_gps_gdf, cols_to_enrich)

        # 7. Atomic Time/Numeric Interpolation
        if self.interpolate_columns:
            snapped_path_gdf = interpolate_along_path(
                snapped_path_gdf, 
                raw_gps_gdf, 
                self.interpolate_columns, 
                self.keep_interpolated_original
            )

        # Final Formatting
        core_cols = ['walk_id', 'sequence_idx', 'node_id', 'exit_edge']
        
        final_df = (
            pd.DataFrame(snapped_path_gdf)
            .drop(columns=['geometry'])
            .assign(
                walk_id=walk_id,
                sequence_idx=lambda df: range(len(df))
            )
        )
        
        # Ensure correct column order, allowing dynamic passthrough attributes
        final_cols = core_cols + [c for c in final_df.columns if c not in core_cols]
        return final_df[final_cols]