from __future__ import annotations

import multiprocessing as mp
import time
from typing import TYPE_CHECKING, cast

import lazy_loader as lazy
from rich.console import Console

from dimos.core.core import rpc
from dimos.utils.logging_config import setup_logger

if TYPE_CHECKING:
    # Avoid runtime import to prevent circular import; ruff's TC001 would otherwise move it.
    from dask.distributed import LocalCluster

    from dimos.core._dask_exports import DimosCluster
    from dimos.core.module import Module
    from dimos.core.rpc_client import ModuleProxy

logger = setup_logger()

__getattr__, __dir__, __all__ = lazy.attach(
    __name__,
    submodules=["colors"],
    submod_attrs={
        "blueprints": ["autoconnect", "Blueprint"],
        "_dask_exports": ["DimosCluster"],
        "_protocol_exports": ["LCMRPC", "RPCSpec", "LCMTF", "TF", "PubSubTF", "TFConfig", "TFSpec"],
        "module": ["Module", "ModuleBase", "ModuleConfig", "ModuleConfigT"],
        "stream": ["In", "Out", "RemoteIn", "RemoteOut", "Transport"],
        "transport": [
            "LCMTransport",
            "SHMTransport",
            "ZenohTransport",
            "pLCMTransport",
            "pSHMTransport",
        ],
    },
)
__all__ += ["DimosCluster", "Module", "rpc", "start", "wait_exit"]


class CudaCleanupPlugin:
    """Dask worker plugin to cleanup CUDA resources on shutdown."""

    def setup(self, worker) -> None:  # type: ignore[no-untyped-def]
        """Called when worker starts."""
        pass

    def teardown(self, worker) -> None:  # type: ignore[no-untyped-def]
        """Clean up CUDA resources when worker shuts down."""
        try:
            import sys

            if "cupy" in sys.modules:
                import cupy as cp  # type: ignore[import-not-found, import-untyped]

                # Clear memory pools
                mempool = cp.get_default_memory_pool()
                pinned_mempool = cp.get_default_pinned_memory_pool()
                mempool.free_all_blocks()
                pinned_mempool.free_all_blocks()
                cp.cuda.Stream.null.synchronize()
                mempool.free_all_blocks()
                pinned_mempool.free_all_blocks()
        except Exception:
            pass


def patch_actor(actor, cls) -> None: ...  # type: ignore[no-untyped-def]


