"""
Data ingestion and export utilities.

Provides helpers for validating schemas, writing Parquet files,
and converting data to the walkingpandas superset format.
"""

import pandas as pd
import geopandas as gpd
import pyarrow as pa
import pyarrow.parquet as pq
from pathlib import Path
from typing import Optional, Union, List, Dict, Any
import warnings
import json
import duckdb

from ..core.network import Network
from ..core.walk import WalkFrame


# Superset schema definition
REQUIRED_COLUMNS = ['walk_id', 'sequence_idx', 'node_id']
OPTIONAL_COLUMNS = ['exit_edge', 'timestamp', 'dwell_time', 'dynamic_attributes']


def validate_schema(df: pd.DataFrame, strict: bool = False) -> bool:
    """Validate that DataFrame matches walkingpandas superset schema.

    Args:
        df: DataFrame to validate.
        strict: If True, raise error on missing required columns.

    Returns:
        bool: True if schema is valid.

    Raises:
        ValueError: If strict=True and required columns are missing.
    """
    df_cols = set(df.columns.str.lower())
    required_lower = [c.lower() for c in REQUIRED_COLUMNS]
    
    missing = [col for col in required_lower if col not in df_cols]
    
    if missing:
        if strict:
            raise ValueError(f"Missing required columns: {missing}")
        else:
            warnings.warn(f"Missing recommended columns: {missing}")
            return False
    
    return True


def convert_to_superset_format(df: pd.DataFrame) -> pd.DataFrame:
    """Convert user DataFrame to walkingpandas superset format.

    Attempts to map common column names to the superset schema.

    Args:
        df: Input DataFrame.

    Returns:
        DataFrame: DataFrame in superset format.

    Examples:
        >>> df = pd.DataFrame({
        ...     'id': [1, 1, 2],
        ...     'step': [0, 1, 0],
        ...     'node': ['A', 'B', 'C']
        ... })
        >>> converted = convert_to_superset_format(df)
    """
    df = df.copy()
    
    # Column name mapping (common variations)
    column_mapping = {
        # walk_id variations
        'id': 'walk_id',
        'walk_id': 'walk_id',
        'path_id': 'walk_id',
        'trajectory_id': 'walk_id',
        
        # sequence_idx variations
        'step': 'sequence_idx',
        'sequence_idx': 'sequence_idx',
        'sequence': 'sequence_idx',
        'order': 'sequence_idx',
        'index': 'sequence_idx',
        
        # node_id variations
        'node': 'node_id',
        'node_id': 'node_id',
        'vertex': 'node_id',
        
        # exit_edge variations
        'edge': 'exit_edge',
        'exit_edge': 'exit_edge',
        'edge_id': 'exit_edge',
        'next_edge': 'exit_edge',
        
        # timestamp variations
        'time': 'timestamp',
        'timestamp': 'timestamp',
        't': 'timestamp',
        'datetime': 'timestamp',
        
        # dwell_time variations
        'dwell': 'dwell_time',
        'dwell_time': 'dwell_time',
        'wait_time': 'dwell_time',
        'duration': 'dwell_time',
    }
    
    # Rename columns
    df = df.rename(columns=column_mapping)
    
    # Ensure required columns exist
    if 'walk_id' not in df.columns:
        raise ValueError("Could not map 'walk_id' column. Provide 'id', 'walk_id', 'path_id', or 'trajectory_id'")
    if 'sequence_idx' not in df.columns:
        raise ValueError("Could not map 'sequence_idx' column. Provide 'step', 'sequence_idx', 'sequence', or 'order'")
    if 'node_id' not in df.columns:
        raise ValueError("Could not map 'node_id' column. Provide 'node', 'node_id', or 'vertex'")
    
    if not pd.api.types.is_integer_dtype(df['sequence_idx']):
        df['sequence_idx'] = pd.to_numeric(df['sequence_idx'], errors='coerce').astype('Int64')

    if not pd.api.types.is_integer_dtype(df['node_id']):
        try:
            df['node_id'] = pd.to_numeric(df['node_id'], errors='raise').astype('Int64')
        except (ValueError, TypeError):
            pass

    if 'exit_edge' in df.columns and not pd.api.types.is_integer_dtype(df['exit_edge']):
        try:
            df['exit_edge'] = pd.to_numeric(df['exit_edge'], errors='raise').astype('Int64')
        except (ValueError, TypeError):
            pass

    # Handle dynamic_attributes: collect all other columns
    base_cols = REQUIRED_COLUMNS + OPTIONAL_COLUMNS
    other_cols = [c for c in df.columns if c not in base_cols]
    
    if other_cols:
        # Convert to JSON/struct (batch via to_dict('records') faster than apply(axis=1))
        if 'dynamic_attributes' not in df.columns:
            df['dynamic_attributes'] = [
                json.dumps(row) for row in df[other_cols].to_dict('records')
            ]
            # Optionally drop original columns
            # df = df.drop(columns=other_cols)
    
    return df


