"""
Converter for city-scale vehicle trajectory datasets (Shenzhen, Jinan).

Transforms raw CSV nodes, edges, and trajectories into the standardized
walkingpandas Parquet schema. Designed to be a standalone, easily editable
script with a decoupled data downloader.

Data source
----------
The script expects data from the **City-scale Vehicle Trajectory Data from
Traffic Camera Videos** dataset (Yu et al., 2023), published in *Scientific
Data*.

- **Source:** Figshare collection (DOI: 10.6084/m9.figshare.c.6676199)
- **URL:** https://figshare.com/collections/City-scale_mobility_dataset/6676199
- **Data:** ~5 million vehicle trajectories recovered from 3,000+ traffic
  cameras across Shenzhen and Jinan (China). Trajectories are extracted from
  traffic camera videos via a cross-camera recovery framework, not from GPS.
- **Access:** Data can be downloaded programmatically via this script
  (FigshareDownloader) or manually from the Figshare website in a browser.

Citation
-------
Yu, F., Yan, H., Chen, R., Zhang, G., Liu, Y., Chen, M., & Li, Y. (2023).
City-scale vehicle trajectory data from traffic camera videos. *Scientific
Data*, 10, 711. https://doi.org/10.1038/s41597-023-02589-y
"""

import re
import shutil
import zipfile
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

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

from walkingpandas.io.ingest import write_walks
from tqdm.auto import tqdm

try:
    import requests
except ImportError:
    requests = None

# Modify these variables to point to your local dataset files.

# 1. Path to the folder containing the raw CSVs (will be created automatically if missing)
DATA_DIR = Path("./raw_data")
    
# 2. Path to where the Parquet files should be exported
OUTPUT_DIR = Path("./processed_data")

# 3. Which configuration preset(s) to use. Can be a single string ('shenzhen'),
# a list of strings (['shenzhen', 'jinan']), or 'all' to process everything.
CITIES = "all"
    
# 4. How many CSV rows to process at once before writing to Parquet
CHUNK_SIZE = 100000
    
# 5. Automatically pull data from Figshare if not available locally
AUTO_DOWNLOAD = True
    
# 6. Merge the generated chunks into a single walks.parquet file and delete the chunk directory
MERGE_CHUNKS = True


class FigshareDownloader:
    """Downloads all dataset files from a Figshare collection dynamically."""
    
    def __init__(self, collection_id: int, data_dir: Union[str, Path]):
        self.collection_id = collection_id
        self.data_dir = Path(data_dir)
        self.api_base = "https://api.figshare.com/v2"

    def download(self):
        if requests is None:
            raise ImportError(
                "The 'requests' library is required to download data. "
                "Run: pip install requests"
            )
            
        self.data_dir.mkdir(parents=True, exist_ok=True)
        
        try:
            print(f"Fetching metadata for Figshare collection {self.collection_id}...")
            articles_url = f"{self.api_base}/collections/{self.collection_id}/articles"
            resp = requests.get(articles_url, timeout=15)
            resp.raise_for_status()
            articles = resp.json()
        except Exception as e:
            print(f"Warning: Could not reach Figshare ({e}). Assuming local files are complete.")
            return
        
        downloaded_any = False

        # 2. Iterate articles and download all files
        for article in articles:
            article_id = article['id']
            
            try:
                files_url = f"{self.api_base}/articles/{article_id}/files"
                f_resp = requests.get(files_url, timeout=15)
                f_resp.raise_for_status()
                files = f_resp.json()
            except Exception as e:
                print(f"Warning: Could not fetch files for article {article_id} ({e}).")
                continue
            
            for f in files:
                download_url = f['download_url']
                target_path = self.data_dir / f['name']
                
                # Check if the file is completely missing OR if it is 0 bytes (corrupted/partial)
                if not target_path.exists() or target_path.stat().st_size == 0:
                    self._download_file(download_url, target_path)
                    downloaded_any = True
                    
                    # Automatically unpack zipped data
                    if target_path.suffix == '.zip':
                        self._extract_zip(target_path)
                            
        if not downloaded_any:
            print(f"All raw files for collection {self.collection_id} appear to be downloaded and valid.")

    def _download_file(self, url: str, target_path: Path):
        try:
            with requests.get(url, stream=True, timeout=15) as r:
                r.raise_for_status()
                total_size = int(r.headers.get('content-length', 0))
                with open(target_path, 'wb') as f, tqdm(
                    desc=target_path.name,
                    total=total_size,
                    unit='iB',
                    unit_scale=True,
                    unit_divisor=1024,
                ) as pbar:
                    for chunk in r.iter_content(chunk_size=8192):
                        size = f.write(chunk)
                        pbar.update(size)
        except Exception as e:
            print(f"\nError downloading {target_path.name}: {e}")
            # If the download failed halfway, delete the partial 0-byte file 
            # so it triggers a fresh redownload on the next run
            if target_path.exists():
                target_path.unlink()

    def _extract_zip(self, zip_path: Path):
        print(f"Extracting {zip_path.name}...")
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            # Flatten the extraction so files land directly in data_dir (ignore folders)
            for member in zip_ref.namelist():
                filename = Path(member).name
                if not filename:  # Skip directories
                    continue
                source = zip_ref.open(member)
                target = open(self.data_dir / filename, "wb")
                with source, target:
                    shutil.copyfileobj(source, target)
                    
        # Note: We intentionally do NOT delete the zip file.
        # Keeping it prevents the downloader from redownloading it next time
        # since it uses the zip file's existence as a completion marker.


