#!/usr/bin/env python3
"""
Convert Shenzhen- or Jinan-style (Yu2023) road network and vehicle trajectory CSVs to walkingpandas Parquet format.

Reads node, edge, and trajectory CSVs from a configurable data directory (outside the project),
converts them to walkingpandas schema with integer IDs, and writes nodes.parquet, edges.parquet,
and walks.parquet (or per-day walks_YYYYMMDD.parquet) under output_dir/<city>/.

Supported cities:
  - shenzhen: node_shenzhen.csv, edge_shenzhen.csv, traj_shenzhen_*.csv (date from filename)
  - jinan:    node_jinan.csv, edge_jinan.csv, traj_jinan.csv (single file; date 2022-10-17)

Expected CSV formats (actual headers from example data):
  - Node:    NodeID, Longitude, Latitude, HasCamera
  - Edge:    Origin, Destination, Class, Geometry, Length
  - Trajectory: VehicleID, TripID, Points, DepartureTime, Duration, Length

Trajectory filenames may include YYYYMMDD for timestamp computation; otherwise use --traj-date (e.g. jinan).
Points column format: "node1-time1_node2-time2_..." with time in seconds from midnight.
"""

from __future__ import annotations

import argparse
from typing import Any
import re
import shutil
from datetime import datetime
from pathlib import Path

import geopandas as gpd
import pandas as pd
import pyarrow.parquet as pq
from shapely.geometry import LineString, Point

# Optional: env defaults
import os

from walkingpandas.io.ingest import validate_schema, write_walks

try:
    from tqdm import tqdm
except ImportError:
    def tqdm(iterable=None, desc=None, total=None, leave=True, **kwargs):
        """No-op progress wrapper when tqdm is not installed. Install with: uv sync --extra scripts"""
        if iterable is None and total is not None:
            return range(total)
        return iterable if iterable is not None else iter([])

# Filename pattern to extract date: traj_*_YYYYMMDD.csv
TRAJ_DATE_PATTERN = re.compile(r"(\d{8})\.csv$", re.IGNORECASE)

# Per-city default file names and trajectory date behavior
CITY_CONFIG = {
    "shenzhen": {
        "nodes": "node_shenzhen.csv",
        "edges": "edge_shenzhen.csv",
        "traj_pattern": "traj_shenzhen_*.csv",
        "default_traj_date": None,  # date from filename
    },
    "jinan": {
        "nodes": "node_jinan.csv",
        "edges": "edge_jinan.csv",
        "traj_pattern": "traj_jinan.csv",  # single file, no date in name
        "default_traj_date": "2022-10-17",  # YYYY-MM-DD
    },
}


def _parse_trajectory_date_from_path(file_path: Path) -> datetime | None:
    """Parse date from trajectory filename (e.g. traj_shenzhen_20201104.csv -> 2020-11-04)."""
    name = file_path.name
    m = TRAJ_DATE_PATTERN.search(name)
    if not m:
        return None
    try:
        return datetime.strptime(m.group(1), "%Y%m%d")
    except ValueError:
        return None


def _parse_geometry_encoded(geom_str: str) -> LineString | None:
    """Parse edge geometry string: 'lon-lat_lon-lat_...' into a Shapely LineString."""
    if not geom_str or pd.isna(geom_str):
        return None
    points = []
    for segment in geom_str.strip().split("_"):
        segment = segment.strip()
        if not segment:
            continue
        parts = segment.split("-")
        if len(parts) >= 2:
            try:
                lon = float(parts[0])
                lat = float(parts[1])
                points.append((lon, lat))
            except (ValueError, TypeError):
                continue
    if len(points) < 2:
        return None
    return LineString(points)


def build_nodes_gdf(
    nodes_path: Path,
    crs: str,
    show_progress: bool = True,
) -> gpd.GeoDataFrame:
    """Read node CSV and build GeoDataFrame with node_id (int), longitude, latitude, has_camera, geometry, CRS."""
    df = pd.read_csv(nodes_path)
    # Map CSV headers to script names
    column_map = {
        "NodeID": "node_id",
        "Longitude": "longitude",
        "Latitude": "latitude",
        "HasCamera": "has_camera",
    }
    df = df.rename(columns=column_map)
    for col in ["node_id"]:
        if col in df.columns:
            df[col] = df[col].astype("Int64")
    df["has_camera"] = df["has_camera"].astype(bool)
    it = zip(df["longitude"], df["latitude"])
    if show_progress:
        it = tqdm(it, total=len(df), desc="Nodes", unit="node")
    geometry = [Point(x, y) for x, y in it]
    gdf = gpd.GeoDataFrame(df, geometry=geometry, crs=crs)
    return gdf


