"""
SHP (Shapefile) reader/writer using geopandas.
"""
import os
from pathlib import Path
import geopandas as gpd
import pandas as pd
from shapely.geometry import Point
from core.data_model import CoordData, CoordPoint


class ShpHandler:
    EXTENSIONS = {".shp"}

    @classmethod
    def read(cls, path: str) -> CoordData:
        gdf = gpd.read_file(path, encoding="utf-8")
        if gdf.empty:
            raise ValueError(f"Shapefile 无数据: {path}")

        crs = gdf.crs.to_string() if gdf.crs else None

        # Extract centroids for non-Point geometries
        geom_col = gdf.geometry.name
        points = []
        for _, row in gdf.iterrows():
            geom = row[geom_col]
            if geom is None or geom.is_empty:
                continue
            centroid = geom.centroid if geom.geom_type != "Point" else geom
            pt = CoordPoint(
                x=centroid.x,
                y=centroid.y,
                z=centroid.z if hasattr(centroid, "z") and centroid.z else None,
                attrs={k: v for k, v in row.items() if k != geom_col and not pd.isna(v)},
            )
            points.append(pt)

        cols = [c for c in gdf.columns if c != geom_col]
        return CoordData(points=points, crs=crs, columns=cols,
                         geometry_type=gdf.geom_type.iloc[0] if len(gdf) > 0 else "Point")

    @classmethod
    def write(cls, data: CoordData, path: str, crs: str | None = None):
        records = data.to_records()
        if not records:
            raise ValueError("无数据可写入")

        # Build geometry column
        geometry = [Point(r.pop("X"), r.pop("Y")) for r in records]

        gdf = gpd.GeoDataFrame(records, geometry=geometry, crs=crs or data.crs)
        gdf.to_file(path, encoding="utf-8")