class CityDatasetConverter:
    """
    Cleans and converts standard road network and trajectory CSVs.
    
    Usage:
        >>> from walkingpandas.scripts.ingest_traffic_cameras import CityDatasetConverter
        >>> converter = CityDatasetConverter(data_dir="raw_data", output_dir="processed", city="shenzhen")
        >>> converter.run()
    """

    # Pre-configured file patterns and settings for different datasets
    CONFIGS = {
        "shenzhen": {
            "nodes": "node_shenzhen.csv",
            "edges": "edge_shenzhen.csv",
            "traj_pattern": "traj_shenzhen_*.csv",
            "default_date": None,
        },
        "jinan": {
            "nodes": "node_jinan.csv",
            "edges": "edge_jinan.csv",
            "traj_pattern": "traj_jinan.csv",
            "default_date": "2022-10-17",
        },
    }

    def __init__(
        self,
        data_dir: Union[str, Path],
        output_dir: Union[str, Path],
        city: str = "shenzhen",
        crs: str = "EPSG:4326"
    ):
        self.data_dir = Path(data_dir)
        self.output_dir = Path(output_dir) / city.lower()
        self.city = city.lower()
        self.crs = crs
        
        if self.city not in self.CONFIGS:
            raise ValueError(f"Unknown city preset '{city}'. Available: {list(self.CONFIGS.keys())}")
            
        self.config = self.CONFIGS[self.city]
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self._edge_map: Dict[Tuple[int, int], int] = {}

    def run(self, chunk_size: int = 100000, merge_chunks: bool = True):
        """Executes the full conversion pipeline."""
        print(f"Starting conversion for {self.city.title()} dataset...")
            
        self.process_nodes()
        self.process_edges()
        self.process_trajectories(chunk_size=chunk_size, merge_chunks=merge_chunks)
        print(f"\nConversion complete. Data saved to: {self.output_dir}")

    def process_nodes(self):
        """Parse node CSV and export to nodes.parquet."""
        nodes_path = self.data_dir / self.config["nodes"]
        print(f"\n1. Processing nodes from {nodes_path.name}...")
        
        if not nodes_path.exists() or nodes_path.stat().st_size == 0:
            print(f"   -> Error: Could not find {nodes_path} or file is empty.")
            return
            
        df = pd.read_csv(nodes_path)
        # Standardize columns to avoid KeyError on spaces/capitalization
        df.columns = df.columns.str.strip().str.lower()
        df = df.rename(columns={
            "nodeid": "node_id",
            "longitude": "longitude",
            "latitude": "latitude",
            "hascamera": "has_camera",
        })
        
        # Cast to nullable types
        df["node_id"] = df["node_id"].astype("Int64")
        if "has_camera" in df.columns:
            df["has_camera"] = df["has_camera"].astype(bool)
        
        geometry = [Point(lon, lat) for lon, lat in zip(df["longitude"], df["latitude"])]
        gdf = gpd.GeoDataFrame(df, geometry=geometry, crs=self.crs)
        
        out_path = self.output_dir / "nodes.parquet"
        gdf.to_parquet(out_path, index=False)
        print(f"   -> Saved {len(gdf)} nodes to {out_path.name}")

    def _parse_linestring(self, geom_str: str) -> Optional[LineString]:
        """Convert underscore-separated coordinate strings to Shapely LineStrings."""
        if not isinstance(geom_str, str) or not geom_str.strip():
            return None
            
        coords = []
        for segment in geom_str.strip().split("_"):
            if not segment:
                continue
            parts = segment.split("-")
            if len(parts) >= 2:
                try:
                    coords.append((float(parts[0]), float(parts[1])))
                except ValueError:
                    continue
                    
        return LineString(coords) if len(coords) >= 2 else None

    def process_edges(self):
        """Parse edge CSV, build geometric linestrings, and export to edges.parquet."""
        edges_path = self.data_dir / self.config["edges"]
        print(f"\n2. Processing edges from {edges_path.name}...")

        if not edges_path.exists() or edges_path.stat().st_size == 0:
            print(f"   -> Error: Could not find {edges_path} or file is empty.")
            return

        df = pd.read_csv(edges_path)
        df.columns = df.columns.str.strip().str.lower()
        df = df.rename(columns={
            "origin": "source_node",
            "destination": "target_node",
            "class": "road_class",
            "geometry": "geometry_encoded",
            "length": "length_m",
        })
        
        # Inject sequential edge ID required by walkingpandas
        df.insert(0, "edge_id", range(len(df)))
        
        df["source_node"] = df["source_node"].astype(int)
        df["target_node"] = df["target_node"].astype(int)
        if "geometry_encoded" in df.columns:
            df["geometry"] = df["geometry_encoded"].apply(self._parse_linestring)
            df = df.drop(columns=["geometry_encoded"])
        
        # Cache the edge mapping for O(1) lookups during trajectory parsing
        self._edge_map = {
            (row.source_node, row.target_node): row.edge_id 
            for row in df.itertuples()
        }
        
        gdf = gpd.GeoDataFrame(df, geometry="geometry", crs=self.crs)
        out_path = self.output_dir / "edges.parquet"
        gdf.to_parquet(out_path, index=False)
        print(f"   -> Saved {len(gdf)} edges to {out_path.name}")

    def _extract_date(self, filepath: Path) -> Optional[datetime]:
        """Extract the base date from a filename, or fallback to the config default."""
        pattern = re.compile(r"(\d{8})\.csv$", re.IGNORECASE)
        match = pattern.search(filepath.name)
        if match:
            return datetime.strptime(match.group(1), "%Y%m%d")
            
        default = self.config["default_date"]
        return datetime.strptime(default, "%Y-%m-%d") if default else None

    def _parse_trajectory_points(self, points_str: str) -> List[Tuple[int, float]]:
        """Parse raw point string into (node_id, time_seconds) tuples."""
        steps = []
        if not isinstance(points_str, str) or not points_str.strip():
            return steps
            
        for segment in points_str.strip().split("_"):
            if not segment:
                continue
            parts = segment.split("-")
            if len(parts) >= 2:
                try:
                    steps.append((int(parts[0]), float(parts[1])))
                except ValueError:
                    continue
        return steps

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

    def process_trajectories(self, chunk_size: int = 100000, merge_chunks: bool = True):
        """Parse and stream trajectories into chunked Parquet files."""
        pattern = self.config["traj_pattern"]
        traj_files = sorted(self.data_dir.glob(pattern))
        
        valid_traj_files = [f for f in traj_files if f.stat().st_size > 0]
        
        if not valid_traj_files:
            print(f"\nWarning: No valid trajectory files found matching {pattern} in {self.data_dir}")
            return

        # Pre-count rows for an accurate progress bar
        total_rows = sum(self._count_csv_rows(p) for p in valid_traj_files)
        print(f"\n3. Processing {len(valid_traj_files)} trajectory file(s) (~{total_rows:,} rows)...")
        
        walks_out_dir = self.output_dir / "walks"
        walks_out_dir.mkdir(exist_ok=True)
        
        walk_id_counter = 0
        total_steps = 0
        chunk_paths = []
        
        with tqdm(total=total_rows, desc="Trajectories", unit="row") as pbar:
            for file_idx, filepath in enumerate(valid_traj_files):
                base_dt = self._extract_date(filepath)
                # Compute base timestamp natively outside the chunk loops
                base_ts = pd.Timestamp(base_dt.date()) if base_dt else None
                
                # Stream large CSVs to bound memory usage
                for chunk_idx, chunk_df in enumerate(pd.read_csv(filepath, chunksize=chunk_size)):
                    
                    # Resolve column indices robustly (ignoring spaces/caps)
                    chunk_df.columns = chunk_df.columns.str.strip().str.lower()
                    cols = list(chunk_df.columns)
                    veh_idx = cols.index("vehicleid") + 1 if "vehicleid" in cols else -1
                    trip_idx = cols.index("tripid") + 1 if "tripid" in cols else -1
                    pts_idx = cols.index("points") + 1 if "points" in cols else -1
                    
                    rows = []
                    for row in chunk_df.itertuples():
                        if pts_idx == -1:
                            continue
                            
                        path_sequence = row[pts_idx]
                        steps = self._parse_trajectory_points(path_sequence if pd.notna(path_sequence) else "")
                        if not steps:
                            continue
                            
                        v_id = int(row[veh_idx]) if veh_idx != -1 else -1
                        t_id = int(row[trip_idx]) if trip_idx != -1 else -1
                        
                        for seq_idx, (node_id, time_seconds) in enumerate(steps):
                            # Determine exit_edge looking ahead to the next step
                            exit_edge = None
                            if seq_idx + 1 < len(steps):
                                next_node = steps[seq_idx + 1][0]
                                exit_edge = self._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,
                                "time_seconds": time_seconds,
                                "vehicle_id": v_id,
                                "trip_id": t_id,
                            })
                            
                        walk_id_counter += 1
                        
                    if rows:
                        out_df = pd.DataFrame(rows)
                        
                        # Vectorize timestamp math (Orders of magnitude faster than building pd.Timestamp in the loop)
                        if base_ts is not None:
                            out_df["timestamp"] = base_ts + pd.to_timedelta(out_df["time_seconds"], unit="s")
                        else:
                            out_df["timestamp"] = pd.NaT
                        
                        out_df = out_df.drop(columns=["time_seconds"])
                        
                        # Standardize to walkingpandas schema format
                        out_df["walk_id"] = out_df["walk_id"].astype("Int64")
                        out_df["sequence_idx"] = out_df["sequence_idx"].astype("Int64")
                        out_df["node_id"] = out_df["node_id"].astype("Int64")
                        out_df["exit_edge"] = out_df["exit_edge"].astype("Int64")
                        out_df["vehicle_id"] = out_df["vehicle_id"].astype("Int64")
                        out_df["trip_id"] = out_df["trip_id"].astype("Int64")
                        
                        total_steps += len(out_df)
                        
                        # Write chunk directly to the directory
                        chunk_path = walks_out_dir / f"chunk_{file_idx:03d}_{chunk_idx:04d}.parquet"
                        write_walks(out_df, chunk_path, schema_validate=False)
                        chunk_paths.append(chunk_path)
                        
                    # Update progress bar AT THE END of processing the chunk
                    pbar.update(len(chunk_df))
                
        if merge_chunks and chunk_paths:
            merged_path = self.output_dir / "walks.parquet"
            print(f"   -> Merging {len(chunk_paths)} chunks into a single {merged_path.name}...")
            
            schema = pq.ParquetFile(chunk_paths[0]).schema_arrow
            with pq.ParquetWriter(merged_path, schema) as writer:
                for p in tqdm(chunk_paths, desc="Merging chunks", unit="file"):
                    writer.write_table(pq.read_table(p))
                    
            print(f"   -> Cleaning up temporary chunk directory...")
            shutil.rmtree(walks_out_dir, ignore_errors=True)
            print(f"   -> Saved {walk_id_counter:,} walks ({total_steps:,} total steps) to {merged_path.name}")
        else:
            print(f"   -> Generated {walk_id_counter:,} walks ({total_steps:,} total steps) in {walks_out_dir.name}/ directory")


