"""
OSMnx Graph Loader module.

Responsible for loading and caching OSMnx processing artifacts.
"""

from __future__ import annotations

import logging
import osmnx as ox
import threading
import time
import tempfile
from collections.abc import Mapping
from pathlib import Path
from typing import Any, Tuple

from filelock import FileLock
import geopandas as gpd

from sibi_flux.core import SecureResource
from ._pbf_handler import PBFHandler


class GraphLoader(SecureResource):
    """
    A generic loader for OSMnx graph artifacts that handles:
    1. checking for artifacts in a remote/shared filesystem (source).
    2. copying them to a local cache directory (with file locking).
    3. building them from PBF if missing (via PBFHandler).
    4. lazy loading into memory (singleton pattern per instance).
    """

    def __init__(self, config: Mapping[str, Any], fs: Any):
        """
        Args:
            config: Configuration dictionary containing keys like 'lock_dir', 'files_prefix', etc.
                    Must be compatible with PBFHandler options.
            fs: The filesystem instance used to check/retrieve remote artifacts.
        """
        # Call SecureResource init to setup logger, fs, and project_root
        super().__init__(fs=fs, **config)
        self.config = config

        # Initialize handler to resolve paths
        _handler_config = dict(config)
        _handler_config["fs"] = self.fs
        self.handler = PBFHandler(**_handler_config)

        # Source paths (from handler)
        self.source_graph_path = self.handler.graph_path
        self.source_nodes_path = self.handler.node_path
        self.source_edges_path = self.handler.edge_path
        self.base_path = self.handler.base

        # Local cache configuration - HARDENED with sandboxing
        default_lock_dir = self.project_root / ".cache" / "osmnx-locks"
        self.lock_dir = self.get_secure_path(config.get("lock_dir", default_lock_dir))
        self.lock_dir.mkdir(parents=True, exist_ok=True)

        _prefix = config.get("files_prefix", "osmnx-data").rstrip("-")
        self.lock_file = self.lock_dir / f".osmnx-{_prefix}.lock"
        self.lock_timeout = int(config.get("lock_timeout", "900"))
        self.lock_max_age = int(config.get("lock_max_age", "3600"))

        # Local artifact paths
        # We assume source paths have filenames we can extract.
        # Since source_graph_path is a string (possibly URL), we use Path(str).name logic
        # or simple split, but fsspec paths are usually forward-slash separated.
        self.local_graph_path = self.lock_dir / self.source_graph_path.split("/")[-1]
        self.local_nodes_path = self.lock_dir / self.source_nodes_path.split("/")[-1]
        self.local_edges_path = self.lock_dir / self.source_edges_path.split("/")[-1]

        # In-memory state
        self._G: Any | None = None
        self._g_nodes: Any | None = None
        self._g_edges: Any | None = None
        self._initialized = False
        self._init_mutex = threading.Lock()

    def _source_artifacts_exist(self) -> bool:
        """Checks if artifacts exist in the primary shared storage (FS)."""
        return (
            self.fs.exists(self.source_graph_path)
            and self.fs.exists(self.source_nodes_path)
            and self.fs.exists(self.source_edges_path)
        )

    def _local_artifacts_exist(self) -> bool:
        """Checks if artifacts exist in the local cache directory."""
        return (
            self.local_graph_path.exists()
            and self.local_nodes_path.exists()
            and self.local_edges_path.exists()
        )

    def _delete_local_artifacts(self) -> None:
        """Removes local artifacts to force a re-fetch."""
        for p in [self.local_graph_path, self.local_nodes_path, self.local_edges_path]:
            if p.exists():
                p.unlink()

    def _delete_source_artifacts(self) -> None:
        """Removes remote/source artifacts to force a rebuild."""
        try:
            if self.fs.exists(self.source_graph_path):
                self.fs.rm(self.source_graph_path)
            if self.fs.exists(self.source_nodes_path):
                self.fs.rm(self.source_nodes_path)
            if self.fs.exists(self.source_edges_path):
                self.fs.rm(self.source_edges_path)
        except Exception as e:
            self.logger.error(f"Failed to delete source artifacts: {e}")

    def _load_artifact(self, path: Path) -> Any:
        """Helper to load a single artifact based on extension."""
        if path.suffix == ".parquet":
            return gpd.read_parquet(path)
        if path.suffix in [".graphml", ".xml"]:
            return ox.load_graphml(path)

        raise ValueError(
            f"Unsupported artifact extension: {path.suffix}. Pickle support has been removed for security."
        )

    def _load_local_artifacts(self) -> Tuple[Any, Any, Any]:
        """Loads artifacts from the local cache directory."""
        self.logger.info(f"Loading artifacts from local cache: {self.lock_dir}")
        G = self._load_artifact(self.local_graph_path)
        g_nodes = self._load_artifact(self.local_nodes_path)
        g_edges = self._load_artifact(self.local_edges_path)
        return G, g_nodes, g_edges

    def _clean_stale_lock(self) -> None:
        try:
            if not self.lock_file.exists():
                return
            if time.time() - self.lock_file.stat().st_mtime > self.lock_max_age:
                self.lock_file.unlink()
        except OSError:
            pass

    def _ensure_local_artifacts(self) -> None:
        """Ensures artifacts are present in the local cache, building/copying if necessary."""
        if self._local_artifacts_exist():
            return

        # FileLock expects string or Path (recent versions support Path, but safe to cast)
        with FileLock(str(self.lock_file), timeout=self.lock_timeout):
            if self._local_artifacts_exist():
                return

            self.logger.info("Local artifacts not found. Checking source location...")
            if not self._source_artifacts_exist():
                self.logger.info("Source artifacts not found. Triggering build...")
                # Trigger build via handler
                self.handler.process_pbf()
                self.handler.save_artifacts()
                self.logger.info("OSMnx data built successfully!")

                if not self._source_artifacts_exist():
                    raise RuntimeError(
                        "OSMnx artifacts were not created by PBFHandler."
                    )

            self.logger.info(
                f"Copying artifacts from {self.base_path} to local cache {self.lock_dir}..."
            )
            # fs.get usually supports Path objects if mapped to local file system,
            # but providing string path is safest for fsspec API.
            self.fs.get(self.source_graph_path, str(self.local_graph_path))
            self.fs.get(self.source_nodes_path, str(self.local_nodes_path))
            self.fs.get(self.source_edges_path, str(self.local_edges_path))

    def ensure_initialized(self) -> None:
        """Public function to explicitly trigger initialization."""
        if self._initialized:
            return

        with self._init_mutex:
            if self._initialized:
                return

            self.logger.info("Initialising OSMnx graph data...")
            self._clean_stale_lock()
            self._ensure_local_artifacts()
            try:
                self._G, self._g_nodes, self._g_edges = self._load_local_artifacts()
            except ValueError as e:
                self.logger.warning(
                    f"Corrupted artifacts detected ({e}). nuking and rebuilding..."
                )
                self._delete_local_artifacts()
                self._delete_source_artifacts()
                # Clear handler build content just in case
                self.handler.graph = None
                self.handler.nodes = None
                self.handler.edges = None

                self._ensure_local_artifacts()  # retry (will trigger build)
                self._G, self._g_nodes, self._g_edges = self._load_local_artifacts()
            self._initialized = True
            self.logger.info("OSMnx graph data initialised.")

    def get_graph_triplet(self) -> Tuple[Any, Any, Any]:
        """Public function to get all graph data; triggers initialization."""
        self.ensure_initialized()
        return self._G, self._g_nodes, self._g_edges

    @property
    def G(self):
        self.ensure_initialized()
        return self._G

    @property
    def nodes(self):
        self.ensure_initialized()
        return self._g_nodes

    @property
    def edges(self):
        self.ensure_initialized()
        return self._g_edges
