# src/sibi_flux/dask_cluster/client_manager.py
from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import multiprocessing
import os
import shutil
import tempfile
import threading
import time
from contextlib import asynccontextmanager, contextmanager, suppress
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Iterator,
    Optional,
)

import psutil

try:
    from dask.distributed import Client, LocalCluster, get_client

    HAS_DISTRIBUTED = True
except ImportError:
    Client = object
    LocalCluster = object
    get_client = None
    HAS_DISTRIBUTED = False

# Lazy import watchdog to allow remote workers (which might lack it) to import store
try:
    from watchdog.events import FileSystemEventHandler
    from watchdog.observers import Observer

    HAS_WATCHDOG = True
except ImportError:
    Observer = object
    FileSystemEventHandler = object
    HAS_WATCHDOG = False

try:
    from sibi_flux.logger import Logger

    _default_logger = Logger.default_logger(logger_name="dask_lifecycle_manager")
except ImportError:
    _default_logger = logging.getLogger(__name__)  # type: ignore
    if not _default_logger.handlers:
        _default_logger.addHandler(logging.StreamHandler())
    _default_logger.setLevel(logging.INFO)

if TYPE_CHECKING:
    from sibi_flux.logger import Logger as SibiLogger

_cores = multiprocessing.cpu_count()
DEFAULT_WORKERS = max(2, min(4, _cores))
DEFAULT_THREADS = 2


class RegistryEventHandler(FileSystemEventHandler):
    """
    Watchdog handler to trigger logic when the registry file changes.
    Used to detect if the shared cluster configuration has been updated by another process.
    """

    def __init__(self, callback: Callable[[], None]) -> None:
        if not HAS_WATCHDOG:
            return
        self.callback = callback

    def on_modified(self, event: Any) -> None:
        if not event.is_directory and event.src_path.endswith(
            "shared_dask_cluster.json"
        ):
            self.callback()