def build_edges_gdf(
    edges_path: Path,
    crs: str,
    show_progress: bool = True,
) -> tuple[gpd.GeoDataFrame, dict[tuple[int, int], int]]:
    """Read edge CSV, build GeoDataFrame with edge_id (int), source_node, target_node, geometry, etc.; return (gdf, (o,d)->edge_id map)."""
    df = pd.read_csv(edges_path)
    column_map = {
        "Origin": "origin_node_id",
        "Destination": "destination_node_id",
        "Class": "road_class",
        "Geometry": "geometry_encoded",
        "Length": "length_m",
    }
    df = df.rename(columns=column_map)
    df["origin_node_id"] = df["origin_node_id"].astype(int)
    df["destination_node_id"] = df["destination_node_id"].astype(int)
    geom_series = df["geometry_encoded"]
    if show_progress:
        geom_series = tqdm(geom_series, total=len(df), desc="Edges", unit="edge")
    df["geometry"] = [_parse_geometry_encoded(s) for s in geom_series]
    df = df.drop(columns=["geometry_encoded"], errors="ignore")
    # Assign unique integer edge_id
    df.insert(0, "edge_id", range(len(df)))
    df["source_node"] = df["origin_node_id"].astype(int)
    df["target_node"] = df["destination_node_id"].astype(int)
    edge_map: dict[tuple[int, int], int] = {}
    it = df.iterrows()
    if show_progress:
        it = tqdm(it, total=len(df), desc="Edge index", unit="edge")
    for _, row in it:
        edge_map[(int(row["source_node"]), int(row["target_node"]))] = int(row["edge_id"])
    # Keep only output columns: edge_id, source_node, target_node, geometry, road_class, length_m
    out_cols = ["edge_id", "source_node", "target_node", "geometry", "road_class", "length_m"]
    df = df[[c for c in out_cols if c in df.columns]]
    gdf = gpd.GeoDataFrame(df, geometry="geometry", crs=crs)
    return gdf, edge_map


def _streaming_merge_parquet(frag_paths: list[Path], output_path: Path) -> None:
    """Merge fragment Parquet files into one using ParquetWriter (no full load)."""
    if not frag_paths:
        return
    schema = pq.ParquetFile(frag_paths[0]).schema_arrow
    with pq.ParquetWriter(output_path, schema) as writer:
        for p in frag_paths:
            writer.write_table(pq.read_table(p))


def _count_csv_rows(path: Path) -> int:
    """Count data rows in a CSV (lines minus header). Fast line-count only."""
    with open(path, "rb") as f:
        return max(0, sum(1 for _ in f) - 1)


def _build_nodes_chunked(
    nodes_path: Path,
    crs: str,
    chunk_size: int,
    temp_dir: Path,
    output_path: Path,
    show_progress: bool,
) -> None:
    """Read node CSV in chunks, write fragment Parquets, streaming merge to output_path."""
    temp_nodes = temp_dir / "nodes"
    temp_nodes.mkdir(parents=True, exist_ok=True)
    column_map = {
        "NodeID": "node_id",
        "Longitude": "longitude",
        "Latitude": "latitude",
        "HasCamera": "has_camera",
    }
    frag_paths: list[Path] = []
    total_nodes = _count_csv_rows(nodes_path)
    chunk_iter = pd.read_csv(nodes_path, chunksize=chunk_size)
    pbar = (
        tqdm(total=total_nodes, desc="Nodes", unit="node")
        if show_progress
        else None
    )
    for i, chunk_df in enumerate(chunk_iter):
        if pbar is not None:
            pbar.update(len(chunk_df))
        df = chunk_df.rename(columns=column_map)
        if "node_id" in df.columns:
            df["node_id"] = df["node_id"].astype("Int64")
        df["has_camera"] = df["has_camera"].astype(bool)
        geometry = [Point(x, y) for x, y in zip(df["longitude"], df["latitude"])]
        gdf = gpd.GeoDataFrame(df, geometry=geometry, crs=crs)
        out_cols = ["node_id", "longitude", "latitude", "has_camera", "geometry"]
        gdf = gdf[[c for c in out_cols if c in gdf.columns]]
        frag_path = temp_nodes / f"frag_{i:06d}.parquet"
        gdf.to_parquet(frag_path, index=False)
        frag_paths.append(frag_path)
    if pbar is not None:
        pbar.close()
    _streaming_merge_parquet(sorted(frag_paths), output_path)
    shutil.rmtree(temp_nodes, ignore_errors=True)