def write_walks(
    df: pd.DataFrame,
    output_path: Union[str, Path],
    partition_by: Optional[Union[str, List[str]]] = None,
    schema_validate: bool = True
) -> None:
    """Write walks DataFrame to Parquet format with optional partitioning.

    Args:
        df: DataFrame in walkingpandas superset format.
        output_path: Output directory or file path.
        partition_by: Column(s) to partition by (e.g., 'year', 'month', 'region').
        schema_validate: If True, validate schema before writing.

    Examples:
        >>> df = pd.DataFrame({...})
        >>> write_walks(df, "data/walks", partition_by=['year', 'month'])
    """
    output_path = Path(output_path)
    
    # Validate schema
    if schema_validate:
        validate_schema(df, strict=True)
    
    # Convert to superset format if needed
    if not all(col in df.columns for col in REQUIRED_COLUMNS):
        df = convert_to_superset_format(df)
    
    # Ensure output directory exists
    if partition_by:
        output_path.mkdir(parents=True, exist_ok=True)
    else:
        output_path.parent.mkdir(parents=True, exist_ok=True)
    
    # Write Parquet
    if partition_by:
        if isinstance(partition_by, str):
            partition_by = [partition_by]
        
        # Use pyarrow for partitioning
        table = pa.Table.from_pandas(df)
        pq.write_to_dataset(
            table,
            root_path=str(output_path),
            partition_cols=partition_by
        )
    else:
        # Simple write
        df.to_parquet(output_path, index=False)


