from __future__ import annotations

from abc import ABCMeta, abstractmethod
from collections.abc import Callable
from numbers import Number
from pathlib import Path
from typing import Any, cast

import torch
from packaging.version import Version

_torch_version_gt_112 = Version(torch.__version__) > Version("1.12.0")


class ComputationModel(metaclass=ABCMeta):
    """Base class for distributed computation models and defines interface methods.
    This class is public and should be used for other custom derived distributed models.
    """

    # this is an additional local rank storage used when idist is setup from existing native torch dist context
    _ext_local_rank: int | None = None

    def __init__(self) -> None:
        self._backend: str | None = None
        self._nproc_per_node: int | None = None
        self._nnodes: int | None = None
        self._node: int | None = None

    def _setup_attrs(self) -> None:
        if self._nproc_per_node is None:
            self._nproc_per_node = self._compute_nproc_per_node() if self.get_world_size() > 1 else 1
        if self._nnodes is None:
            self._nnodes = self.get_world_size() // self._nproc_per_node
        if self._node is None:
            self._node = self.get_rank() // self._nproc_per_node

    @abstractmethod
    def _compute_nproc_per_node(self) -> int:
        pass

    @abstractmethod
    def get_local_rank(self) -> int:
        pass

    @abstractmethod
    def get_rank(self) -> int:
        pass

    @abstractmethod
    def get_world_size(self) -> int:
        pass

    @abstractmethod
    def get_nproc_per_node(self) -> int:
        pass

    @abstractmethod
    def get_nnodes(self) -> int:
        pass

    @abstractmethod
    def get_node_rank(self) -> int:
        pass

    @abstractmethod
    def device(self) -> torch.device:
        pass

    @abstractmethod
    def backend(self) -> str | None:
        pass

    @abstractmethod
    def finalize(self) -> None:
        pass

    @staticmethod
    @abstractmethod
    def create_from_context() -> ComputationModel | None:
        pass

    @staticmethod
    @abstractmethod
    def create_from_backend(backend: str, **kwargs: Any) -> ComputationModel:
        pass

    @staticmethod
    @abstractmethod
    def spawn(*args: Any, **kwargs: Any) -> None:
        pass

    _collective_op_dtype: Any = None

    @staticmethod
    def _encode_str(x: str, device: torch.device, size: int) -> torch.Tensor:
        name = torch.tensor(bytearray(x, "utf-8")).to(device)
        padded_x = torch.zeros(size + 1, device=device, dtype=torch.long)
        padded_x[: len(name)] = name
        padded_x[-1] = len(name)
        # output is tensor of shape (1, size + 1)
        return padded_x.unsqueeze(0)

    def _get_max_length(self, x: str, device: torch.device) -> int:
        size = torch.tensor([len(x)], device=device)
        size = self._do_all_reduce(size, op="MAX")
        return cast(int, size.item())

    @staticmethod
    def _encode_input_data(x: torch.Tensor | float | str | Path | None, is_src: bool) -> list[int]:
        encoded_msg = [-1] * 512
        if not is_src:
            # Discard input type if not source
            return encoded_msg

        if isinstance(x, torch.Tensor):
            shape = x.shape
            dtype = str(x.dtype)
            msg = [0, len(shape), *shape, len(dtype), *list(bytearray(dtype, "utf-8"))]
            encoded_msg[: len(msg)] = msg
        elif isinstance(x, Number):
            encoded_msg[0] = 1
        elif isinstance(x, str):
            encoded_msg[0] = 2
        elif isinstance(x, Path):
            # Keep a distinct tag so non-source ranks in safe mode can reconstruct
            # Path placeholders and return Path values rather than strings.
            encoded_msg[0] = 3
        return encoded_msg

    @staticmethod
    def _decode_as_placeholder(encoded_msg: list[int], device: torch.device) -> torch.Tensor | float | str | Path:
        if encoded_msg[0] == 0:
            len_shape = encoded_msg[1]
            le = 2 + len_shape
            shape = encoded_msg[2:le] if len_shape > 0 else []
            len_dtype = encoded_msg[le]
            dtype_str = bytearray(encoded_msg[le + 1 : le + 1 + len_dtype]).decode("utf-8")
            dtype = eval(dtype_str)
            return torch.empty(shape, device=device, dtype=dtype)
        elif encoded_msg[0] == 1:
            return 0.0
        elif encoded_msg[0] == 2:
            return ""
        elif encoded_msg[0] == 3:
            return Path("")
        else:
            raise RuntimeError(f"Internal error: unhandled dtype {encoded_msg[0]}, encoded_msg={encoded_msg}")

    def _setup_placeholder(
        self, x: torch.Tensor | float | str | Path | None, device: torch.device, is_src: bool
    ) -> torch.Tensor | float | str | Path:
        encoded_msg_per_rank = self._encode_input_data(x, is_src)
        encoded_msg_all_ranks = self._do_all_reduce(torch.tensor(encoded_msg_per_rank, device=device), op="MAX")

        if is_src:
            if x is None:
                raise RuntimeError("Internal error, x is None. Please, file an issue if you encounter this error.")
            return x

        encoded_msg = encoded_msg_all_ranks.cpu().tolist()
        return self._decode_as_placeholder(encoded_msg, device)

    @staticmethod
    def _decode_str(xs: torch.Tensor) -> list[str]:
        # xs.shape = (n, size + 1), e.g. (world_size, size + 1)
        out = [bytearray(x[: x[-1]].tolist()).decode("utf-8") for x in xs]
        return out

    def _apply_op(
        self, tensor: torch.Tensor, device: torch.device, fn: Callable, *args: Any, **kwargs: Any
    ) -> torch.Tensor:
        out_dtype = None
        tensor_device = None

        # check if the tensor is at specified device
        if tensor.device != device:
            tensor_device = tensor.device
            tensor = tensor.to(device)

        if self._collective_op_dtype is not None:
            # cast to _collective_op_dtype if current type is not floatX
            if tensor.dtype not in (torch.float32, torch.float64):
                out_dtype = tensor.dtype
                tensor = tensor.to(self._collective_op_dtype)

        tensor = fn(tensor, *args, **kwargs)

        if out_dtype is not None and tensor_device is not None:
            return tensor.to(dtype=out_dtype, device=tensor_device)
        if out_dtype is not None:
            return tensor.to(dtype=out_dtype)
        if tensor_device is not None:
            return tensor.to(device=tensor_device)
        return tensor

    def _collective_op(
        self, tensor: torch.Tensor | Number | str, fn: Callable, *args: Any, **kwargs: Any
    ) -> torch.Tensor | float | list[float] | list[str]:
        tensor_to_number = tensor_to_str = False
        device = self.device()
        if isinstance(tensor, (Number, float)):
            tensor_to_number = True
            dtype = self._collective_op_dtype
            if dtype is None and isinstance(tensor, float):
                dtype = torch.double
            tensor = torch.tensor(tensor, device=device, dtype=dtype)
        elif isinstance(tensor, str):
            tensor_to_str = True
            max_length = self._get_max_length(tensor, device)
            tensor = self._encode_str(tensor, device, size=max_length)

        tensor = self._apply_op(tensor, device, fn, *args, **kwargs)

        if tensor_to_number:
            return tensor.tolist()
        elif tensor_to_str:
            return self._decode_str(tensor)
        return tensor

    def all_reduce(
        self, tensor: torch.Tensor | float, op: str = "sum", group: Any | None = None
    ) -> torch.Tensor | float:
        if not isinstance(tensor, (torch.Tensor, Number)):
            raise TypeError(f"Unhandled input type {type(tensor)}")

        return cast(torch.Tensor | float, self._collective_op(tensor, self._do_all_reduce, op, group=group))

    def all_gather(
        self, tensor: torch.Tensor | float | str | Any, group: Any | None = None
    ) -> torch.Tensor | float | list[float] | list[str] | list[Any]:
        if not isinstance(tensor, (torch.Tensor, Number, str)):
            return self._do_all_gather_object(tensor, group=group)

        return self._collective_op(tensor, self._do_all_gather, group=group)

    def new_group(self, ranks: list[int], **kwargs: Any) -> Any:
        if isinstance(ranks, list) and all(isinstance(item, int) for item in ranks):
            return self._do_new_group(ranks, **kwargs)
        else:
            raise ValueError("Argument ranks should be list of int")

    def broadcast(
        self, tensor: torch.Tensor | float | str | Path | None, src: int = 0, safe_mode: bool = False
    ) -> torch.Tensor | float | str | Path:
        if not (isinstance(tensor, (torch.Tensor, Number, str, Path)) or tensor is None):
            raise TypeError(f"Unhandled input type {type(tensor)}")

        rank = self.get_rank()
        if tensor is None:
            if rank == src:
                raise ValueError("Source data can not be None")
            elif not safe_mode:
                raise ValueError("Argument safe_mode should be True if None is passed for non src rank")

        device = self.device()
        tensor_to_number = tensor_to_str = tensor_to_path = False

        if safe_mode:
            tensor = self._setup_placeholder(tensor, device, rank == src)

        if tensor is None:
            raise RuntimeError("Internal error, tensor is None. Please, file an issue if you encounter this error.")

        if isinstance(tensor, (Number, float)):  # have to use Number and float to satisfy mypy
            tensor_to_number = True
            if rank != src:
                tensor = torch.empty(1, device=device, dtype=torch.float)
            else:
                tensor = torch.tensor([tensor], device=device, dtype=torch.float)
        elif isinstance(tensor, (str, Path)):
            tensor_to_str = isinstance(tensor, str)
            tensor_to_path = isinstance(tensor, Path)
            tensor_str = str(tensor)
            max_length = self._get_max_length(tensor_str, device)
            if rank != src:
                tensor = torch.empty(1, max_length + 1, device=device, dtype=torch.long)
            else:
                tensor = self._encode_str(tensor_str, device, size=max_length)

        tensor = self._apply_op(tensor, device, self._do_broadcast, src)

        if tensor_to_number:
            return tensor.item()
        if tensor_to_str or tensor_to_path:
            list_str = self._decode_str(tensor)
            result = list_str[0]
            return Path(result) if tensor_to_path else result
        return tensor

    @abstractmethod
    def _do_all_reduce(self, tensor: torch.Tensor, op: str = "SUM", group: Any | None = None) -> torch.Tensor:
        pass

    @abstractmethod
    def _do_all_gather(self, tensor: torch.Tensor, group: Any | None = None) -> torch.Tensor:
        pass

    @abstractmethod
    def _do_all_gather_object(self, tensor: Any, group: Any | None = None) -> list[Any]:
        pass

    @abstractmethod
    def _do_broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
        pass

    @abstractmethod
    def barrier(self) -> None:
        pass

    @abstractmethod
    def _do_new_group(self, ranks: list[int], **kwargs: Any) -> Any:
        pass

    @abstractmethod
    def _rank_not_in_group(self, group: Any) -> bool:
        pass


