"""
SkyTorch PyTorch backend storage manager - Tracks remote storage allocations.

This module manages storage allocations for sky tensors, tracking storage IDs
and their metadata. Storage IDs are used as proxy data pointers in the
allocator, avoiding actual memory allocation on the client side.
"""

from __future__ import annotations

import asyncio
import logging
import weakref
from collections import defaultdict
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Optional
from weakref import WeakValueDictionary

import torch

from skytorch.torch.backend._async import get_event_loop
from skytorch.torch.client.tensor import get_storage_id, get_tensor_id

if TYPE_CHECKING:
    from skytorch.client.compute import Compute

logger = logging.getLogger(__name__)


@dataclass
class StorageInfo:
    """Information about a storage allocation."""

    nbytes: int
    device_index: int
    _compute_ref: weakref.ref[Compute] = field(repr=False)

    @property
    def compute(self) -> Optional[Compute]:
        """Get the associated Compute, or None if garbage collected."""
        return self._compute_ref()


def _delete_tensors_after_gc(
    loop: asyncio.AbstractEventLoop,
    compute: "Compute",
    tensor_ids: list[int],
) -> None:
    """
    Schedule the deletion coroutine on the event loop.

    This function schedules tensor deletion asynchronously rather than blocking.
    This is critical to avoid a deadlock when garbage collection triggers during
    an async operation:

    1. Main thread is in `run_async()` -> `run_until_complete()` -> `_run_once()`
       -> `selector.select()` waiting for I/O
    2. During this wait, garbage collection triggers and finalizes tensor storage
    3. GC calls `free_storage()` which previously called `run_async(delete_tensors(...))`
    4. The nested `run_async()` also tries to enter `selector.select()`
    5. Deadlock: Both calls are blocked on `selector.select()` on the same selector

    By using `call_soon_threadsafe()` + `ensure_future()`, we add the deletion to
    the event loop's callback queue and return immediately. The deletion runs when
    the loop is ready, not during `selector.select()`, avoiding the deadlock.
    """
    from skytorch.torch.backend._client import delete_tensors

    async def do_delete():
        try:
            await delete_tensors(compute, tensor_ids)
        except Exception as e:
            logger.warning(
                f"Failed to delete tensor(s) {tensor_ids} after GC: {e}"
            )

    asyncio.ensure_future(do_delete(), loop=loop)


class StorageManager:
    """
    Manages storage allocations for sky tensors.

    Storage IDs are simple integers that serve as references to remote storage.
    The data pointer in PyTorch tensors is set to the storage ID cast to void*,
    allowing efficient tracking without actual memory allocation.

    IMPORTANT: Storage registration is lazy - storage IDs are generated by the
    C++ allocator using an atomic counter, but registration with StorageManager
    (including Compute association) is deferred until first tensor use. This
    keeps the allocator GIL-free for PyTorch 2.10 compatibility.
    """

    def __init__(self):
        self._storages: dict[int, StorageInfo] = {}
        self._tensor_to_storage: WeakValueDictionary[
            int, torch.UntypedStorage
        ] = WeakValueDictionary()
        self._storage_to_tensors: dict[int, set[int]] = defaultdict(set)

    def register_storage(
        self,
        storage_id: int,
        nbytes: int,
        device_index: int,
    ) -> None:
        """
        Register storage lazily if not already registered.

        This is called at first tensor use to register storage that was
        allocated by the C++ allocator. The allocator generates storage IDs
        without calling Python (GIL-free), so registration is deferred here.

        Args:
            storage_id: Storage ID from the C++ allocator
            nbytes: Size of the storage in bytes
            device_index: Device index for the storage
        """
        if storage_id in self._storages:
            return  # Already registered

        from skytorch.torch.backend._device import device_manager

        compute = device_manager.get_compute(device_index)
        if compute is None:
            raise RuntimeError(
                f"No Compute registered for device index {device_index}. "
                "Ensure you have called compute.device() to register the device."
            )

        info = StorageInfo(nbytes, device_index, _compute_ref=weakref.ref(compute))
        self._storages[storage_id] = info

    def free_storage(self, storage_id: int) -> None:
        """
        Free a storage allocation.

        Handles both registered and unregistered storage IDs.
        If the storage was never used (lazy allocation), this is a no-op.
        Schedules tensor deletion on the Compute's event loop (non-blocking).

        Args:
            storage_id: ID of the storage to free
        """
        tensor_ids = list(self._storage_to_tensors.pop(storage_id, set()))
        info = self._storages.pop(storage_id, None)
        if info is None or not tensor_ids:
            return

        # Get the Compute (may be None if garbage collected)
        compute = info.compute
        if compute is None:
            logger.warning(
                f"Compute was garbage collected, skipping deletion of tensor(s) "
                f"{tensor_ids}"
            )
            return

        # # Get the gRPC client's primary loop (if available)
        # grpc_client = getattr(compute, '_grpc_client', None)
        # if grpc_client is None:
        #     logger.warning(
        #         f"No gRPC client available, skipping deletion of tensor(s) "
        #         f"{tensor_ids}"
        #     )
        #     return
        #
        # loop = getattr(grpc_client, '_primary_loop', None)
        # if loop is None or loop.is_closed():
        #     logger.warning(
        #         f"Event loop is closed or unavailable, skipping deletion of "
        #         f"tensor(s) {tensor_ids}"
        #     )
        #     return

        loop = get_event_loop()

        # Schedule deletion on the primary loop - non-blocking, fire-and-forget
        try:
            loop.call_soon_threadsafe(
                _delete_tensors_after_gc, loop, compute, tensor_ids
            )
        except RuntimeError as e:
            logger.warning(
                f"Failed to schedule deletion for tensor(s) {tensor_ids}: {e}"
            )

    def resize_storage(self, storage_id: int, new_nbytes: int) -> None:
        """
        Resize a storage allocation.

        Args:
            storage_id: ID of the storage to resize
            new_nbytes: New size in bytes
        """
        if storage_id in self._storages:
            self._storages[storage_id].nbytes = new_nbytes

    def get_storage(self, storage_id: int) -> Optional[StorageInfo]:
        """
        Get storage info by ID.

        Args:
            storage_id: ID of the storage

        Returns:
            StorageInfo or None if not found
        """
        return self._storages.get(storage_id)

    def register_tensor(self, tensor: torch.Tensor) -> int:
        """
        Register a tensor with its associated storage.

        Args:
            tensor: The sky tensor to register

        Returns:
            The tensor ID
        """
        tensor_id = get_tensor_id(tensor)
        storage_id = get_storage_id(tensor)
        self._tensor_to_storage[tensor_id] = tensor.untyped_storage()
        self._storage_to_tensors[storage_id].add(tensor_id)
        return tensor_id

    def tensor_ref(self, tensor: torch.Tensor) -> Optional[int]:
        """
        Get a registered tensor ID for this tensor or its storage.

        Args:
            tensor: The tensor to look up

        Returns:
            The tensor's ID if registered, otherwise the first registered
            tensor ID sharing the same storage, otherwise None
        """
        tensor_id = get_tensor_id(tensor)
        if tensor_id in self._tensor_to_storage:
            return tensor_id

        storage_id = get_storage_id(tensor)
        tensor_ids = self._storage_to_tensors.get(storage_id)
        if tensor_ids:
            return next(iter(tensor_ids))

        return None


# Global storage manager singleton
storage_manager = StorageManager()