def _build_edges_chunked(
    edges_path: Path,
    crs: str,
    chunk_size: int,
    temp_dir: Path,
    output_path: Path,
    show_progress: bool,
) -> dict[tuple[int, int], int]:
    """Read edge CSV in chunks, build edge_map incrementally, write fragment Parquets, streaming merge. Returns edge_map."""
    temp_edges = temp_dir / "edges"
    temp_edges.mkdir(parents=True, exist_ok=True)
    column_map = {
        "Origin": "origin_node_id",
        "Destination": "destination_node_id",
        "Class": "road_class",
        "Geometry": "geometry_encoded",
        "Length": "length_m",
    }
    out_cols = ["edge_id", "source_node", "target_node", "geometry", "road_class", "length_m"]
    edge_map: dict[tuple[int, int], int] = {}
    edge_id_counter = 0
    frag_paths: list[Path] = []
    total_edges = _count_csv_rows(edges_path)
    chunk_iter = pd.read_csv(edges_path, chunksize=chunk_size)
    pbar = (
        tqdm(total=total_edges, desc="Edges", unit="edge")
        if show_progress
        else None
    )
    for i, chunk_df in enumerate(chunk_iter):
        if pbar is not None:
            pbar.update(len(chunk_df))
        df = chunk_df.rename(columns=column_map)
        df["origin_node_id"] = df["origin_node_id"].astype(int)
        df["destination_node_id"] = df["destination_node_id"].astype(int)
        df["geometry"] = [_parse_geometry_encoded(s) for s in df["geometry_encoded"]]
        df = df.drop(columns=["geometry_encoded"], errors="ignore")
        n = len(df)
        df.insert(0, "edge_id", range(edge_id_counter, edge_id_counter + n))
        edge_id_counter += n
        df["source_node"] = df["origin_node_id"].astype(int)
        df["target_node"] = df["destination_node_id"].astype(int)
        for _, row in df.iterrows():
            edge_map[(int(row["source_node"]), int(row["target_node"]))] = int(row["edge_id"])
        df = df[[c for c in out_cols if c in df.columns]]
        gdf = gpd.GeoDataFrame(df, geometry="geometry", crs=crs)
        frag_path = temp_edges / f"frag_{i:06d}.parquet"
        gdf.to_parquet(frag_path, index=False)
        frag_paths.append(frag_path)
    if pbar is not None:
        pbar.close()
    _streaming_merge_parquet(sorted(frag_paths), output_path)
    shutil.rmtree(temp_edges, ignore_errors=True)
    return edge_map


