"""
09 — Persistence: Network.save(), Network.load(), database=, connect(), write_duckdb.

This example builds a network, saves it to a temp .duckdb file, loads it back
with Network.load(), creates a Network with database= path from the start,
uses connect(duckdb_path, walks_path=...) to load Network + WalkFrame from
Parquet, and demonstrates single-file write/read with write_duckdb and
connect() or WalkFrame.from_duckdb(). (read_duckdb is deprecated.)
"""

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_simple_network_gdf, make_simple_walks_df, write_to_temp_parquet


def main():
    print("=" * 60)
    print("09 — Persistence")
    print("=" * 60)

    nodes_gdf, edges_gdf = make_simple_network_gdf()
    walks_df = make_simple_walks_df()
    temp_dir = write_to_temp_parquet(nodes_gdf, edges_gdf, walks_df)
    persist_dir = Path(tempfile.mkdtemp())
    db_path = persist_dir / "net.duckdb"
    db_path_2 = persist_dir / "net2.duckdb"
    db_onefile = persist_dir / "onefile.duckdb"

    try:
        # 1. Build network and save to .duckdb
        network = wp.Network(nodes=nodes_gdf, edges=edges_gdf)
        network.save(db_path)
        print("\n1. Network.save(path) — saved to", db_path.name)

        # 2. Load with Network.load()
        loaded = wp.Network.load(db_path)
        nodes = loaded.get_nodes(as_geopandas=False)
        edges = loaded.get_edges(as_geopandas=False)
        print(f"2. Network.load(path) — got {len(nodes)} nodes, {len(edges)} edges")

        # 3. Network(..., database=path) — create with persistent DB from the start
        net_persistent = wp.Network(nodes=nodes_gdf, edges=edges_gdf, database=db_path_2)
        n = len(net_persistent.get_nodes())
        print(f"3. Network(..., database=path) — created persistent DB, {n} nodes")

        # 4. connect(duckdb_path, walks_path=...) — load Network + WalkFrame from Parquet
        walks = wp.connect(db_path, walks_path=temp_dir / "walks.parquet")
        result = walks.limit(2).compute()
        print("\n4. connect(duckdb_path, walks_path=...) then .limit(2).compute()")
        print(result[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # 5. Single-file DuckDB: write_duckdb then load with connect() or from_duckdb
        net_one = wp.Network(nodes=nodes_gdf, edges=edges_gdf)
        wp.write_duckdb(net_one, walks_df, db_onefile)
        walks_one = wp.connect(db_onefile, walks_table="wp_walks")
        result_one = walks_one.limit(2).compute()
        print("\n5. connect(path, walks_table='wp_walks') — one file with network + wp_walks table")
        print(result_one[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))
        # Same via WalkFrame.from_duckdb()
        walks_one2 = wp.WalkFrame.from_duckdb(db_onefile)
        print("   WalkFrame.from_duckdb(path) — same result")

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


if __name__ == "__main__":
    main()