def main():
    """
    Internal script configuration.
    """

    print("=" * 50)
    print(f" WalkingPandas Dataset Cleaner")
    print("=" * 50)
    
    # Parse requested cities
    if isinstance(CITIES, str):
        if CITIES.lower() == "all":
            cities_to_process = list(CityDatasetConverter.CONFIGS.keys())
        else:
            cities_to_process = [CITIES.lower()]
    else:
        cities_to_process = [c.lower() for c in CITIES]

    # --- PHASE 1: Download Data (Optional) ---
    if AUTO_DOWNLOAD:
        print(f"\nVerifying raw data in {DATA_DIR} via Figshare...")
        # 6676199 is the Yu2023 City-scale Vehicle Trajectory Data Collection ID
        downloader = FigshareDownloader(collection_id=6676199, data_dir=DATA_DIR)
        downloader.download()
        
        # Verify files are present and valid post-download
        for city in cities_to_process:
            config = CityDatasetConverter.CONFIGS.get(city)
            if not config:
                continue
                
            nodes_path = DATA_DIR / config["nodes"]
            edges_path = DATA_DIR / config["edges"]
            traj_files = list(DATA_DIR.glob(config["traj_pattern"]))
            valid_traj_files = [f for f in traj_files if f.stat().st_size > 0]
            
            if not nodes_path.exists() or nodes_path.stat().st_size == 0 or \
               not edges_path.exists() or edges_path.stat().st_size == 0 or \
               not valid_traj_files:
                raise FileNotFoundError(
                    f"Failed to locate valid required files for {city}. "
                    f"Please verify the files are present in {DATA_DIR} and are not empty."
                )

    # --- PHASE 2: Clean and Convert Data ---
    for city in cities_to_process:
        if city not in CityDatasetConverter.CONFIGS:
            print(f"Unknown city preset: {city}. Skipping.")
            continue
            
        print(f"\n{'=' * 40}")
        print(f" Processing {city.title()}")
        print(f"{'=' * 40}")
        
        converter = CityDatasetConverter(
            data_dir=DATA_DIR,
            output_dir=OUTPUT_DIR,
            city=city
        )
        converter.run(chunk_size=CHUNK_SIZE, merge_chunks=MERGE_CHUNKS)


if __name__ == "__main__":
    main()