def _expand_chunk_to_walk_rows(
    chunk_df: pd.DataFrame,
    traj_date: datetime | None,
    edge_map: dict[tuple[int, int], int],
    walk_id_start: int,
    column_map: dict[str, str],
    pbar: Any = None,
) -> tuple[list[dict], int]:
    """Expand a trajectory chunk DataFrame to walk rows. Returns (rows, next_walk_id)."""
    chunk_df = chunk_df.rename(columns=column_map)
    chunk_df["vehicle_id"] = chunk_df["vehicle_id"].astype(int)
    chunk_df["trip_id"] = chunk_df["trip_id"].astype(int)
    base_date = traj_date.date() if traj_date else None
    rows: list[dict] = []
    walk_id_counter = walk_id_start
    for _, row in chunk_df.iterrows():
        if pbar is not None:
            pbar.update(1)
        vehicle_id = int(row["vehicle_id"])
        trip_id = int(row["trip_id"])
        path_sequence = row["path_sequence"]
        if pd.isna(path_sequence) or not str(path_sequence).strip():
            continue
        steps = []
        for segment in str(path_sequence).strip().split("_"):
            segment = segment.strip()
            if not segment:
                continue
            parts = segment.split("-")
            if len(parts) >= 2:
                try:
                    nid = int(parts[0])
                    tval = float(parts[1])
                    steps.append((nid, tval))
                except (ValueError, TypeError):
                    continue
        if not steps:
            continue
        for seq_idx, (node_id, time_value) in enumerate(steps):
            ts = (
                pd.Timestamp(year=base_date.year, month=base_date.month, day=base_date.day)
                + pd.Timedelta(seconds=float(time_value))
                if base_date
                else pd.NaT
            )
            exit_edge = None
            if seq_idx + 1 < len(steps):
                next_node = steps[seq_idx + 1][0]
                exit_edge = edge_map.get((node_id, next_node))
            rows.append({
                "walk_id": walk_id_counter,
                "sequence_idx": seq_idx,
                "node_id": node_id,
                "exit_edge": exit_edge,
                "timestamp": ts,
                "vehicle_id": vehicle_id,
                "trip_id": trip_id,
            })
        walk_id_counter += 1
    return rows, walk_id_counter


def _build_walks_chunked(
    traj_paths: list[Path],
    edge_map: dict[tuple[int, int], int],
    default_traj_date: str | datetime | None,
    chunk_size: int,
    temp_dir: Path,
    output_dir: Path,
    no_merge: bool,
    walks_per_file: bool,
    show_progress: bool,
) -> tuple[int, int]:
    """Read trajectory CSVs in chunks, write frags; streaming merge or leave as directory (--no-merge). Returns (total_steps, total_walks)."""
    default_dt = default_traj_date if isinstance(default_traj_date, datetime) else _parse_default_traj_date(
        str(default_traj_date) if default_traj_date else None
    )
    column_map = {
        "VehicleID": "vehicle_id",
        "TripID": "trip_id",
        "Points": "path_sequence",
        "DepartureTime": "departure_time",
        "Duration": "duration_seconds",
        "Length": "length_m",
    }
    walk_id_counter = 0
    total_steps = 0
    all_frag_paths: list[Path] = []
    per_file_frags: list[tuple[str, list[Path]]] = []
    global_frag_i = 0

    if no_merge:
        writes_dir = output_dir / "walks"
        writes_dir.mkdir(parents=True, exist_ok=True)
    else:
        temp_walks = temp_dir / "walks"
        temp_walks.mkdir(parents=True, exist_ok=True)

    total_paths = sum(_count_csv_rows(p) for p in traj_paths)
    path_progress_bar = (
        tqdm(total=total_paths, desc="Paths", unit="path")
        if show_progress
        else None
    )
    path_iter = tqdm(traj_paths, desc="Trajectory files", unit="file") if show_progress else traj_paths
    for file_idx, path in enumerate(path_iter):
        traj_date = _parse_trajectory_date_from_path(path) or default_dt
        date_str = traj_date.strftime("%Y%m%d") if traj_date else "unknown"
        file_frags: list[Path] = []
        chunk_iter = pd.read_csv(path, chunksize=chunk_size)
        for chunk_i, chunk_df in enumerate(chunk_iter):
            rows, walk_id_counter = _expand_chunk_to_walk_rows(
                chunk_df, traj_date, edge_map, walk_id_counter, column_map, pbar=path_progress_bar
            )
            if not rows:
                continue
            total_steps += len(rows)
            out = pd.DataFrame(rows)
            out["walk_id"] = out["walk_id"].astype("Int64")
            out["sequence_idx"] = out["sequence_idx"].astype("Int64")
            out["node_id"] = out["node_id"].astype("Int64")
            out["vehicle_id"] = out["vehicle_id"].astype("Int64")
            out["trip_id"] = out["trip_id"].astype("Int64")
            out["exit_edge"] = out["exit_edge"].astype("Int64")
            out["timestamp"] = pd.to_datetime(out["timestamp"], utc=False).astype("datetime64[ns]")
            if no_merge:
                frag_path = writes_dir / f"chunk_{global_frag_i:06d}.parquet"
                global_frag_i += 1
            else:
                frag_path = temp_walks / f"f_{file_idx}_{chunk_i:06d}.parquet"
            out.to_parquet(frag_path, index=False)
            all_frag_paths.append(frag_path)
            file_frags.append(frag_path)

        if walks_per_file:
            per_file_frags.append((date_str, file_frags))

    if path_progress_bar is not None:
        path_progress_bar.close()

    if no_merge:
        if all_frag_paths and show_progress:
            print(f"Wrote {len(all_frag_paths)} walk chunk(s) to {output_dir / 'walks'} (use WalkFrame('walks/*.parquet', network=...))")
        return (total_steps, walk_id_counter)

    # Validate schema on first fragment before merge
    if all_frag_paths:
        validate_schema(pd.read_parquet(all_frag_paths[0]), strict=True)
    if walks_per_file:
        for date_str, paths in per_file_frags:
            if not paths:
                continue
            out_path = output_dir / f"walks_{date_str}.parquet"
            _streaming_merge_parquet(sorted(paths), out_path)
            print(f"Wrote {out_path}")
    else:
        _streaming_merge_parquet(sorted(all_frag_paths), output_dir / "walks.parquet")
        print(f"Wrote {output_dir / 'walks.parquet'}")
    shutil.rmtree(temp_dir / "walks", ignore_errors=True)
    return (total_steps, walk_id_counter)