class DaskClientMixin:
    """
    Stabilised Dask lifecycle manager.

    This class manages a shared Dask LocalCluster instance, enabling:
    - Singleton cluster management across processes via a file-based registry.
    - Automatic healing and reconnection if the cluster becomes unresponsive.
    - Event-driven monitoring using `watchdog` to detect external registry updates.
    """

    REGISTRY_PATH = Path(tempfile.gettempdir()) / "shared_dask_cluster.json"
    REGISTRY_LOCK_PATH = REGISTRY_PATH.parent / (REGISTRY_PATH.name + ".lock")
    WATCHDOG_INTERVAL = 30
    HEAL_COOLDOWN = 5.0

    def __init__(self, **kwargs: Any) -> None:
        self.dask_client: Optional[Client] = None
        self.own_dask_client = False
        self.logger: SibiLogger | logging.Logger = kwargs.get("logger") or _default_logger
        self._watchdog_task: Optional[asyncio.Task[None]] = None
        self._watchdog_stop = asyncio.Event()
        self._fs_observer: Optional[Observer] = None
        self._init_params: dict[str, Any] = {}
        self._last_heal_time = 0.0
        self._last_registry_write = 0.0

    def _get_pkg_hash(self, pkg_path: Path) -> str:
        """Calculates a recursive hash of the package directory."""
        hasher = hashlib.md5()
        for root, _, files in os.walk(pkg_path):
            for file in sorted(files):
                if file.endswith(".pyc") or "__pycache__" in root:
                    continue
                path = Path(root) / file
                hasher.update(path.name.encode())
                try:
                    hasher.update(path.read_bytes())
                except (OSError, IOError):
                    continue
        return hasher.hexdigest()

    def _deploy_local_code(self, client: Client) -> None:
        """
        Zips the local 'sibi_flux' package and uploads it to the scheduler.
        Uses a hash-based check to skip redundant uploads.
        """
        try:
            import sibi_flux

            pkg_path = Path(sibi_flux.__file__).parent
            src_root = pkg_path.parent
            current_hash = self._get_pkg_hash(pkg_path)

            hash_file = Path(tempfile.gettempdir()) / ".sibi_flux_deploy_hash"
            if hash_file.exists() and hash_file.read_text().strip() == current_hash:
                self.logger.info("Local code already deployed; skipping upload.")
                return

            with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
                archive_path = Path(tmp.name)

            shutil.make_archive(
                str(archive_path.with_suffix("")),
                "zip",
                root_dir=str(src_root),
                base_dir="sibi_flux",
            )

            self.logger.info(f"Deploying local code to cluster: {archive_path}")
            client.upload_file(str(archive_path))

            # Update hash
            hash_file.write_text(current_hash)

            # Cleanup
            archive_path.unlink(missing_ok=True)

        except Exception as e:
            self.logger.warning(f"Failed to deploy local code: {e}")

    def _verify_connectivity(self, client: Client) -> bool:
        """
        Verifies that the client can actually run tasks on the cluster.
        """
        try:
            client.wait_for_workers(1, timeout=2.0)

            def environment_probe() -> bool:
                return True

            f = client.submit(environment_probe)
            for attempt in range(3):
                try:
                    if f.result(timeout=5.0) is True:
                        return True
                except Exception:
                    if attempt < 2:
                        # Non-blocking backoff if we could, but here we are usually in init
                        time.sleep(1.0 + attempt)
                        continue
            return False
        except Exception as e:
            self.logger.warning(
                f"Cluster verification failed: {e}",
                extra={"error": str(e)},
            )
            return False

    def _remove_registry(self) -> None:
        with suppress(FileNotFoundError):
            self.REGISTRY_PATH.unlink()

    def _cleanup_stale_registry(self) -> None:
        """Removes the registry file if the owning process is no longer running."""
        reg = self._read_registry()
        if reg and reg.get("pid"):
            if not psutil.pid_exists(reg["pid"]):
                self.logger.warning(
                    "Cleaning zombie registry", extra={"pid": reg["pid"]}
                )
                self._remove_registry()

    def _retire_all_workers(self) -> None:
        """Attempts to gracefully retire all workers."""
        if self.dask_client:
            with suppress(Exception):
                self.dask_client.retire_workers()

    def _has_inflight(self) -> bool:
        return False

    @contextmanager
    def _registry_lock(self, timeout: float = 10.0) -> Iterator[None]:
        """
        Acquires an exclusive file lock to ensure atomic registry updates.
        """
        start_time = time.time()
        while True:
            try:
                self.REGISTRY_LOCK_PATH.touch(exist_ok=False)
                try:
                    yield
                finally:
                    with suppress(OSError):
                        self.REGISTRY_LOCK_PATH.unlink()
                return
            except FileExistsError:
                if time.time() - start_time > timeout:
                    raise TimeoutError("Could not acquire registry lock")
                # Randomized jitter to avoid thundering herd on lock release
                import random
                time.sleep(0.01 + random.random() * 0.05)

    def _read_registry(self) -> Optional[dict[str, Any]]:
        """Reads and parses the registry file JSON."""
        if not self.REGISTRY_PATH.exists():
            return None
        try:
            with self.REGISTRY_PATH.open("r") as f:
                return json.load(f)
        except (json.JSONDecodeError, OSError):
            return None

    def _write_registry(self, data: dict[str, Any]) -> None:
        """
        Writes data to the registry file atomically.
        """
        self._last_registry_write = time.time()
        tmp = self.REGISTRY_PATH.with_suffix(".tmp")
        with tmp.open("w") as f:
            json.dump(data, f)
        tmp.replace(self.REGISTRY_PATH)

    def _init_dask_client(self, **kwargs: Any) -> None:
        self._init_params = kwargs
        if not HAS_DISTRIBUTED:
            self.logger.info("Dask Distributed not installed. Skipping cluster init.")
            return

        if kwargs.get("dask_client"):
            self.dask_client = kwargs["dask_client"]
            return
        if kwargs.get("reuse_context_client", True):
            with suppress(ValueError):
                self.dask_client = get_client()
                return
        self._connect_or_create()
        if kwargs.get("watchdog", True):
            self._start_watchdog()

    def _connect_or_create(self) -> None:
        """
        Connects to an existing Dask cluster or creates a new one.
        """
        p = self._init_params
        try:
            explicit_address = p.get("scheduler_address")
            if explicit_address:
                try:
                    candidate_client = Client(
                        address=explicit_address,
                        timeout=p.get("timeout", 5),
                        set_as_default=p.get("set_as_default", True),
                    )
                    if self._verify_connectivity(candidate_client):
                        self.dask_client = candidate_client
                        self.own_dask_client = False
                        return
                    candidate_client.close()
                except Exception as e:
                    self.logger.error(f"Failed to connect to specified scheduler: {e}")

            env_address = os.environ.get("DASK_SCHEDULER_ADDRESS")
            if env_address:
                try:
                    candidate_client = Client(
                        address=env_address,
                        timeout=p.get("timeout", 5),
                        set_as_default=p.get("set_as_default", True),
                    )
                    if p.get("deploy_mode", False):
                        self._deploy_local_code(candidate_client)

                    if self._verify_connectivity(candidate_client):
                        self.dask_client = candidate_client
                        self.own_dask_client = False
                        return
                    candidate_client.close()
                except Exception as e:
                    self.logger.error(f"Failed to connect to env-specified scheduler: {e}")

            try:
                with self._registry_lock():
                    self._cleanup_stale_registry()
                    reg = self._read_registry()

                    if reg and not reg.get("closing"):
                        try:
                            self.dask_client = Client(
                                address=reg["address"],
                                timeout=p.get("timeout", 30),
                                set_as_default=p.get("set_as_default", True),
                            )
                            reg["refcount"] = int(reg.get("refcount", 0)) + 1
                            self._write_registry(reg)
                            self.own_dask_client = False
                            return
                        except Exception:
                            pass

                    self._bootstrap_new_cluster(p)
            except TimeoutError:
                self._bootstrap_new_cluster(p)

        except Exception as e:
            self.logger.error("Dask initialization failed", extra={"error": str(e)})

    def _bootstrap_new_cluster(self, p: dict[str, Any]) -> None:
        self.logger.info("Starting NEW LOCAL Dask cluster...")
        cluster = LocalCluster(
            n_workers=p.get("n_workers", DEFAULT_WORKERS),
            threads_per_worker=p.get("threads_per_worker", DEFAULT_THREADS),
            processes=p.get("processes", False),
            dashboard_address=p.get("dashboard_address", ":0"),
        )
        self.dask_client = Client(cluster, set_as_default=p.get("set_as_default", True))
        self.own_dask_client = True
        self._write_registry(
            {
                "address": cluster.scheduler_address,
                "dashboard_link": cluster.dashboard_link,
                "refcount": 1,
                "closing": False,
                "pid": os.getpid(),
                "created_at": time.time(),
            }
        )

    def _close_dask_client(self) -> None:
        if self._fs_observer:
            self._fs_observer.stop()
            self._fs_observer.join(timeout=2.0)
            self._fs_observer = None

        if self._watchdog_task:
            self._watchdog_stop.set()

        if not self.dask_client:
            return

        if not self.own_dask_client:
            self.dask_client = None
            return

        try:
            with self._registry_lock():
                if self._has_inflight():
                    return

                reg = self._read_registry()
                if reg:
                    reg["refcount"] = max(0, int(reg.get("refcount", 1)) - 1)
                    if reg["refcount"] > 0:
                        self._write_registry(reg)
                        self.dask_client.close()
                        return
                    reg["closing"] = True
                    self._write_registry(reg)
                    self._retire_all_workers()
                    self._remove_registry()

                cluster = getattr(self.dask_client, "cluster", None)
                self.dask_client.close()
                if cluster:
                    cluster.close()
                self.dask_client = None
        except Exception as e:
            self.logger.error("Shutdown error", extra={"error": str(e)})

    def _start_watchdog(self) -> None:
        if not HAS_WATCHDOG:
            return

        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            return

        def on_fs_event() -> None:
            if time.time() - self._last_registry_write < 1.5:
                return
            loop.call_soon_threadsafe(
                lambda: asyncio.create_task(self._heal_if_disconnected())
            )

        handler = RegistryEventHandler(on_fs_event)
        self._fs_observer = Observer()
        self._fs_observer.schedule(
            handler, str(self.REGISTRY_PATH.parent), recursive=False
        )
        self._fs_observer.start()

        async def liveness_checker() -> None:
            while not self._watchdog_stop.is_set():
                await asyncio.sleep(self.WATCHDOG_INTERVAL)
                await self._heal_if_disconnected()

        self._watchdog_task = loop.create_task(liveness_checker())

    async def _heal_if_disconnected(self) -> None:
        now = time.time()
        if now - self._last_heal_time < self.HEAL_COOLDOWN:
            return
        is_healthy = False
        if self.dask_client:
            try:
                if getattr(self.dask_client, "asynchronous", False):
                    await self.dask_client.scheduler_info()
                else:
                    self.dask_client.scheduler_info()
                is_healthy = True
            except Exception:
                pass
        if not is_healthy:
            self._last_heal_time = now
            self.logger.warning("Dask heartbeat lost. Healing.")
            with suppress(Exception):
                if self.dask_client:
                    await self.dask_client.close(timeout=1)
            self.dask_client = None
            self._connect_or_create()


