"""
08 — Ingestion and I/O: validate_schema, convert_to_superset_format, write_walks, read_dataset.

This example shows validate_schema (required vs missing columns; strict=True vs False),
convert_to_superset_format (mapping id/step/node or path_id to superset names, optional
dynamic_attributes from extra columns), write_walks (single file and partition_by),
and read_dataset as a one-liner to get a WalkFrame and run a simple query.
"""

import shutil
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import pandas as pd
import walkingpandas as wp
from walkingpandas.io.ingest import (
    validate_schema,
    convert_to_superset_format,
    write_walks,
    read_dataset,
)

from _data_helpers import (
    make_simple_network_gdf,
    make_simple_walks_df,
    make_string_id_network_gdf,
    make_string_id_walks_df,
    write_to_temp_parquet,
)


def main():
    print("=" * 60)
    print("08 — Ingestion and I/O")
    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)

    try:
        # validate_schema: required columns present
        ok = validate_schema(walks_df, strict=True)
        print("\n1. validate_schema(df, strict=True) →", ok)

        # validate_schema: missing columns (strict=False warns)
        bad_df = pd.DataFrame({"walk_id": [1], "sequence_idx": [0]})
        ok_bad = validate_schema(bad_df, strict=False)
        print("2. validate_schema(df with missing 'node_id', strict=False) →", ok_bad, "(warning above)")

        # convert_to_superset_format: alternative column names
        alt_df = pd.DataFrame({
            "id": ["p1", "p1", "p2"],
            "step": [0, 1, 0],
            "node": [0, 1, 1],
        })
        converted = convert_to_superset_format(alt_df)
        print("\n3. convert_to_superset_format(alt names id/step/node)")
        print(converted[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # write_walks: single file
        out_dir = temp_dir / "out"
        write_walks(walks_df, out_dir / "walks.parquet", schema_validate=True)
        print("\n4. write_walks(df, path) — single file written")

        # write_walks: with partition_by (e.g. year/month) — need a column; add dummy
        walks_with_part = walks_df.copy()
        walks_with_part["year"] = 2024
        walks_with_part["month"] = 1
        write_walks(walks_with_part, out_dir / "partitioned", partition_by=["year", "month"])
        print("5. write_walks(..., partition_by=['year','month']) — partitioned dir written")

        # read_dataset: one-liner to get WalkFrame
        walks = read_dataset(
            nodes=temp_dir / "nodes.parquet",
            edges=temp_dir / "edges.parquet",
            walks=temp_dir / "walks.parquet",
        )
        result = walks.limit(3).compute()
        print("\n6. read_dataset(nodes, edges, walks) then .limit(3).compute()")
        print(result[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # --- String IDs: from_dataframe auto-translates to integers ---
        print("\n--- String-ID ingestion ---")
        str_nodes, str_edges = make_string_id_network_gdf()
        str_net = wp.Network()
        str_net.add_nodes(str_nodes)
        str_net.add_edges(str_edges)

        str_walks = make_string_id_walks_df()
        str_wf = wp.WalkFrame.from_dataframe(str_walks, network=str_net)
        result = str_wf.filter(walk_id="walk1").compute(decode=True)
        print("7. WalkFrame.from_dataframe with string IDs → compute(decode=True)")
        print(result[["walk_id", "node_id", "exit_edge"]].to_string(index=False))

        raw = str_wf.filter(walk_id="walk1").compute(decode=False)
        print("\n8. Same query with decode=False → raw integers")
        print(raw[["walk_id", "node_id", "exit_edge"]].to_string(index=False))

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


if __name__ == "__main__":
    main()
