import os
import tempfile
import logging
import warnings
import geopandas as gpd
import osmnx as ox
import httpx
from typing import Optional, Any
from fsspec.core import url_to_fs
from sibi_flux.core import SecureResource

# --- Production Tuning ---
# Mitigate the "79x Area" warning by allowing larger query chunks and longer timeouts
ox.settings.max_query_area_size = 2_500_000_000
ox.settings.requests_timeout = 600
# Silence the warning in logs once we've tuned the settings
warnings.filterwarnings("ignore", category=UserWarning, module="osmnx._overpass")


class PBFHandler(SecureResource):
    """
    Handles OSMnx graph lifecycle: Build (Overpass or Local PBF),
    Atomic Persistence (fsspec), and Memory-efficient Loading.

    Optimizations:
    - Uses Parquet for GeoDataFrames (nodes/edges).
    - Uses GraphML for NetworkX graph (safer, portable).
    """

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self.kwargs = kwargs
        self.graph = None
        self.nodes: Optional[gpd.GeoDataFrame] = None
        self.edges: Optional[gpd.GeoDataFrame] = None

        self.rebuild: bool = kwargs.get("rebuild", False)
        self.place: str = kwargs.get("place", "Costa Rica")
        self.network_type: str = kwargs.get("network_type", "all")

        # Filesystem Abstraction
        base_url: str = kwargs.get("data_path", "default_depot/pbf_files")
        prefix: str = kwargs.get("files_prefix", "costa-rica-").rstrip("-") + "-"

        # We inherit self.fs and self.base resolution logic if passed, 
        # but PBFHandler has its own fs/base logic.
        fs = kwargs.get("fs")
        if fs is not None:
            self.fs, self.base = fs, base_url.rstrip("/")
        else:
            self.fs, self.base = url_to_fs(base_url)

        self.fs.mkdirs(self.base, exist_ok=True)
        self._prefix = prefix

    @property
    def graph_path(self) -> str:
        return f"{self.base}/{self._prefix}graph.graphml"

    @property
    def node_path(self) -> str:
        return f"{self.base}/{self._prefix}nodes.parquet"

    @property
    def edge_path(self) -> str:
        return f"{self.base}/{self._prefix}edges.parquet"

    def load(self) -> None:
        """Main entry point for retrieving graph data."""
        if self.rebuild:
            self._delete_artifacts()

        if not self._artifacts_exist():
            self.process_pbf()
            self.save_artifacts()
        else:
            self.load_artifacts()

    def process_pbf(self) -> None:
        """Builds graph from local/remote .pbf source if available, else hits Overpass."""
        self.logger.info(f"Building graph: {self.place}")

        pbf_path = self.kwargs.get("pbf_path")
        pbf_url = self.kwargs.get("pbf_url")

        if pbf_path:
            # 1. Check if file exists in our FS (S3 or local)
            if not self.fs.exists(pbf_path):
                # 2. If valid URL provided, download and save to FS
                if pbf_url:
                    self.logger.info(
                        f"PBF missing at {pbf_path}, downloading from {pbf_url}..."
                    )
                    self._download_and_save_pbf(pbf_url, pbf_path)

            # 3. If file exists (pre-existing or just downloaded), process it
            if self.fs.exists(pbf_path):
                self.logger.info(f"Using source: {pbf_path}")

                # OSMnx requires a local file path.
                # If FS is remote (NetworkFileSystem), we must download to temp local file.
                # If FS is local (LocalFileSystem), we use the direct path.
                try:
                    # fsspec's helper `open_local` might be useful, but explicit tempfile is safer for cleanup
                    import tempfile
                    from fsspec.implementations.local import LocalFileSystem

                    if isinstance(self.fs, LocalFileSystem):
                        # Use actual path if local
                        local_path = pbf_path
                        # Handle case where pbf_path might be relative to self.base or arbitrary
                        # If pbf_path is absolute/relative on local fs, just use it.
                        self.graph = ox.graph_from_xml(local_path)
                    else:
                        # Remote FS: Stream to temp file
                        # Stream PBF to temp file for osmium
                        with tempfile.NamedTemporaryFile(suffix=".osm.pbf") as tmp:
                            self.logger.info(
                                f"Streaming remote PBF from {pbf_path} to local temporary file..."
                            )
                            with self.fs.open(pbf_path, "rb") as remote_f:
                                while True:
                                    chunk = remote_f.read(1024 * 1024)
                                    if not chunk:
                                        break
                                    tmp.write(chunk)
                            tmp.flush()

                            self.logger.info(
                                f"Processing local PBF file: {tmp.name} with Osmium..."
                            )
                            self._process_pbf_with_osmium(tmp.name)

                except Exception as e:
                    self.logger.error(
                        f"Error processing PBF {pbf_path}: {e}", exc_info=True
                    )
                    self._fallback_to_place()
            else:
                self._fallback_to_place()
        else:
            self._fallback_to_place()

    def _process_pbf_with_osmium(self, pbf_path: str) -> None:
        """Uses osmium to read PBF and build GeoDataFrames, then Graph."""
        import osmium
        from shapely.geometry import Point, LineString
        import pandas as pd

        class GraphHandler(osmium.SimpleHandler):
            def __init__(self):
                super().__init__()
                # Use lists of primitives for scalar data (highly memory efficient)
                self.nodes_osmid = []
                self.nodes_x = []
                self.nodes_y = []
                
                self.edge_us = []
                self.edge_vs = []
                self.edge_osmids = []
                self.edge_highways = []
                self.edge_oneways = []
                self.edge_geoms = []
                # Metadata tags
                self.edge_names = []
                self.edge_maxspeeds = []
                self.edge_lanes = []
                self.edge_refs = []

            def way(self, w):
                if "highway" in w.tags:
                    hw = w.tags["highway"]
                    if hw in [
                        "motorway", "trunk", "primary", "secondary", 
                        "tertiary", "unclassified", "residential", "service",
                    ]:
                        try:
                            # 1. Collect coords for geometry immediately
                            coords = [(n.lon, n.lat) for n in w.nodes]
                            if len(coords) < 2:
                                return

                            # 2. Append edge scalars
                            self.edge_us.append(w.nodes[0].ref)
                            self.edge_vs.append(w.nodes[-1].ref)
                            self.edge_osmids.append(w.id)
                            self.edge_highways.append(hw)
                            self.edge_oneways.append(w.tags.get("oneway") == "yes")
                            self.edge_geoms.append(LineString(coords))
                            
                            # 3. Handle optional tags using None for missing (keeps lists aligned)
                            self.edge_names.append(w.tags.get("name"))
                            self.edge_maxspeeds.append(w.tags.get("maxspeed"))
                            self.edge_lanes.append(w.tags.get("lanes"))
                            self.edge_refs.append(w.tags.get("ref"))

                            # 4. Collect node scalars (deduplication happens in DataFrame)
                            for n in w.nodes:
                                self.nodes_osmid.append(n.ref)
                                self.nodes_x.append(n.lon)
                                self.nodes_y.append(n.lat)

                        except (osmium.InvalidLocationError, Exception):
                            pass

        self.logger.info("Initializing Osmium handler...")
        handler = GraphHandler()
        # locations=True stores node coords in memory (essential for way processing)
        # For Costa Rica (~40MB PBF), this is fine for memory.
        self.logger.info("Applying PBF file to handler (extracting nodes/ways)...")
        handler.apply_file(pbf_path, locations=True)

        self.logger.info(
            f"Extracted {len(handler.nodes_osmid)} raw node points and {len(handler.edge_us)} edges."
        )

        if not handler.edge_us:
            raise ValueError("No edges extracted from PBF.")

        self.logger.info("Building GeoDataFrames (Vectorized)...")
        # 1. Build Nodes GDF using vectorized points_from_xy (huge memory win)
        df_nodes = pd.DataFrame({
            "osmid": handler.nodes_osmid,
            "x": handler.nodes_x,
            "y": handler.nodes_y,
        }).drop_duplicates(subset=["osmid"]).set_index("osmid")
        
        # Explicitly release raw node lists immediately
        del handler.nodes_osmid
        del handler.nodes_x
        del handler.nodes_y
        import gc; gc.collect()

        gdf_nodes = gpd.GeoDataFrame(
            df_nodes, 
            geometry=gpd.points_from_xy(df_nodes.x, df_nodes.y), 
            crs="EPSG:4326"
        )
        self.nodes = gdf_nodes

        # 2. Build Edges GDF
        df_edges = pd.DataFrame({
            "u": handler.edge_us,
            "v": handler.edge_vs,
            "osmid": handler.edge_osmids,
            "highway": handler.edge_highways,
            "oneway": handler.edge_oneways,
            "geometry": handler.edge_geoms,
            "name": handler.edge_names,
            "maxspeed": handler.edge_maxspeeds,
            "lanes": handler.edge_lanes,
            "ref": handler.edge_refs,
        })
        
        # Release raw edge lists
        del handler.edge_us
        del handler.edge_vs
        del handler.edge_osmids
        del handler.edge_highways
        del handler.edge_oneways
        del handler.edge_geoms
        del handler.edge_names
        del handler.edge_maxspeeds
        del handler.edge_lanes
        del handler.edge_refs
        gc.collect()

        if not df_edges.empty:
            df_edges["key"] = df_edges.groupby(["u", "v"]).cumcount()
            gdf_edges = gpd.GeoDataFrame(df_edges, geometry="geometry", crs="EPSG:4326")
            gdf_edges.set_index(["u", "v", "key"], inplace=True)
            self.edges = gdf_edges

            self.logger.info("Constructing NetworkX graph from GDFs (this creates a copy)...")
            self.graph = ox.graph_from_gdfs(gdf_nodes, gdf_edges)
            
            # REDUNDANCY REMOVAL: The Graph object now contains all node/edge data.
            # We nullify the DataFrames immediately to free ~50% of peak building memory.
            self.nodes = None
            self.edges = None
            del gdf_nodes
            del gdf_edges
            del df_nodes
            del df_edges
            gc.collect()

        # Calculate edge lengths (crucial for routing/stats)
        self.logger.info("Calculating edge lengths...")
        self.graph = ox.distance.add_edge_lengths(self.graph)

        self.logger.info("Graph construction complete.")

    def _fallback_to_place(self):
        self.logger.info(
            f"Falling back to Overpass API for {self.place} (this may be slow)..."
        )
        self.graph = ox.graph_from_place(self.place, network_type=self.network_type)
        self.logger.info("Converting graph to GeoDataFrames...")
        self.nodes, self.edges = ox.graph_to_gdfs(self.graph)

    def _download_and_save_pbf(self, url: str, dest_path: str):
        try:
            self.logger.info(f"Downloading PBF from {url}...")
            # Stream download to avoid memory spikes, follow redirects for stability
            with httpx.Client(follow_redirects=True) as client:
                with client.stream("GET", url) as r:
                    r.raise_for_status()
                    
                    # Security Check: Enforce a maximum download size (e.g., 500MB)
                    max_bytes = self.kwargs.get("max_pbf_size", 500 * 1024 * 1024)
                    content_length = r.headers.get("Content-Length")
                    if content_length and int(content_length) > max_bytes:
                        raise ValueError(f"PBF file too large: {content_length} bytes (limit: {max_bytes})")

                    # Write directly to FS (upload stream)
                    self.logger.info(f"Uploading PBF to destination: {dest_path}...")
                    downloaded_bytes = 0
                    with self.fs.open(dest_path, "wb") as f:
                        for chunk in r.iter_bytes(chunk_size=8192):
                            downloaded_bytes += len(chunk)
                            if downloaded_bytes > max_bytes:
                                raise ValueError(f"PBF download exceeded limit: {max_bytes} bytes")
                            f.write(chunk)
            self.logger.info(f"Successfully saved PBF to {dest_path}")
        except Exception as e:
            self.logger.error(f"Failed to download PBF from {url}: {e}")

    def save_artifacts(self) -> None:
        """
        Atomic save of graph (pickle) and GDFs (parquet).
        Renamed from save_to_pickle to reflect mixed formats.
        """
        self.logger.info("Saving artifacts to storage...")
        # 1. Save Graph (Pickle)
        if self.graph is not None:
            self.logger.info(f"Saving Graph to {self.graph_path}...")
            self._atomic_write_graphml(self.graph_path, self.graph)

        # 2. Save Nodes/Edges (Parquet)
        if self.nodes is not None:
            self.logger.info(f"Saving Nodes to {self.node_path}...")
            self._atomic_write_parquet(self.node_path, self.nodes)

        if self.edges is not None:
            self.logger.info(f"Saving Edges to {self.edge_path}...")
            self._atomic_write_parquet(self.edge_path, self.edges)
        self.logger.info("Artifacts saved successfully.")

    def load_artifacts(self) -> None:
        """Loads all artifacts from the filesystem into memory."""
        self.graph = self._load_graphml(self.graph_path)
        self.nodes = self._load_parquet(self.node_path)
        self.edges = self._load_parquet(self.edge_path)

    # --- Internal Helpers ---

    def _atomic_write_graphml(self, path: str, G: Any) -> None:
        """
        Safely writes GraphML by first writing to a local temp file,
        uploading to a remote .tmp file, and then renaming.
        Required because ox.save_graphml expects a file path string or local file handler,
        and s3fs/remote file objects aren't always fully supported by networkx/osmnx.
        """
        tmp_dir = self.project_root / ".cache" / "tmp"
        tmp_dir.mkdir(parents=True, exist_ok=True)

        # Create a local temp file to satisfy ox.save_graphml requirement for a path
        with tempfile.NamedTemporaryFile(suffix=".graphml", dir=tmp_dir, delete=False) as local_tmp:
            local_tmp_path = self.get_secure_path(local_tmp.name)

        try:
            # 1. Write to local temp file
            ox.save_graphml(G, local_tmp_path)

            # 2. Upload local temp file to remote storage
            self.fs.put(local_tmp_path, tmp_path)

            # 3. Atomic rename on remote
            self.fs.rename(tmp_path, path)

        except Exception as e:
            if self.fs.exists(tmp_path):
                self.fs.rm(tmp_path)
            raise IOError(f"Failed atomic graphml write to {path}: {e}")
        finally:
            # Cleanup local temp file
            if os.path.exists(local_tmp_path):
                os.remove(local_tmp_path)

    def _atomic_write_parquet(self, path: str, gdf: gpd.GeoDataFrame) -> None:
        tmp_path = f"{path}.tmp"
        try:
            # Geopandas to_parquet handles filesystem URLs if supported,
            # but fsspec support varies. Safe bet is opening file via fs.
            with self.fs.open(tmp_path, "wb") as f:
                gdf.to_parquet(f)
            self.fs.rename(tmp_path, path)
        except Exception as e:
            if self.fs.exists(tmp_path):
                self.fs.rm(tmp_path)
            raise IOError(f"Failed atomic parquet write to {path}: {e}")

    def _load_graphml(self, path: str) -> Any:
        tmp_dir = self.project_root / ".cache" / "tmp"
        tmp_dir.mkdir(parents=True, exist_ok=True)
        # For robustness with remote filesystems, download to local temp first
        with tempfile.NamedTemporaryFile(suffix=".graphml", dir=tmp_dir, delete=False) as local_tmp:
            local_tmp_path = self.get_secure_path(local_tmp.name)

        try:
            self.fs.get(path, local_tmp_path)
            return ox.load_graphml(local_tmp_path)
        finally:
            if os.path.exists(local_tmp_path):
                os.remove(local_tmp_path)

    def _load_parquet(self, path: str) -> gpd.GeoDataFrame:
        with self.fs.open(path, "rb") as f:
            return gpd.read_parquet(f)

    def _artifacts_exist(self) -> bool:
        return all(
            self.fs.exists(p) for p in (self.graph_path, self.node_path, self.edge_path)
        )

    def _delete_artifacts(self) -> None:
        for p in (self.graph_path, self.node_path, self.edge_path):
            if self.fs.exists(p):
                self.fs.rm(p)

    def clear_cache(self) -> None:
        """Clears in-memory data to free resources after persistence."""
        self.logger.info("Clearing PBFHandler in-memory cache...")
        self.graph = None
        self.nodes = None
        self.edges = None
        import gc
        gc.collect()
