"""
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 logging
import threading
import weakref
import collections
from collections import defaultdict
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Optional
from weakref import WeakValueDictionary

import torch

from skytorch.torch.client.loop 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__)


def _delete_tensors_after_gc(compute, tensor_ids, nbytes=0):
    try:
        client = compute._grpc_client
        if client is not None and client.stream is not None:
            # Route through _mt_ops for FIFO ordering with ATen ops.
            # Pending ATen ops in _mt_ops must reach the server before
            # the delete message; _drain_mt_ops processes them in order
            # and _flush_deferred flushes batch buffers before sending deletes.
            stream = client.stream
            with stream._mt_lock:
                stream._mt_ops.append(("ff", tensor_ids))
                if not stream._mt_wake_pending:
                    stream._mt_wake_pending = True
                    stream._loop.call_soon_threadsafe(stream._drain_mt_ops)
    except Exception as e:
        logger.warning(f"Failed to defer deletion of tensor(s) {tensor_ids}: {e}")


@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()


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._lock = threading.RLock()
        self._storages: dict[int, StorageInfo] = {}
        self._tensor_to_storage: WeakValueDictionary[int, torch.UntypedStorage] = (
            WeakValueDictionary()
        )
        self._storage_to_tensors: dict[int, set[int]] = defaultdict(set)
        # Iterative free queue: prevents stack overflow when gc.collect()
        # triggers a chain of tensor deallocations, each calling free_storage.
        self._pending_frees: collections.deque[int] = collections.deque()
        self._processing_frees: bool = False
        # Monotonic byte counters for memory-pressure GC trigger
        self._bytes_registered: int = 0
        self._bytes_freed: int = 0

    @property
    def live_storage_bytes(self) -> int:
        return self._bytes_registered - self._bytes_freed

    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
        """
        with self._lock:
            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
            self._bytes_registered += nbytes

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

        Uses an iterative queue to prevent stack overflow: gc.collect()
        can trigger a chain of tensor deallocations where each calls
        free_storage, growing the stack until RecursionError.  By queuing
        the storage_id and processing iteratively, the stack depth stays
        constant regardless of how many storages are freed at once.

        Args:
            storage_id: ID of the storage to free
        """
        self._pending_frees.append(storage_id)
        if self._processing_frees:
            return  # Already in the processing loop
        self._processing_frees = True
        # Batch deletes per Compute — one call_soon_threadsafe per Compute
        # instead of one per freed storage.  Critical when gc.collect()
        # frees hundreds of storages at once.
        batched: dict[int, tuple] = {}  # id(compute) → (compute, [tensor_ids])
        try:
            while self._pending_frees:
                sid = self._pending_frees.popleft()
                self._free_storage_one(sid, batched)
        finally:
            self._processing_frees = False
        if batched:
            loop = get_event_loop()
            for compute, tensor_ids in batched.values():
                try:
                    loop.call_soon_threadsafe(
                        _delete_tensors_after_gc, compute, tensor_ids
                    )
                except RuntimeError:
                    pass

    def _free_storage_one(self, storage_id: int, batched: dict) -> None:
        """Process a single storage free (called iteratively, not recursively)."""
        with self._lock:
            tensor_ids = list(self._storage_to_tensors.pop(storage_id, set()))
            info = self._storages.pop(storage_id, None)
            if info is not None:
                self._bytes_freed += info.nbytes

        # Unregister all tensor_ids from C++ tracking
        if tensor_ids:
            try:
                from skytorch.torch.backend import _C

                for tid in tensor_ids:
                    _C._unregister_tensor_id(tid)
                _C._unregister_storage_tensor_mapping(storage_id)
            except (ImportError, AttributeError):
                pass

        if info is None or not tensor_ids:
            return

        compute = info.compute
        if compute is None:
            return

        cid = id(compute)
        if cid in batched:
            batched[cid][1].extend(tensor_ids)
        else:
            batched[cid] = (compute, list(tensor_ids))

    def drain_compute(self, compute: "Compute") -> list[int]:
        """Remove all storage entries for a Compute and return their tensor IDs."""
        compute_id = id(compute)
        tensor_ids = []

        with self._lock:
            storage_ids_to_remove = [
                sid
                for sid, info in self._storages.items()
                if info.compute is not None and id(info.compute) == compute_id
            ]

            for sid in storage_ids_to_remove:
                tensor_ids.extend(self._storage_to_tensors.pop(sid, set()))
                self._storages.pop(sid, None)

        return tensor_ids

    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
        """
        with self._lock:
            if storage_id in self._storages:
                old_nbytes = self._storages[storage_id].nbytes
                self._storages[storage_id].nbytes = new_nbytes
                self._bytes_registered += new_nbytes - old_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
        """
        with self._lock:
            return self._storages.get(storage_id)

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

        Args:
            tensor: The sky tensor to register
            skip_cpp: If True, skip C++ fast-path registration. Used by callers
                that need to defer C++ registration until after submission to
                avoid marking tensors as "known" before the server has them.

        Returns:
            The tensor ID
        """
        tensor_id = get_tensor_id(tensor)
        storage_id = get_storage_id(tensor)
        with self._lock:
            self._tensor_to_storage[tensor_id] = tensor.untyped_storage()
            self._storage_to_tensors[storage_id].add(tensor_id)
        if not skip_cpp:
            # Sync to C++ registration set for fast path lookup
            try:
                from skytorch.torch.backend import _C

                _C._register_tensor_id(tensor_id)
                # Also register storage_id → tensor_id mapping for view detection
                _C._register_storage_tensor_mapping(storage_id, tensor_id)
            except (ImportError, AttributeError):
                pass
        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)
        with self._lock:
            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()
