"""WalkFrame class for lazy-evaluated walk/path data analysis.

Provides a pandas-like API for querying massive-scale traversal data
stored in Parquet files, with lazy evaluation and DuckDB backend.
"""

import copy
import tempfile
import duckdb
import pandas as pd
import geopandas as gpd
from pathlib import Path
from typing import TYPE_CHECKING, Optional, Union, Tuple, List, Dict, Any

if TYPE_CHECKING:
    from .network import Network
    from shapely.geometry import Point, Polygon, LineString
from datetime import datetime, time
import warnings
from .._utils import _escape_sql_string, _format_id_list_sql
from ..graph.validate import only_simple_paths, only_cycles, only_complex_walks
from ..spatial.filter import passing_through, within_bounds, intersecting, near_point
from ..temporal.edge import (
    edge_frequencies, node_frequencies, walk_lengths,
    edge_traffic_at, filter_time_range
)


class WalkFrame:
    """Lazy-evaluated wrapper around Parquet files or a DuckDB table containing walk/path data.

    WalkFrame provides a pandas-like API for querying massive-scale
    traversal data. Queries are built up lazily and only executed
    when ``.compute()`` is called.

    Walks can be loaded from Parquet file(s) or from a table in the same
    DuckDB as the network (e.g. for single-file workflows). You can provide
    either an existing Network or nodes/edges directly; the underlying
    Network is available as the ``.network`` property for saving or sharing.

    Args:
        parquet_path: Path pattern to Parquet files (supports glob patterns).
            Omit or ``None`` when using *walks_table*.
        network: Network instance providing the underlying topology.
            Required if *nodes* is not provided.
        view_name: Name for the DuckDB view (default: ``'wp_walks'``).
        walks_table: Name of the table in the Network's DuckDB that
            contains walk rows (superset schema). Exactly one of
            *parquet_path* or *walks_table* must be provided.
        nodes: Nodes data source. If provided (and *network* is not),
            a Network is built from *nodes* and *edges*.
        edges: Edges data source. Optional when *nodes* is used;
            only nodes are loaded if omitted.

    Examples:
        >>> network = Network(nodes="nodes.parquet", edges="edges.parquet")
        >>> walks = WalkFrame("data/walks/*.parquet", network=network)
        >>> result = walks.filter(time_range=('08:00', '09:00')).compute()

        Using nodes and edges directly (network built internally):

        >>> walks = WalkFrame("data/walks/*.parquet", nodes=nodes_gdf, edges=edges_gdf)
        >>> walks.network.get_nodes()  # access the underlying network

        Single-file DuckDB (network + walks table in one file):

        >>> network = Network.load("data.duckdb")
        >>> walks = WalkFrame(network=network, walks_table="wp_walks")

    Ways to create a WalkFrame:

    - **Parquet path(s) for walks + network or nodes** --
      ``WalkFrame.from_parquet(walks_path, network=...)`` or ``read_dataset(...)``.
    - **In-memory DataFrame of walks** --
      ``WalkFrame.from_dataframe(walks_df, network=...)`` or ``nodes=.../edges=...``.
    - **Only Parquet walks (infer network)** --
      ``WalkFrame.from_walks(walks_path, infer_network=True)``.
    - **Single DuckDB file (network + walks table)** --
      ``connect(duckdb_path, walks_table=...)`` or ``WalkFrame.from_duckdb(duckdb_path)``.
    - **Network in DuckDB + walks in Parquet** --
      ``connect(duckdb_path, walks_path=walks_path)``.

    Note: ``read_duckdb`` is deprecated; use ``connect()`` or ``WalkFrame.from_duckdb()``.
    """
    
    # Expected schema columns
    REQUIRED_COLUMNS = ['walk_id', 'sequence_idx', 'node_id']
    OPTIONAL_COLUMNS = ['exit_edge', 'timestamp', 'dwell_time', 'attributes']
    
    def __init__(
        self,
        parquet_path: Optional[Union[str, Path]] = None,
        network: Optional['Network'] = None,
        view_name: str = 'wp_walks',
        *,
        walks_table: Optional[str] = None,
        nodes: Optional[Union[str, Path, gpd.GeoDataFrame, pd.DataFrame]] = None,
        edges: Optional[Union[str, Path, gpd.GeoDataFrame, pd.DataFrame]] = None,
    ):
        """Initialize WalkFrame.

        Args:
            parquet_path: Path to Parquet file(s) — supports glob patterns.
                Required if *walks_table* is not set.
            network: Network instance. Required if *nodes* is not provided.
            view_name: Name for DuckDB view.
            walks_table: Name of table in the Network's DuckDB with walk data.
                Required if *parquet_path* is not set.
            nodes: Nodes data source. If provided (and *network* is not),
                a Network is built from *nodes* and *edges*.
            edges: Edges data source. Optional when *nodes* is used.
        """
        has_path = parquet_path is not None and (not isinstance(parquet_path, str) or parquet_path.strip() != '')
        if not has_path and not walks_table:
            raise ValueError("Provide either parquet_path or walks_table.")
        if has_path and walks_table is not None:
            raise ValueError("Provide exactly one of parquet_path or walks_table.")
        if network is not None:
            if nodes is not None or edges is not None:
                raise ValueError("Provide either network= or (nodes=, edges=), not both.")
            self._network = network
        elif nodes is not None:
            from .network import Network
            self._network = Network(nodes=nodes, edges=edges)
        else:
            raise ValueError("Provide either network= or nodes= (and optionally edges=).")
        self._parquet_path = str(parquet_path) if has_path else None
        self._walks_table = walks_table
        self._view_name = view_name
        self._conn = self._network.connection
        
        # Query state for lazy evaluation
        self._query_state = {
            'filters': [],
            'validations': [],
            'aggregations': [],
            'select_cols': None,
            'group_by': None,
            'order_by': None,
            'limit': None,
            'custom_sql': None,  # For complex queries that override normal SQL building
        }
        
        # Create view over Parquet files or DuckDB table
        self._create_view()
        
        # Validate schema
        self._validate_schema()

    @property
    def network(self) -> 'Network':
        """The underlying Network (for saving, sharing, or extending)."""
        return self._network

    @classmethod
    def from_walks(
        cls,
        parquet_path: Union[str, Path],
        *,
        network: Optional['Network'] = None,
        infer_network: bool = False,
        database: Optional[Union[str, Path]] = None,
        **kwargs: Any,
    ) -> 'WalkFrame':
        """Create a WalkFrame over Parquet walk data, optionally inferring the network.

        Args:
            parquet_path: Path (or glob) to Parquet file(s) with walk data.
            network: If provided, use this network (same as
                ``WalkFrame(parquet_path, network=network, **kwargs)``).
            infer_network: If ``True``, infer nodes and edges from the walk
                data via ``Network.from_walks(parquet_path, database=database)``,
                then build the WalkFrame with that network.
            database: Passed to ``Network.from_walks`` when *infer_network*
                is ``True``.
            **kwargs: Additional arguments for the WalkFrame constructor
                (e.g. *view_name*).

        Returns:
            WalkFrame: Lazy-evaluated walk frame over the given data.
        """
        if network is not None:
            return cls(parquet_path, network=network, **kwargs)
        if infer_network:
            from .network import Network
            net = Network.from_walks(parquet_path, database=database)
            return cls(parquet_path, network=net, **kwargs)
        raise ValueError("Provide either network= or infer_network=True.")

    @classmethod
    def from_parquet(
        cls,
        walks_path: Union[str, Path],
        *,
        network: Optional['Network'] = None,
        nodes: Optional[Union[str, Path, gpd.GeoDataFrame, pd.DataFrame]] = None,
        edges: Optional[Union[str, Path, gpd.GeoDataFrame, pd.DataFrame]] = None,
        view_name: str = 'wp_walks',
        **kwargs: Any,
    ) -> 'WalkFrame':
        """Create a WalkFrame from Parquet file(s) containing walk data.

        Preferred way to create a WalkFrame when you have a Parquet path and
        either an existing Network or nodes/edges to build the network from.

        Args:
            walks_path: Path (or glob) to Parquet file(s) with walk data.
            network: Existing Network instance. Provide this or *nodes*
                (and optionally *edges*).
            nodes: Nodes data source. If provided (and *network* is not),
                a Network is built from *nodes* and *edges*.
            edges: Edges data source. Optional when *nodes* is used.
            view_name: Name for the DuckDB view (default: ``'wp_walks'``).
            **kwargs: Additional arguments passed to the WalkFrame
                constructor.

        Returns:
            WalkFrame: Lazy-evaluated walk frame over the given Parquet data.
        """
        return cls(
            walks_path,
            network=network,
            nodes=nodes,
            edges=edges,
            view_name=view_name,
            **kwargs
        )

    @classmethod
    def from_duckdb(
        cls,
        duckdb_path: Union[str, Path],
        *,
        walks_table: str = 'wp_walks',
        **kwargs: Any,
    ) -> 'WalkFrame':
        """Create a WalkFrame from a single DuckDB file containing the network and walks table.

        Loads the network from the DuckDB file and builds a WalkFrame over
        the given table. For "network in DuckDB + walks in Parquet" use the
        top-level ``connect()`` function instead.

        Args:
            duckdb_path: Path to the ``.duckdb`` file (must contain a
                network from ``Network.save()`` and a walks table).
            walks_table: Name of the table containing walk rows
                (default: ``'wp_walks'``).
            **kwargs: Additional arguments passed to DuckDB when loading
                the network.

        Returns:
            WalkFrame: Lazy-evaluated walk frame backed by the DuckDB file.
        """
        from .network import Network
        network = Network.load(duckdb_path, **kwargs)
        return cls(network=network, walks_table=walks_table)

    @classmethod
    def from_dataframe(
        cls,
        walks_df: pd.DataFrame,
        *,
        network: Optional['Network'] = None,
        nodes: Optional[Union[str, Path, gpd.GeoDataFrame, pd.DataFrame]] = None,
        edges: Optional[Union[str, Path, gpd.GeoDataFrame, pd.DataFrame]] = None,
        **kwargs: Any,
    ) -> 'WalkFrame':
        """Create a WalkFrame from an in-memory pandas DataFrame of walk data.

        The DataFrame is converted to superset format if needed (e.g.
        ``id`` -> ``walk_id``, ``step`` -> ``sequence_idx``,
        ``node`` -> ``node_id``), validated, written to a temporary Parquet
        file, and the returned WalkFrame uses that file. Temporary files
        are not deleted so the WalkFrame remains valid for its lifetime.

        Args:
            walks_df: In-memory DataFrame of walk data (will be converted
                to superset schema if needed).
            network: Existing Network instance. Provide this or *nodes*
                (and optionally *edges*).
            nodes: Nodes data source. If provided (and *network* is not),
                a Network is built from *nodes* and *edges*.
            edges: Edges data source. Optional when *nodes* is used.
            **kwargs: Additional arguments passed to the WalkFrame
                constructor.

        Returns:
            WalkFrame: Lazy-evaluated walk frame over the temporary
            Parquet file.
        """
        from ..io.ingest import convert_to_superset_format, validate_schema, write_walks

        if network is not None and (nodes is not None or edges is not None):
            raise ValueError("Provide either network= or (nodes=, edges=), not both.")
        if network is None and nodes is None:
            raise ValueError("Provide either network= or nodes= (and optionally edges=).")

        df = convert_to_superset_format(walks_df)
        validate_schema(df, strict=True)

        resolved_net = network
        if resolved_net is None and nodes is not None:
            from .network import Network as _Net
            resolved_net = _Net(nodes=nodes, edges=edges)

        if resolved_net is not None and resolved_net.has_string_mapping:
            nmap = resolved_net._node_original_to_int
            emap = resolved_net._edge_original_to_int
            if 'node_id' in df.columns and pd.api.types.is_string_dtype(df['node_id']):
                df['node_id'] = df['node_id'].astype(str).map(nmap).astype(pd.Int64Dtype())
            if 'exit_edge' in df.columns and pd.api.types.is_string_dtype(df['exit_edge']):
                df['exit_edge'] = df['exit_edge'].astype(str).map(emap).astype(pd.Int64Dtype())

        tmp_dir = tempfile.mkdtemp(prefix='walkingpandas_')
        temp_path = Path(tmp_dir) / 'walks.parquet'
        write_walks(df, temp_path, schema_validate=False)

        return cls(
            str(temp_path),
            network=resolved_net,
            **kwargs
        )

    def _create_view(self):
        """Create DuckDB view over Parquet files or an existing table."""
        if self._walks_table is not None:
            # If table name equals view name, use the table directly (no view needed)
            if self._walks_table == self._view_name:
                return
            try:
                self._conn.execute(f"DROP VIEW IF EXISTS {self._view_name}")
            except Exception:
                pass
            quoted = '"' + self._walks_table.replace('"', '""') + '"'
            self._conn.execute(f"""
                CREATE VIEW {self._view_name} AS
                SELECT * FROM {quoted}
            """)
        else:
            try:
                self._conn.execute(f"DROP VIEW IF EXISTS {self._view_name}")
            except Exception:
                pass
            # Create view with hive partitioning support
            try:
                self._conn.execute(f"""
                    CREATE VIEW {self._view_name} AS 
                    SELECT * FROM read_parquet('{self._parquet_path}', hive_partitioning=1)
                """)
            except Exception as e:
                try:
                    self._conn.execute(f"""
                        CREATE VIEW {self._view_name} AS 
                        SELECT * FROM read_parquet('{self._parquet_path}')
                    """)
                except Exception as e2:
                    raise ValueError(f"Failed to create view over Parquet files: {e2}")
    
    def _validate_schema(self):
        """Validate that Parquet files have required schema columns."""
        try:
            # Get schema from view
            schema_query = f"DESCRIBE {self._view_name}"
            schema_df = self._conn.execute(schema_query).df()
            available_cols = set(schema_df['column_name'].str.lower())
            
            # Check required columns
            missing = [col for col in self.REQUIRED_COLUMNS if col.lower() not in available_cols]
            if missing:
                raise ValueError(
                    f"Missing required columns in walk data schema: {missing}. "
                    f"Available columns: {list(available_cols)}"
                )
        except Exception as e:
            warnings.warn(f"Schema validation failed: {e}")
    
    def _build_sql(self) -> str:
        """Build SQL query from accumulated query state."""
        # Check for custom SQL (complex queries that override normal building)
        if self._query_state.get('custom_sql'):
            return self._query_state['custom_sql']
        
        # Start with base SELECT
        if self._query_state['select_cols']:
            select_clause = f"SELECT {self._query_state['select_cols']}"
        else:
            select_clause = f"SELECT *"
        
        sql_parts = [f"{select_clause} FROM {self._view_name}"]
        
        # Add WHERE clauses
        # Handle subqueries in filters (they may contain newlines)
        filters = []
        for f in self._query_state['filters']:
            if '\n' in f or 'SELECT' in f.upper():
                # This is a subquery, wrap it properly
                filters.append(f"({f.strip()})")
            else:
                filters.append(f)
        
        if filters:
            where_clause = " AND ".join(filters)
            sql_parts.append(f"WHERE {where_clause}")
        
        # Add GROUP BY
        if self._query_state['group_by']:
            sql_parts.append(f"GROUP BY {self._query_state['group_by']}")
        
        # Add ORDER BY
        if self._query_state['order_by']:
            sql_parts.append(f"ORDER BY {self._query_state['order_by']}")
        
        # Add LIMIT
        if self._query_state['limit']:
            sql_parts.append(f"LIMIT {self._query_state['limit']}")
        
        return " ".join(sql_parts)
    
    def compute(self, decode: bool = True) -> pd.DataFrame:
        """Execute the accumulated query and return result as DataFrame.

        Args:
            decode: When the network uses a string translation layer,
                join the result with the mapping tables so the returned
                DataFrame shows the original string identifiers instead
                of integers.

        Returns:
            DataFrame: Query result as pandas DataFrame.
        """
        sql = self._build_sql()
        result = self._conn.execute(sql).df()

        if decode and self._network.has_string_mapping:
            result = self._decode_result(result)

        self._reset_query_state()
        return result

    @staticmethod
    def _make_decoder(mapping: dict):
        """Return a function that maps integer keys to original labels."""
        def _decode(v):
            if pd.isna(v):
                return v
            try:
                return mapping.get(int(v), v)
            except (ValueError, TypeError):
                return v
        return _decode

    def _decode_result(self, df: pd.DataFrame) -> pd.DataFrame:
        """Replace integer IDs with original string labels using in-memory maps."""
        net = self._network
        df = df.copy()
        nmap = net._node_int_to_original
        emap = net._edge_int_to_original
        if 'node_id' in df.columns and nmap:
            df['node_id'] = df['node_id'].map(self._make_decoder(nmap))
        if 'exit_edge' in df.columns and emap:
            df['exit_edge'] = df['exit_edge'].map(self._make_decoder(emap))
        if 'edge_id' in df.columns and emap:
            df['edge_id'] = df['edge_id'].map(self._make_decoder(emap))
        return df

    def to_pandas(self) -> pd.DataFrame:
        """Alias for ``compute()``."""
        return self.compute()
    
    def to_geopandas(self) -> gpd.GeoDataFrame:
        """Execute query and return as GeoDataFrame with geometries.

        Joins with the Network to get geometries for nodes/edges.

        Returns:
            GeoDataFrame: Query result with geometries.
        """
        # For now, return regular DataFrame
        # Full geometry joining can be added later
        df = self.compute()
        return gpd.GeoDataFrame(df)
    
    def _reset_query_state(self):
        """Reset query state to initial state."""
        self._query_state = {
            'filters': [],
            'validations': [],
            'aggregations': [],
            'select_cols': None,
            'group_by': None,
            'order_by': None,
            'limit': None,
            'custom_sql': None,
        }

    def _branch(self) -> 'WalkFrame':
        """Return a new WalkFrame sharing data sources but with independent query state."""
        new_self = copy.copy(self)
        new_self._query_state = copy.deepcopy(self._query_state)
        return new_self

    def filter(self, time_range: Optional[Tuple] = None, **kwargs: Any) -> 'WalkFrame':
        """Add filter conditions to the query.

        Args:
            time_range: ``(start_time, end_time)`` tuple for temporal
                filtering.
            **kwargs: Additional column-based filters.

        Returns:
            WalkFrame: New WalkFrame with filters applied (for method
            chaining).
        """
        new_self = self._branch()
        if time_range:
            filter_time_range(new_self, time_range)
            return new_self
        for col, value in kwargs.items():
            resolved = self._resolve_filter_value(col, value)
            if isinstance(resolved, list):
                values_sql = _format_id_list_sql(resolved)
                new_self._query_state['filters'].append(f"{col} IN ({values_sql})")
            else:
                if isinstance(resolved, str):
                    value_str = f"'{_escape_sql_string(resolved)}'"
                elif isinstance(resolved, (int, float)) and not isinstance(resolved, bool):
                    value_str = str(int(resolved))
                else:
                    value_str = str(resolved)
                new_self._query_state['filters'].append(f"{col} = {value_str}")
        return new_self

    def _resolve_filter_value(self, col: str, value):
        """Resolve string IDs through the network mapping for ID columns."""
        net = self._network
        has_map = getattr(net, 'has_string_mapping', False)
        if not has_map:
            return value

        def _resolve_one(v, col_name):
            if isinstance(v, str):
                if col_name in ('node_id',):
                    return net._node_original_to_int.get(v, v)
                if col_name in ('exit_edge', 'edge_id'):
                    return net._edge_original_to_int.get(v, v)
            return v

        if isinstance(value, (list, tuple)):
            return [_resolve_one(v, col) for v in value]
        return _resolve_one(value, col)
    
    def edge_frequencies(self) -> 'WalkFrame':
        """Count distinct walks per edge.

        Returns:
            WalkFrame: New WalkFrame with the aggregation applied.
        """
        new_self = self._branch()
        edge_frequencies(new_self)
        return new_self

    def node_frequencies(self) -> 'WalkFrame':
        """Count distinct walks per node.

        Returns:
            WalkFrame: New WalkFrame with the aggregation applied.
        """
        new_self = self._branch()
        node_frequencies(new_self)
        return new_self

    def walk_lengths(self) -> 'WalkFrame':
        """Calculate length (number of steps) for each walk.

        Returns:
            WalkFrame: New WalkFrame with the aggregation applied.
        """
        new_self = self._branch()
        walk_lengths(new_self)
        return new_self

    def edge_traffic_at(self, time: Union[datetime, time, str]) -> 'WalkFrame':
        """Get traffic on edges at a specific time.

        Args:
            time: Time to query traffic at.

        Returns:
            WalkFrame: New WalkFrame with the traffic filter applied.
        """
        new_self = self._branch()
        edge_traffic_at(new_self, time)
        return new_self

    def limit(self, n: int) -> 'WalkFrame':
        """Limit the number of results.

        Args:
            n: Maximum number of rows to return.

        Returns:
            WalkFrame: New WalkFrame with the limit applied.
        """
        new_self = self._branch()
        new_self._query_state['limit'] = n
        return new_self

    def only_simple_paths(self) -> 'WalkFrame':
        """Filter to only simple paths (no node visited twice).

        Returns:
            WalkFrame: New WalkFrame with the filter applied.
        """
        new_self = self._branch()
        only_simple_paths(new_self)
        return new_self

    def only_cycles(self) -> 'WalkFrame':
        """Filter to only cycles (walks where first node equals last node).

        Returns:
            WalkFrame: New WalkFrame with the filter applied.
        """
        new_self = self._branch()
        only_cycles(new_self)
        return new_self

    def only_complex_walks(self) -> 'WalkFrame':
        """Filter to only complex walks (walks that revisit at least one node).

        Returns:
            WalkFrame: New WalkFrame with the filter applied.
        """
        new_self = self._branch()
        only_complex_walks(new_self)
        return new_self

    def only_walks(self) -> 'WalkFrame':
        """Deprecated. Use ``only_complex_walks()`` instead."""
        warnings.warn(
            "only_walks() is deprecated, use only_complex_walks().",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.only_complex_walks()

    def passing_through(self, nodes: Union[int, str, List]) -> 'WalkFrame':
        """Filter walks that pass through specified node(s).

        Args:
            nodes: Node ID(s) to filter by (strings resolved via
                translation table).

        Returns:
            WalkFrame: New WalkFrame with the filter applied.
        """
        new_self = self._branch()
        passing_through(new_self, nodes)
        return new_self

    def within_bounds(self, bbox: Tuple[float, float, float, float]) -> 'WalkFrame':
        """Filter walks within a bounding box.

        Args:
            bbox: ``(minx, miny, maxx, maxy)`` bounding box.

        Returns:
            WalkFrame: New WalkFrame with the spatial filter applied.
        """
        new_self = self._branch()
        within_bounds(new_self, bbox)
        return new_self

    def intersecting(self, geometry: Union["Polygon", "LineString", "Point"]) -> 'WalkFrame':
        """Filter walks intersecting a geometry.

        Args:
            geometry: Shapely geometry (Polygon, LineString, or Point)
                to intersect with.

        Returns:
            WalkFrame: New WalkFrame with the spatial filter applied.
        """
        new_self = self._branch()
        intersecting(new_self, geometry)
        return new_self

    def near_point(self, point: "Point", distance: float) -> 'WalkFrame':
        """Filter walks near a point within specified distance.

        Args:
            point: Shapely Point geometry.
            distance: Distance threshold.

        Returns:
            WalkFrame: New WalkFrame with the spatial filter applied.
        """
        new_self = self._branch()
        near_point(new_self, point, distance)
        return new_self
    
    def get_orphaned_records(self) -> pd.DataFrame:
        """Find records with node_id or exit_edge that don't exist in Network.

        Returns:
            DataFrame: Orphaned records.
        """
        # Check for orphaned node_ids
        orphan_nodes_sql = f"""
            SELECT w.*
            FROM {self._view_name} w
            LEFT JOIN wp_nodes n ON w.node_id = n.node_id
            WHERE n.node_id IS NULL
        """
        
        # Check for orphaned exit_edges
        orphan_edges_sql = f"""
            SELECT w.*
            FROM {self._view_name} w
            LEFT JOIN wp_edges e ON w.exit_edge = e.edge_id
            WHERE w.exit_edge IS NOT NULL AND e.edge_id IS NULL
        """
        
        orphan_nodes = self._conn.execute(orphan_nodes_sql).df()
        orphan_edges = self._conn.execute(orphan_edges_sql).df()
        
        # Combine and deduplicate
        if not orphan_nodes.empty and not orphan_edges.empty:
            return pd.concat([orphan_nodes, orphan_edges]).drop_duplicates()
        elif not orphan_nodes.empty:
            return orphan_nodes
        elif not orphan_edges.empty:
            return orphan_edges
        else:
            return pd.DataFrame()
