import logging
import os
from typing import Any, Dict, List, Optional, Tuple, Union

import torch

logger = logging.getLogger(__name__)


class NixlBackendSelection:
    """Handles NIXL backend selection and creation."""

    # Priority order for File-based plugins in case of auto selection
    FILE_PLUGINS = ["3FS", "POSIX", "GDS_MT", "GDS"]
    # Priority order for File-based plugins in case of auto selection (add more as needed)
    OBJ_PLUGINS = ["OBJ"]  # Based on Amazon S3 SDK

    def __init__(self, plugin: str = "auto"):
        """Initialize backend selection.
        Args:
            plugin: Plugin to use (default "auto" selects best available).
                   Can be a file plugin (3FS, POSIX, GDS, GDS_MT) or
                   an object plugin (OBJ).
        """
        self.plugin = plugin
        self.backend_name = None
        self.mem_type = None

    def set_bucket(self, bucket_name: str) -> None:
        """Set AWS bucket name in environment variable."""
        os.environ["AWS_DEFAULT_BUCKET"] = bucket_name
        logger.debug(f"Set AWS bucket name to: {bucket_name}")

    def create_backend(self, agent) -> bool:
        """Create the appropriate NIXL backend based on configuration."""
        try:
            plugin_list = agent.get_plugin_list()
            logger.debug(f"Available NIXL plugins: {plugin_list}")

            # Handle explicit plugin selection or auto priority
            if self.plugin == "auto":
                # Try all file plugins first
                for plugin in self.FILE_PLUGINS:
                    if plugin in plugin_list:
                        self.backend_name = plugin
                        break
                # If no file plugin found, try object plugins
                if not self.backend_name:
                    for plugin in self.OBJ_PLUGINS:
                        if plugin in plugin_list:
                            self.backend_name = plugin
                            break
            else:
                # Use explicitly requested plugin
                self.backend_name = self.plugin

            if self.backend_name not in plugin_list:
                logger.error(
                    f"Backend {self.backend_name} not available in plugins: {plugin_list}"
                )
                return False

            # Create backend and set memory type
            if self.backend_name in self.OBJ_PLUGINS:
                bucket = os.environ.get("AWS_DEFAULT_BUCKET")
                if not bucket:
                    logger.error(
                        "AWS_DEFAULT_BUCKET environment variable must be set for object storage"
                    )
                    return False
                agent.create_backend(self.backend_name, {"bucket": bucket})
            else:
                agent.create_backend(self.backend_name)

            self.mem_type = "OBJ" if self.backend_name in self.OBJ_PLUGINS else "FILE"
            logger.debug(
                f"Created NIXL backend: {self.backend_name} with memory type: {self.mem_type}"
            )
            return True

        except Exception as e:
            logger.error(f"Failed to create NIXL backend: {e}")
            return False


class NixlRegistration:
    """Handles NIXL memory registration."""

    def __init__(self, agent):
        self.agent = agent

    def create_query_tuples(
        self, key: str, mem_type: str, file_manager=None
    ) -> List[Tuple]:
        """Create NIXL tuples for querying memory.
        Args:
            key: Key to query (file path for FILE or object key for OBJ)
            mem_type: Memory type ("FILE" or "OBJ")
            file_manager: Optional NixlFileManager for FILE memory type
        Returns:
            List of NIXL tuples for querying
        """
        if mem_type == "FILE":
            if file_manager is None:
                logger.error("file_manager required for FILE memory type")
                return []
            return [(0, 0, 0, file_manager.get_file_path(key))]
        else:  # OBJ
            return [(0, 0, key)]

    def _register_memory(
        self, items: Union[List[tuple], List[torch.Tensor]], mem_type: str, desc: str
    ) -> Optional[Any]:
        """Common registration logic for files, objects, and buffers.
        Args:
            items: List of tuples or tensors to register
            mem_type: Memory type ("FILE", "OBJ", "DRAM", "VRAM")
            desc: Description for logging
        """
        try:
            if not items:
                return None

            reg_descs = self.agent.get_reg_descs(items, mem_type)
            if reg_descs is None:
                logger.error("Failed to create registration descriptors")
                return None

            registered_memory = self.agent.register_memory(reg_descs)
            if registered_memory:
                return registered_memory
            else:
                logger.error("Failed to register with NIXL")
                return None

        except Exception as e:
            logger.error(f"Failed to register {desc}: {e}")
            return None

    def register_buffers(
        self, buffers: Union[torch.Tensor, List[torch.Tensor]]
    ) -> Optional[Any]:
        """Register tensors/buffers with NIXL."""
        if isinstance(buffers, torch.Tensor):
            buffers = [buffers]

        if not buffers:
            return None

        # Determine memory type based on tensor device
        mem_type = "VRAM" if buffers[0].device.type == "cuda" else "DRAM"
        return self._register_memory(buffers, mem_type, "buffers")

    def register_files(self, tuples: List[tuple]) -> Optional[Any]:
        """Register files with NIXL using (0, 0, fd, file_path) tuples."""
        return self._register_memory(tuples, "FILE", "files")

    def register_objects(
        self, keys: List[str], tensors: Optional[List[torch.Tensor]] = None
    ) -> Optional[Any]:
        """Register objects with NIXL."""
        if not keys:
            return None

        # Create object tuples with proper sizes
        tuples = [
            (0, tensor.element_size() * tensor.numel() if tensor else 0, key)
            for key, tensor in zip(keys, tensors or [None] * len(keys))
        ]
        return self._register_memory(tuples, "OBJ", "objects")


class NixlFileManager:
    """Handles file system operations for NIXL."""

    def __init__(self, base_dir: str):
        """
        Initialize file manager.
        Args:
            base_dir: Base directory for storing tensor files
        """
        self.base_dir = base_dir
        if base_dir == "":
            logger.debug(f"Initialized file manager without a base directory")
        else:
            os.makedirs(base_dir, exist_ok=True)
            logger.debug(f"Initialized file manager with base directory: {base_dir}")

    def get_file_path(self, key: str) -> str:
        """Get full file path for a given key."""
        return os.path.join(self.base_dir, key)

    def create_file(self, file_path: str) -> bool:
        """Create a file if it doesn't exist."""
        try:
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
            if not os.path.exists(file_path):
                with open(file_path, "wb") as f:
                    pass  # Create empty file
            return True
        except Exception as e:
            logger.error(f"Failed to create file {file_path}: {e}")
            return False

    def open_file(self, file_path: str) -> Optional[int]:
        """Open a file and return its file descriptor."""
        try:
            fd = os.open(file_path, os.O_RDWR)
            return fd
        except Exception as e:
            logger.error(f"Failed to open file {file_path}: {e}")
            return None

    def close_file(self, fd: int) -> bool:
        """Close a file descriptor."""
        try:
            os.close(fd)
            return True
        except Exception as e:
            logger.error(f"Failed to close file descriptor {fd}: {e}")
            return False

    def files_to_nixl_tuples(
        self, file_paths: List[str], open_file: bool = True
    ) -> List[Tuple[int, int, int, str]]:
        """Create NIXL tuples (offset, length, fd, file_path) for given files."""
        if not open_file:
            return [(0, 0, 0, path) for path in file_paths]

        tuples = []
        for path in file_paths:
            if (fd := self.open_file(path)) is None:
                # Clean up on failure
                for t in tuples:
                    self.close_file(t[2])
                return []
            tuples.append((0, 0, fd, path))
        return tuples
