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__()
                self.nodes_list = []
                self.edges_list = []

            def way(self, w):
                # Filter for highways (simplified robust filter)
                if "highway" in w.tags:
                    # Basic tag filtering (can be expanded to match osmnx settings)
                    hw = w.tags["highway"]
                    if hw in [
                        "motorway",
                        "trunk",
                        "primary",
                        "secondary",
                        "tertiary",
                        "unclassified",
                        "residential",
                        "service",
                    ]:
                        try:
                            # Extract geometry coordinates
                            coords = [(n.lon, n.lat) for n in w.nodes]
                            if len(coords) < 2:
                                return

                            # Add Edge
                            geom = LineString(coords)
                            # Approximate u, v ID logic (using coords or arbitrary since we reconstruct)
                            # Actually, osmnx relies on node IDs.
                            # w.nodes gives us node IDs!
                            u_id = w.nodes[0].ref
                            v_id = w.nodes[-1].ref

                            edge_data = {
                                "u": u_id,
                                "v": v_id,
                                "key": 0,
                                "osmid": w.id,
                                "highway": hw,
                                "geometry": geom,
                                "oneway": w.tags.get("oneway") == "yes",
                                "length": None,  # osmnx will calculate
                            }
                            # Add extra useful tags
                            for tag in ["name", "maxspeed", "lanes", "ref"]:
                                if tag in w.tags:
                                    edge_data[tag] = w.tags[tag]

                            self.edges_list.append(edge_data)

                            # Add Nodes (all nodes in way)
                            # Note: This list will have duplicates, we dedup later
                            for n in w.nodes:
                                self.nodes_list.append(
                                    {
                                        "osmid": n.ref,
                                        "x": n.lon,
                                        "y": n.lat,
                                        "geometry": Point(n.lon, n.lat),
                                    }
                                )
                        except osmium.InvalidLocationError:
                            # Node missing in PBF extract
                            pass
                        except 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_list)} nodes and {len(handler.edges_list)} edges."
        )

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

        self.logger.info("Building GeoDataFrames...")
        # Dedup nodes
        df_nodes = (
            pd.DataFrame(handler.nodes_list)
            .drop_duplicates(subset=["osmid"])
            .set_index("osmid")
        )
        gdf_nodes = gpd.GeoDataFrame(df_nodes, geometry="geometry", crs="EPSG:4326")

        df_edges = pd.DataFrame(handler.edges_list)
        # Generate keys for parallel edges
        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")
        if not gdf_edges.empty:
            gdf_edges.set_index(["u", "v", "key"], inplace=True)

        # Assign to self properties
        self.nodes = gdf_nodes
        self.edges = gdf_edges

        self.logger.info("Constructing NetworkX graph from GDFs...")
        self.graph = ox.graph_from_gdfs(gdf_nodes, gdf_edges)

        # 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)