def write_duckdb(
    network: Network,
    walks_df: pd.DataFrame,
    output_path: Union[str, Path],
    walks_table: str = 'wp_walks',
    schema_validate: bool = True,
) -> None:
    """Write a single DuckDB file containing the network and a walks table.

    Creates a .duckdb file with wp_nodes and wp_edges (from the network) and
    a table of walk rows (superset schema). The file can be read with
    read_duckdb(output_path) or connect(output_path, walks_table=walks_table).
    Overwrites the file if it already exists.

    Args:
        network: Network instance (nodes and edges will be saved to the file).
        walks_df: DataFrame in walkingpandas superset format
            (walk_id, sequence_idx, node_id, ...).
        output_path: Path to the .duckdb file to create.
        walks_table: Name of the table to create for walk data (default: 'wp_walks').
        schema_validate: If True, validate schema and convert to superset format
            if needed (default: True).

    Examples:
        >>> network = wp.Network(nodes=nodes_gdf, edges=edges_gdf)
        >>> write_duckdb(network, walks_df, "data.duckdb")
        >>> walks = wp.read_duckdb("data.duckdb")
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if schema_validate:
        if not all(col in walks_df.columns for col in REQUIRED_COLUMNS):
            walks_df = convert_to_superset_format(walks_df)
        validate_schema(walks_df, strict=True)

    network.save(output_path)

    conn = duckdb.connect(str(output_path.resolve()))
    try:
        quoted_table = '"' + walks_table.replace('"', '""') + '"'
        conn.register('_wp_walks_src', walks_df)
        conn.execute(f'CREATE TABLE {quoted_table} AS SELECT * FROM _wp_walks_src')
    finally:
        conn.close()


def read_dataset(
    nodes: Union[str, Path, gpd.GeoDataFrame],
    edges: Union[str, Path, gpd.GeoDataFrame],
    walks: Union[str, Path],
    **kwargs: Any,
) -> WalkFrame:
    """Convenience wrapper to create Network and WalkFrame in one call.

    Equivalent to building a Network from nodes/edges and then
    ``WalkFrame.from_parquet(walks, network=network)``.

    Args:
        nodes: Nodes data source.
        edges: Edges data source.
        walks: Path to Parquet file(s) with walks data.
        **kwargs: Additional arguments passed to Network.

    Returns:
        WalkFrame: Initialized WalkFrame instance.

    Examples:
        >>> walks = read_dataset(
        ...     nodes="data/network/nodes.parquet",
        ...     edges="data/network/edges.parquet",
        ...     walks="data/walks/*.parquet"
        ... )
    """
    network = Network(nodes=nodes, edges=edges, **kwargs)
    return WalkFrame(walks, network=network)


def read_duckdb(
    duckdb_path: Union[str, Path],
    walks_table: str = 'wp_walks',
    **kwargs: Any,
) -> WalkFrame:
    """Load network and walks from a single DuckDB file.

    .. deprecated::
        Use :func:`connect` or :meth:`WalkFrame.from_duckdb` instead.
        ``read_duckdb(path, walks_table=...)`` is equivalent to
        ``connect(path, walks_table=...)`` or ``WalkFrame.from_duckdb(path, walks_table=...)``.

    The file must contain the network (wp_nodes, wp_edges, e.g. from
    Network.save()) and a table with walk data in superset schema
    (walk_id, sequence_idx, node_id, and optional exit_edge, timestamp, etc.).

    Args:
        duckdb_path: Path to the DuckDB database file (.duckdb).
        walks_table: Name of the table containing walk rows (default: 'wp_walks').
        **kwargs: Additional arguments passed to DuckDB when loading the network.

    Returns:
        WalkFrame: WalkFrame reading from the given table in the same database.

    Examples:
        >>> walks = connect("data.duckdb", walks_table="wp_walks")
        >>> walks = WalkFrame.from_duckdb("data.duckdb", walks_table="trajectories")
    """
    warnings.warn(
        "read_duckdb is deprecated; use connect() or WalkFrame.from_duckdb() instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    return connect(duckdb_path, walks_table=walks_table, **kwargs)


def connect(
    duckdb_path: Union[str, Path],
    walks_path: Optional[Union[str, Path]] = None,
    walks_table: Optional[str] = None,
    **kwargs: Any,
) -> WalkFrame:
    """Load Network from DuckDB database file and create WalkFrame.

    This function loads a Network that was previously saved using Network.save().
    Provide either walks_path (Parquet file(s)) or walks_table (table name in
    the same .duckdb) to create the WalkFrame.

    For single-file DuckDB (network + walks table in one file), you can use
    ``connect(path, walks_table='wp_walks')`` or ``WalkFrame.from_duckdb(path)``.

    Args:
        duckdb_path: Path to DuckDB database file (.duckdb) containing Network data.
        walks_path: Path to Parquet file(s) with walks data.
        walks_table: Name of table in the same DuckDB file with walk data
            (superset schema). Use for single-file workflows.
        **kwargs: Additional arguments passed to DuckDB when loading the network.

    Returns:
        WalkFrame: WalkFrame instance.

    Examples:
        >>> # Parquet walks
        >>> walks = connect("my_network.duckdb", walks_path="data/walks/*.parquet")
        >>> # Single-file: walks table in same .duckdb
        >>> walks = connect("my_network.duckdb", walks_table="wp_walks")
    """
    network = Network.load(duckdb_path, **kwargs)
    
    if walks_path is not None and walks_table is not None:
        raise ValueError("Provide exactly one of walks_path or walks_table.")
    if walks_path is None and walks_table is None:
        raise ValueError(
            "Provide either walks_path (Parquet file(s)) or walks_table (table name in the .duckdb). "
            "Or use Network.load() and create WalkFrame separately."
        )
    if walks_table is not None:
        return WalkFrame(network=network, walks_table=walks_table)
    return WalkFrame(walks_path, network=network)