_persistent_mixin: Optional[DaskClientMixin] = None
_singleton_lock = threading.Lock()


def get_persistent_client(**kwargs: Any) -> Optional[Client]:
    """
    Retrieves or creates a singleton Dask client instance.
    """
    if not HAS_DISTRIBUTED:
        return None
    global _persistent_mixin
    with _singleton_lock:
        if _persistent_mixin is None:
            _persistent_mixin = DaskClientMixin(logger=kwargs.get("logger"))

        if _persistent_mixin.dask_client is None:
            _persistent_mixin._init_dask_client(**kwargs)
        else:
            current_local = _persistent_mixin.own_dask_client
            want_deploy = kwargs.get("deploy_mode", False)
            _persistent_mixin._init_params.update(kwargs)

            if want_deploy and current_local:
                _persistent_mixin._close_dask_client()
                _persistent_mixin._init_dask_client(**kwargs)

    return _persistent_mixin.dask_client


def force_close_persistent_client() -> None:
    global _persistent_mixin
    with _singleton_lock:
        if _persistent_mixin:
            _persistent_mixin._close_dask_client()
            _persistent_mixin = None


def shared_dask_session(*, async_mode: bool = True, **kwargs: Any) -> Any:
    """
    Context manager for a shared Dask session.
    """
    mixin = DaskClientMixin(logger=kwargs.get("logger"))
    mixin._init_dask_client(**kwargs)
    if async_mode:

        @asynccontextmanager
        async def _async_manager() -> Iterator[Optional[Client]]:
            try:
                yield mixin.dask_client
            finally:
                mixin._close_dask_client()

        return _async_manager()
    else:

        @contextmanager
        def _sync_manager() -> Iterator[Optional[Client]]:
            try:
                yield mixin.dask_client
            finally:
                mixin._close_dask_client()

        return _sync_manager()
