from typing import TYPE_CHECKING, Literal, Union

from ._base import TensorLike, TensorLikeBackend

if TYPE_CHECKING:
    from numpy.typing import NDArray
    from torch import Tensor

    from ._numpy import NumpyBackend
    from ._torch import TorchBackend

_NUMPY_BACKEND_INSTANCE: 'NumpyBackend | None' = None
_TORCH_BACKEND_INSTANCE: 'TorchBackend | None' = None


def _get_backend_instance(backend: Literal['numpy', 'torch']) -> TensorLikeBackend:
    """获取后端实例"""
    from ._numpy import NumpyBackend
    from ._torch import TorchBackend

    if backend == 'numpy':
        global _NUMPY_BACKEND_INSTANCE
        if _NUMPY_BACKEND_INSTANCE is None:
            _NUMPY_BACKEND_INSTANCE = NumpyBackend()
        return _NUMPY_BACKEND_INSTANCE
    elif backend == 'torch':
        global _TORCH_BACKEND_INSTANCE
        if _TORCH_BACKEND_INSTANCE is None:
            _TORCH_BACKEND_INSTANCE = TorchBackend()
        return _TORCH_BACKEND_INSTANCE
    else:
        raise ValueError(f'Unsupported backend: {backend}')


def get_tensorlike_backend(
    source: Union[str, 'NDArray', 'Tensor', TensorLikeBackend],
) -> TensorLikeBackend:
    """获取合适的张量数据的处理后端

    该函数作为工厂，返回一个单例的后端实例，用于处理给定的数据类型。
    使用的时候传入 'numpy' 或 'torch' 作为字符串，或者直接传入 numpy 数组或 torch 张量。

    Parameters
    ----------
    - `data_or_name` : 传入 'numpy' 或 'torch' 作为字符串，或者直接传入 numpy 数组或 torch 张量。

    Returns
    ----------
    - `TensorLikeBackend`
        - 返回一个单例的后端实例，用于处理给定的数据类型。

    Raises
    ----------
    - `ImportError`
        - _description_
    - `ValueError`
        - _description_
    - `ValueError`
        - _description_
    """
    import numpy as np
    import torch

    if isinstance(source, str):
        name = source.lower()
        if name == 'numpy':
            return _get_backend_instance('numpy')
        elif name == 'torch' or name == 'pytorch':
            return _get_backend_instance('torch')
        else:
            raise ValueError(f'Unsupported backend: {source}')

    if isinstance(source, np.ndarray):
        return _get_backend_instance('numpy')

    if isinstance(source, (list, tuple)):
        return _get_backend_instance('numpy')

    if isinstance(source, torch.Tensor):
        return _get_backend_instance('torch')

    if isinstance(source, TensorLikeBackend):
        return source

    if source is None:
        return _get_backend_instance('numpy')

    raise ValueError(f'Unsupported data type: {type(source)}')
