"""
12 — Dynamic attribute columns: native columns on nodes, edges, and walks.

Any extra columns you add to node, edge, or walk DataFrames are stored as
native DuckDB columns — not packed into JSON.  This means they are fully
queryable via SQL, can be filtered in the Python API, and survive
save/load and Parquet roundtrips.
"""

import shutil
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import walkingpandas as wp

from _data_helpers import (
    make_network_with_attributes_gdf,
    make_walks_with_extra_cols_df,
    write_to_temp_parquet,
)


def main():
    print("=" * 60)
    print("12 — Dynamic attribute columns")
    print("=" * 60)

    nodes_gdf, edges_gdf = make_network_with_attributes_gdf()
    walks_df = make_walks_with_extra_cols_df()
    persist_dir = Path(tempfile.mkdtemp())

    try:
        # --- 1. Create network with attribute-rich GeoDataFrames ---
        print("\n1. Create network with extra attribute columns")
        net = wp.Network.from_gdfs(nodes_gdf, edges_gdf)
        nodes = net.get_nodes(as_geopandas=False)
        edges = net.get_edges(as_geopandas=False)
        print(f"   Node columns: {list(nodes.columns)}")
        print(f"   Edge columns: {list(edges.columns)}")
        print(f"   → 'highway' and 'crossing' are real node columns")
        print(f"   → 'way_id', 'oneway', 'name', 'length' are real edge columns")

        # --- 2. Direct SQL queries on attribute columns ---
        print("\n2. SQL queries on attribute columns via network.connection")
        conn = net.connection
        crossing_nodes = conn.execute(
            "SELECT node_id, highway, crossing FROM wp_nodes "
            "WHERE highway = 'crossing'"
        ).df()
        print("   Nodes where highway='crossing':")
        print("   " + crossing_nodes.to_string(index=False).replace("\n", "\n   "))

        oneway_edges = conn.execute(
            "SELECT edge_id, name, oneway FROM wp_edges WHERE oneway = 'yes'"
        ).df()
        print("\n   Edges where oneway='yes':")
        print("   " + oneway_edges.to_string(index=False).replace("\n", "\n   "))

        # --- 3. Create walks with extra columns ---
        print("\n3. Create walks with extra columns (speed, confidence)")
        temp_dir = write_to_temp_parquet(nodes_gdf, edges_gdf, walks_df)
        wf = wp.WalkFrame(
            temp_dir / "walks.parquet",
            nodes=nodes_gdf,
            edges=edges_gdf,
        )
        result = wf.compute()
        print(f"   Walk columns: {list(result.columns)}")
        print(f"   → 'speed' and 'confidence' are preserved as native columns")

        # --- 4. Filter walks by extra columns ---
        print("\n4. Filter walks by extra columns")
        fast_steps = result[result["speed"] > 1.0]
        print(f"   Steps with speed > 1.0: {len(fast_steps)} rows")
        print("   " + fast_steps[["walk_id", "node_id", "speed", "confidence"]].to_string(index=False).replace("\n", "\n   "))

        # --- 5. save/load roundtrip preserves attribute columns ---
        print("\n5. save/load roundtrip preserves attribute columns")
        db_path = persist_dir / "attr_network.duckdb"
        net.save(db_path)
        loaded = wp.Network.load(db_path)
        loaded_nodes = loaded.get_nodes(as_geopandas=False)
        loaded_edges = loaded.get_edges(as_geopandas=False)
        assert "highway" in loaded_nodes.columns
        assert "crossing" in loaded_nodes.columns
        assert "way_id" in loaded_edges.columns
        assert "name" in loaded_edges.columns
        print(f"   Loaded node columns: {list(loaded_nodes.columns)}")
        print(f"   Loaded edge columns: {list(loaded_edges.columns)}")
        print("   All attribute columns survived DuckDB roundtrip ✓")

        # --- 6. export_parquet roundtrip preserves attribute columns ---
        print("\n6. export_parquet roundtrip preserves attribute columns")
        parquet_dir = persist_dir / "parquet_export"
        net.export_parquet(parquet_dir)
        reimported = wp.Network(
            nodes=parquet_dir / "nodes.parquet",
            edges=parquet_dir / "edges.parquet",
        )
        ri_nodes = reimported.get_nodes(as_geopandas=False)
        ri_edges = reimported.get_edges(as_geopandas=False)
        assert "highway" in ri_nodes.columns
        assert "way_id" in ri_edges.columns
        print(f"   Re-imported node columns: {list(ri_nodes.columns)}")
        print(f"   Re-imported edge columns: {list(ri_edges.columns)}")
        print("   All attribute columns survived Parquet roundtrip ✓")

        print("\nDone.")
    finally:
        shutil.rmtree(persist_dir, ignore_errors=True)
        if "temp_dir" in dir():
            shutil.rmtree(temp_dir, ignore_errors=True)


if __name__ == "__main__":
    main()