class _SerialModel(ComputationModel):
    """Private class defines non-distributed computation model for code compatibility with other distributed models."""

    name = "serial"
    available_backends = ()

    def __init__(self, _backend: str | None = None, **kwargs: Any) -> None:
        super().__init__()

    def get_local_rank(self) -> int:
        return 0

    def get_rank(self) -> int:
        return 0

    def get_world_size(self) -> int:
        return 1

    def get_nproc_per_node(self) -> int:
        return 1

    def get_nnodes(self) -> int:
        return 1

    def get_node_rank(self) -> int:
        return 0

    def device(self) -> torch.device:
        if torch.cuda.is_available():
            return torch.device("cuda")
        if _torch_version_gt_112 and torch.backends.mps.is_available():
            return torch.device("mps")
        return torch.device("cpu")

    def backend(self) -> str | None:
        return None

    def finalize(self) -> None:
        pass

    def _compute_nproc_per_node(self) -> int:
        return 1

    @staticmethod
    def create_from_context() -> _SerialModel:
        return _SerialModel()

    @staticmethod
    def create_from_backend(backend: str | None = None, **kwargs: Any) -> _SerialModel:
        return _SerialModel()

    @staticmethod
    def spawn(*args: Any, **kwargs: Any) -> None:
        raise NotImplementedError("Serial computation model does not implement spawn method")

    def all_reduce(
        self, tensor: torch.Tensor | float, op: str = "SUM", group: Any | None = None
    ) -> torch.Tensor | float:
        return tensor

    def all_gather(
        self, tensor: torch.Tensor | float | str | Any, group: Any | None = None
    ) -> torch.Tensor | float | list[float] | list[str] | list[Any]:
        if isinstance(tensor, torch.Tensor):
            return tensor
        return cast(list[float] | list[str] | list[Any], [tensor])

    def broadcast(
        self, tensor: torch.Tensor | float | str | Path | None, src: int = 0, safe_mode: bool = False
    ) -> torch.Tensor | float | str | Path:
        if tensor is None:
            raise ValueError("Argument tensor should not be None")
        return tensor

    def _do_all_reduce(self, tensor: torch.Tensor, op: str = "SUM", group: Any | None = None) -> torch.Tensor:
        return tensor

    def _do_all_gather(self, tensor: torch.Tensor, group: Any | None = None) -> torch.Tensor:
        return tensor

    def _do_all_gather_object(self, tensor: Any, group: Any | None = None) -> Any:
        return tensor

    def _do_new_group(self, ranks: list[int], **kwargs: Any) -> Any:
        return ranks

    def _do_broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
        return tensor

    def barrier(self) -> None:
        pass

    def new_group(self, ranks: list[int], **kwargs: Any) -> Any:
        if isinstance(ranks, list) and all(isinstance(item, int) for item in ranks):
            return self._do_new_group(ranks, **kwargs)
        else:
            raise ValueError("Argument ranks should be list of int")

    def _rank_not_in_group(self, group: Any) -> bool:
        return False