def _parse_default_traj_date(value: str | None):
    """Parse default traj date string YYYY-MM-DD to datetime (date part)."""
    if not value:
        return None
    try:
        return datetime.strptime(value.strip()[:10], "%Y-%m-%d")
    except ValueError:
        return None


def build_walks_df(
    traj_paths: list[Path],
    edge_map: dict[tuple[int, int], int],
    show_progress: bool = True,
    default_traj_date: str | datetime | None = None,
) -> pd.DataFrame:
    """Read trajectory CSV(s), parse Points with date from filename or default_traj_date; build one walks DataFrame."""
    rows = []
    walk_id_counter = 0
    default_dt = default_traj_date if isinstance(default_traj_date, datetime) else _parse_default_traj_date(str(default_traj_date) if default_traj_date else None)

    path_iter = traj_paths
    if show_progress:
        path_iter = tqdm(traj_paths, desc="Trajectory files", unit="file")
    for path in path_iter:
        traj_date = _parse_trajectory_date_from_path(path) or default_dt
        df = pd.read_csv(path)
        column_map = {
            "VehicleID": "vehicle_id",
            "TripID": "trip_id",
            "Points": "path_sequence",
            "DepartureTime": "departure_time",
            "Duration": "duration_seconds",
            "Length": "length_m",
        }
        df = df.rename(columns=column_map)
        df["vehicle_id"] = df["vehicle_id"].astype(int)
        df["trip_id"] = df["trip_id"].astype(int)

        row_iter = df.iterrows()
        if show_progress:
            row_iter = tqdm(row_iter, total=len(df), desc="Paths", unit="path", leave=False)
        for _, row in row_iter:
            vehicle_id = int(row["vehicle_id"])
            trip_id = int(row["trip_id"])
            path_sequence = row["path_sequence"]
            if pd.isna(path_sequence) or not str(path_sequence).strip():
                continue
            # Parse "node1-time1_node2-time2_..."
            steps = []
            for segment in str(path_sequence).strip().split("_"):
                segment = segment.strip()
                if not segment:
                    continue
                parts = segment.split("-")
                if len(parts) >= 2:
                    try:
                        nid = int(parts[0])
                        tval = float(parts[1])
                        steps.append((nid, tval))
                    except (ValueError, TypeError):
                        continue
            if not steps:
                continue

            # Timestamp: combine file date with seconds-from-midnight
            base_date = traj_date.date() if traj_date else None

            for seq_idx, (node_id, time_value) in enumerate(steps):
                if base_date is not None:
                    # time_value is seconds from midnight
                    ts = pd.Timestamp(year=base_date.year, month=base_date.month, day=base_date.day) + pd.Timedelta(seconds=float(time_value))
                else:
                    ts = pd.NaT
                exit_edge = None
                if seq_idx + 1 < len(steps):
                    next_node = steps[seq_idx + 1][0]
                    key = (node_id, next_node)
                    exit_edge = edge_map.get(key)
                rows.append({
                    "walk_id": walk_id_counter,
                    "sequence_idx": seq_idx,
                    "node_id": node_id,
                    "exit_edge": exit_edge,
                    "timestamp": ts,
                    "vehicle_id": vehicle_id,
                    "trip_id": trip_id,
                })
            walk_id_counter += 1

    if not rows:
        return pd.DataFrame(columns=["walk_id", "sequence_idx", "node_id", "exit_edge", "timestamp", "vehicle_id", "trip_id"])
    out = pd.DataFrame(rows)
    out["walk_id"] = out["walk_id"].astype("Int64")
    out["sequence_idx"] = out["sequence_idx"].astype("Int64")
    out["node_id"] = out["node_id"].astype("Int64")
    out["vehicle_id"] = out["vehicle_id"].astype("Int64")
    out["trip_id"] = out["trip_id"].astype("Int64")
    out["exit_edge"] = out["exit_edge"].astype("Int64")  # nullable int
    return out