def patchdask(dask_client: DimosCluster, local_cluster: LocalCluster) -> DimosCluster:
    from dimos.core.rpc_client import RPCClient
    from dimos.utils.actor_registry import ActorRegistry

    def deploy(  # type: ignore[no-untyped-def]
        actor_class: type[Module],
        *args,
        **kwargs,
    ) -> ModuleProxy:
        from dimos.core.docker_runner import DockerModule, is_docker_module

        # Check if this module should run in Docker (based on its default_config)
        if is_docker_module(actor_class):
            logger.info("Deploying module in Docker.", module=actor_class.__name__)
            dm = DockerModule(actor_class, *args, **kwargs)
            dm.start()  # Explicit start - follows create -> configure -> start lifecycle
            dask_client._docker_modules.append(dm)  # type: ignore[attr-defined]
            return dm  # type: ignore[return-value]

        logger.info("Deploying module.", module=actor_class.__name__)
        actor = dask_client.submit(  # type: ignore[no-untyped-call]
            actor_class,
            *args,
            **kwargs,
            actor=True,
        ).result()

        worker = actor.set_ref(actor).result()
        logger.info("Deployed module.", module=actor._cls.__name__, worker_id=worker)

        # Register actor deployment in shared memory
        ActorRegistry.update(str(actor), str(worker))

        return cast("ModuleProxy", RPCClient(actor, actor_class))

    def check_worker_memory() -> None:
        """Check memory usage of all workers."""
        info = dask_client.scheduler_info()

        console = Console()
        total_workers = len(info.get("workers", {}))
        total_memory_used = 0
        total_memory_limit = 0

        for worker_addr, worker_info in info.get("workers", {}).items():
            metrics = worker_info.get("metrics", {})
            memory_used = metrics.get("memory", 0)
            memory_limit = worker_info.get("memory_limit", 0)

            cpu_percent = metrics.get("cpu", 0)
            managed_bytes = metrics.get("managed_bytes", 0)
            spilled = metrics.get("spilled_bytes", {}).get("memory", 0)
            worker_status = worker_info.get("status", "unknown")
            worker_id = worker_info.get("id", "?")

            memory_used_gb = memory_used / 1e9
            memory_limit_gb = memory_limit / 1e9
            managed_gb = managed_bytes / 1e9
            spilled / 1e9

            total_memory_used += memory_used
            total_memory_limit += memory_limit

            percentage = (memory_used_gb / memory_limit_gb * 100) if memory_limit_gb > 0 else 0

            if worker_status == "paused":
                status = "[red]PAUSED"
            elif percentage >= 95:
                status = "[red]CRITICAL"
            elif percentage >= 80:
                status = "[yellow]WARNING"
            else:
                status = "[green]OK"

            console.print(
                f"Worker-{worker_id} {worker_addr}: "
                f"{memory_used_gb:.2f}/{memory_limit_gb:.2f}GB ({percentage:.1f}%) "
                f"CPU:{cpu_percent:.0f}% Managed:{managed_gb:.2f}GB "
                f"{status}"
            )

        if total_workers > 0:
            total_used_gb = total_memory_used / 1e9
            total_limit_gb = total_memory_limit / 1e9
            total_percentage = (total_used_gb / total_limit_gb * 100) if total_limit_gb > 0 else 0
            console.print(
                f"[bold]Total: {total_used_gb:.2f}/{total_limit_gb:.2f}GB ({total_percentage:.1f}%) across {total_workers} workers[/bold]"
            )

    def close_all() -> None:
        # Prevents multiple calls to close_all
        if hasattr(dask_client, "_closed") and dask_client._closed:
            return
        dask_client._closed = True  # type: ignore[attr-defined]

        # Stop all Docker modules (in reverse order of deployment)
        for dm in reversed(dask_client._docker_modules):  # type: ignore[attr-defined]
            try:
                dm.stop()
            except Exception:
                pass
        dask_client._docker_modules.clear()  # type: ignore[attr-defined]

        # Stop all SharedMemory transports before closing Dask
        # This prevents the "leaked shared_memory objects" warning and hangs
        try:
            import gc

            from dimos.protocol.pubsub.impl import shmpubsub

            for obj in gc.get_objects():
                if isinstance(obj, shmpubsub.SharedMemoryPubSubBase):
                    try:
                        obj.stop()
                    except Exception:
                        pass
        except Exception:
            pass

        # Get the event loop before shutting down
        loop = dask_client.loop

        # Clear the actor registry
        ActorRegistry.clear()

        # Close cluster and client with reasonable timeout
        # The CudaCleanupPlugin will handle CUDA cleanup on each worker
        try:
            local_cluster.close(timeout=5)
        except Exception:
            pass

        try:
            dask_client.close(timeout=5)  # type: ignore[no-untyped-call]
        except Exception:
            pass

        if loop and hasattr(loop, "add_callback") and hasattr(loop, "stop"):
            try:
                loop.add_callback(loop.stop)
            except Exception:
                pass

        # Note: We do NOT shutdown the _offload_executor here because it's a global
        # module-level ThreadPoolExecutor shared across all Dask clients in the process.
        # Shutting it down here would break subsequent Dask client usage (e.g., in tests).
        # The executor will be cleaned up when the Python process exits.

        # Give threads time to clean up
        # Dask's IO loop and Profile threads are daemon threads
        # that will be cleaned up when the process exits
        # This is needed, solves race condition in CI thread check
        time.sleep(0.1)

    dask_client._docker_modules = []  # type: ignore[attr-defined]
    dask_client.deploy = deploy  # type: ignore[attr-defined]
    dask_client.check_worker_memory = check_worker_memory  # type: ignore[attr-defined]
    dask_client.stop = lambda: dask_client.close()  # type: ignore[attr-defined, no-untyped-call]
    dask_client.close_all = close_all  # type: ignore[attr-defined]
    return dask_client  # type: ignore[return-value]


def start(n: int | None = None, memory_limit: str = "auto") -> DimosCluster:
    """Start a Dask LocalCluster with specified workers and memory limits.

    Args:
        n: Number of workers (defaults to CPU count)
        memory_limit: Memory limit per worker (e.g., '4GB', '2GiB', or 'auto' for Dask's default)

    Returns:
        DimosCluster: A patched Dask client with deploy(), check_worker_memory(), stop(), and close_all() methods
    """

    from dask.distributed import Client, LocalCluster

    console = Console()
    if not n:
        n = mp.cpu_count()
    with console.status(
        f"[green]Initializing dimos local cluster with [bright_blue]{n} workers", spinner="arc"
    ):
        cluster = LocalCluster(  # type: ignore[no-untyped-call]
            n_workers=n,
            threads_per_worker=4,
            memory_limit=memory_limit,
            plugins=[CudaCleanupPlugin()],  # Register CUDA cleanup plugin
        )
        client = Client(cluster)  # type: ignore[no-untyped-call]

    console.print(
        f"[green]Initialized dimos local cluster with [bright_blue]{n} workers, memory limit: {memory_limit}"
    )

    patched_client = patchdask(client, cluster)

    return patched_client


def wait_exit() -> None:
    while True:
        try:
            time.sleep(1)
        except KeyboardInterrupt:
            print("exiting...")
            return