def build_walks_dfs_per_file(
    traj_paths: list[Path],
    edge_map: dict[tuple[int, int], int],
    show_progress: bool = True,
    default_traj_date: str | datetime | None = None,
) -> list[tuple[str, pd.DataFrame]]:
    """Read trajectory CSV(s) and return one (date_str, walks_df) per file. date_str is YYYYMMDD for writing walks_YYYYMMDD.parquet."""
    default_dt = default_traj_date if isinstance(default_traj_date, datetime) else _parse_default_traj_date(str(default_traj_date) if default_traj_date else None)
    result: list[tuple[str, pd.DataFrame]] = []
    path_iter = tqdm(traj_paths, desc="Trajectory files", unit="file") if show_progress else traj_paths
    for path in path_iter:
        traj_date = _parse_trajectory_date_from_path(path) or default_dt
        date_str = traj_date.strftime("%Y%m%d") if traj_date else "unknown"
        file_df = pd.read_csv(path)
        column_map = {
            "VehicleID": "vehicle_id",
            "TripID": "trip_id",
            "Points": "path_sequence",
            "DepartureTime": "departure_time",
            "Duration": "duration_seconds",
            "Length": "length_m",
        }
        file_df = file_df.rename(columns=column_map)
        file_df["vehicle_id"] = file_df["vehicle_id"].astype(int)
        file_df["trip_id"] = file_df["trip_id"].astype(int)
        rows: list[dict] = []
        walk_id_counter = 0
        base_date = traj_date.date() if traj_date else None
        row_iter = tqdm(file_df.iterrows(), total=len(file_df), desc="Paths", unit="path", leave=False) if show_progress else file_df.iterrows()
        for _, row in row_iter:
            vehicle_id = int(row["vehicle_id"])
            trip_id = int(row["trip_id"])
            path_sequence = row["path_sequence"]
            if pd.isna(path_sequence) or not str(path_sequence).strip():
                continue
            steps = []
            for segment in str(path_sequence).strip().split("_"):
                segment = segment.strip()
                if not segment:
                    continue
                parts = segment.split("-")
                if len(parts) >= 2:
                    try:
                        nid = int(parts[0])
                        tval = float(parts[1])
                        steps.append((nid, tval))
                    except (ValueError, TypeError):
                        continue
            if not steps:
                continue
            for seq_idx, (node_id, time_value) in enumerate(steps):
                ts = pd.Timestamp(year=base_date.year, month=base_date.month, day=base_date.day) + pd.Timedelta(seconds=float(time_value)) if base_date else pd.NaT
                exit_edge = None
                if seq_idx + 1 < len(steps):
                    next_node = steps[seq_idx + 1][0]
                    exit_edge = edge_map.get((node_id, next_node))
                rows.append({
                    "walk_id": walk_id_counter,
                    "sequence_idx": seq_idx,
                    "node_id": node_id,
                    "exit_edge": exit_edge,
                    "timestamp": ts,
                    "vehicle_id": vehicle_id,
                    "trip_id": trip_id,
                })
            walk_id_counter += 1
        if not rows:
            result.append((date_str, pd.DataFrame(columns=["walk_id", "sequence_idx", "node_id", "exit_edge", "timestamp", "vehicle_id", "trip_id"])))
        else:
            out = pd.DataFrame(rows)
            out["walk_id"] = out["walk_id"].astype("Int64")
            out["sequence_idx"] = out["sequence_idx"].astype("Int64")
            out["node_id"] = out["node_id"].astype("Int64")
            out["vehicle_id"] = out["vehicle_id"].astype("Int64")
            out["trip_id"] = out["trip_id"].astype("Int64")
            out["exit_edge"] = out["exit_edge"].astype("Int64")
            result.append((date_str, out))
    return result


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Convert Shenzhen- or Jinan-style node/edge/trajectory CSVs to walkingpandas Parquet (nodes, edges, walks).",
    )
    parser.add_argument(
        "--city",
        type=str,
        choices=list(CITY_CONFIG),
        default="shenzhen",
        help="City preset for default node/edge/trajectory file names (default: shenzhen)",
    )
    parser.add_argument(
        "--data-dir",
        type=Path,
        default=Path(os.environ.get("WP_DATA_DIR", ".")),
        help="Directory containing node, edge, and trajectory CSVs (default: WP_DATA_DIR or .)",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(os.environ.get("WP_OUTPUT_DIR", ".")),
        help="Base directory for output; writes to <output-dir>/<city>/ (default: WP_OUTPUT_DIR or .)",
    )
    parser.add_argument(
        "--crs",
        type=str,
        default="EPSG:4326",
        help="Coordinate reference system for node/edge geometry (default: EPSG:4326)",
    )
    parser.add_argument(
        "--nodes-csv",
        type=str,
        default=None,
        help="Node CSV filename (overrides --city default)",
    )
    parser.add_argument(
        "--edges-csv",
        type=str,
        default=None,
        help="Edge CSV filename (overrides --city default)",
    )
    parser.add_argument(
        "--traj-csv",
        type=str,
        default=None,
        help="Trajectory CSV file or glob pattern (overrides --city default)",
    )
    parser.add_argument(
        "--traj-date",
        type=str,
        default=None,
        help="Default trajectory date YYYY-MM-DD when filename has no date (e.g. jinan: 2022-10-17)",
    )
    parser.add_argument(
        "--walks-per-file",
        action="store_true",
        help="Write each trajectory file to a separate walks_YYYYMMDD.parquet (default: single walks.parquet)",
    )
    parser.add_argument(
        "--crs-sidecar",
        action="store_true",
        help="Write crs.txt in output directory with the CRS string",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Only check input files exist and print paths; do not write",
    )
    parser.add_argument(
        "--no-progress",
        action="store_true",
        help="Disable progress bars (e.g. for CI or piping)",
    )
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=0,
        metavar="N",
        help="Read CSVs in chunks of N rows; write fragment Parquets and merge with PyArrow (no full load). 0 = disabled.",
    )
    parser.add_argument(
        "--temp-dir",
        type=Path,
        default=None,
        help="Temp directory for chunk Parquets (default: <output-dir>/<city>/.chunks). Cleaned after merge.",
    )
    parser.add_argument(
        "--no-merge",
        action="store_true",
        help="With --chunk-size: do not merge walk chunks; write walks/chunk_*.parquet for downstream glob read.",
    )
    args = parser.parse_args()

    data_dir = args.data_dir.resolve()
    city = args.city
    config = CITY_CONFIG[city]
    # Output under output_dir/city/ (e.g. output_dir/shenzhen/ or output_dir/jinan/)
    output_dir = (args.output_dir.resolve() / city)

    nodes_name = args.nodes_csv or config["nodes"]
    edges_name = args.edges_csv or config["edges"]
    nodes_path = data_dir / nodes_name
    edges_path = data_dir / edges_name
    default_traj_date = args.traj_date or config.get("default_traj_date")

    if not nodes_path.exists():
        print(f"Error: Node CSV not found: {nodes_path}")
        return 1
    if not edges_path.exists():
        print(f"Error: Edge CSV not found: {edges_path}")
        return 1

    # Resolve trajectory file(s)
    if args.traj_csv:
        p = data_dir / args.traj_csv
        if p.exists() and p.is_file():
            traj_paths = [p]
        else:
            traj_paths = sorted(data_dir.glob(args.traj_csv))
    else:
        pattern = config["traj_pattern"]
        if "*" in pattern:
            traj_paths = sorted(data_dir.glob(pattern))
        else:
            # Single file (e.g. jinan: traj_jinan.csv)
            single = data_dir / pattern
            traj_paths = [single] if single.exists() else []

    if not traj_paths:
        print(f"Error: No trajectory file(s) found in {data_dir} (tried pattern or --traj-csv)")
        return 1

    if args.dry_run:
        print("Dry run: input files OK")
        print(f"  City: {city}")
        print(f"  Nodes: {nodes_path}")
        print(f"  Edges: {edges_path}")
        print(f"  Trajectories: {traj_paths}")
        print(f"  Output dir: {output_dir}")
        writes = "walks_YYYYMMDD.parquet (per file)" if args.walks_per_file else "walks.parquet"
        chunk_note = f" (chunked: {args.chunk_size})" if args.chunk_size else ""
        print(f"  Would write: nodes.parquet, edges.parquet, {writes}{chunk_note}" + (" + crs.txt" if args.crs_sidecar else ""))
        return 0

    show_progress = not args.no_progress
    output_dir.mkdir(parents=True, exist_ok=True)
    nodes_out = output_dir / "nodes.parquet"
    edges_out = output_dir / "edges.parquet"
    temp_dir = args.temp_dir or (output_dir / ".chunks")
    use_chunked = args.chunk_size > 0

    if use_chunked:
        temp_dir.mkdir(parents=True, exist_ok=True)
        _build_nodes_chunked(
            nodes_path, args.crs, args.chunk_size, temp_dir, nodes_out, show_progress
        )
        edge_map = _build_edges_chunked(
            edges_path, args.crs, args.chunk_size, temp_dir, edges_out, show_progress
        )
        walk_steps, walk_count = _build_walks_chunked(
            traj_paths,
            edge_map,
            default_traj_date,
            args.chunk_size,
            temp_dir,
            output_dir,
            no_merge=args.no_merge,
            walks_per_file=args.walks_per_file,
            show_progress=show_progress,
        )
        if temp_dir.exists():
            shutil.rmtree(temp_dir, ignore_errors=True)
        node_count = len(pd.read_parquet(nodes_out))
        edge_count = len(pd.read_parquet(edges_out))
        print(f"Wrote {nodes_out} ({node_count} nodes)")
        print(f"Wrote {edges_out} ({edge_count} edges)")
        if args.no_merge:
            print(f"Wrote {output_dir / 'walks'} ({walk_steps} walk steps, {walk_count} walks)")
        elif args.walks_per_file:
            print(f"Walks: {walk_steps} steps, {walk_count} walks")
        else:
            print(f"Wrote {output_dir / 'walks.parquet'} ({walk_steps} walk steps, {walk_count} walks)")
    else:
        nodes_gdf = build_nodes_gdf(nodes_path, args.crs, show_progress=show_progress)
        edges_gdf, edge_map = build_edges_gdf(edges_path, args.crs, show_progress=show_progress)
        nodes_gdf.to_parquet(nodes_out, index=False)
        edges_gdf.to_parquet(edges_out, index=False)
        print(f"Wrote {nodes_out} ({len(nodes_gdf)} nodes)")
        print(f"Wrote {edges_out} ({len(edges_gdf)} edges)")

        if args.walks_per_file:
            per_file = build_walks_dfs_per_file(
                traj_paths, edge_map, show_progress=show_progress, default_traj_date=default_traj_date
            )
            for date_str, wdf in per_file:
                if wdf.empty:
                    continue
                validate_schema(wdf, strict=True)
                walks_out = output_dir / f"walks_{date_str}.parquet"
                write_walks(wdf, walks_out, schema_validate=True)
                print(f"Wrote {walks_out} ({len(wdf)} walk steps, {wdf['walk_id'].nunique()} walks)")
        else:
            walks_df = build_walks_df(
                traj_paths, edge_map, show_progress=show_progress, default_traj_date=default_traj_date
            )
            validate_schema(walks_df, strict=True)
            walks_out = output_dir / "walks.parquet"
            write_walks(walks_df, walks_out, schema_validate=True)
            print(f"Wrote {walks_out} ({len(walks_df)} walk steps, {walks_df['walk_id'].nunique()} walks)")

    if args.crs_sidecar:
        (output_dir / "crs.txt").write_text(args.crs, encoding="utf-8")
        print(f"Wrote {output_dir / 'crs.txt'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
