from ..entities import (
    EngineSettings,
    ModelConfig,
    ParallelStrategy,
    SentenceTransformerModelConfig,
    TokenizerConfig,
    WeightType,
)
from ..model import (
    ModelAlreadyLoadedException,
    TokenizerAlreadyLoadedException,
    TokenizerNotSupportedException,
)
from ..model.vendor import TextGenerationVendor

import asyncio
from abc import ABC, abstractmethod
from contextlib import AsyncExitStack
from importlib import import_module
from importlib.util import find_spec
from logging import ERROR, Logger, getLogger
from sys import modules
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast

if TYPE_CHECKING:
    from diffusers import DiffusionPipeline as DiffusionPipeline
    from torch import dtype
    from transformers import PreTrainedModel as PreTrainedModel
    from transformers import (
        PreTrainedTokenizer,
        PreTrainedTokenizerFast,
    )
else:
    dtype: TypeAlias = Any

    class PreTrainedTokenizer:  # noqa: D101
        pass

    class PreTrainedTokenizerFast:  # noqa: D101
        pass

    class DiffusionPipeline:  # noqa: D101
        pass

    class PreTrainedModel:  # noqa: D101
        pass


class _LazyModule:
    def __init__(self, module_name: str) -> None:
        self._module_name = module_name

    def __getattr__(self, name: str) -> Any:
        return getattr(import_module(self._module_name), name)


class _TransformersLoggingProxy:
    def get_logger(self) -> Logger:
        logging_module = import_module("transformers.utils.logging")
        return cast(Logger, getattr(logging_module, "get_logger")())

    def set_verbosity_error(self) -> None:
        logging_module = import_module("transformers.utils.logging")
        cast(Any, getattr(logging_module, "set_verbosity_error"))()

    def __getattr__(self, name: str) -> Any:
        return getattr(import_module("transformers.utils.logging"), name)


cuda = _LazyModule("torch.cuda")
mps = _LazyModule("torch.backends.mps")
transformers_logging = _TransformersLoggingProxy()
_DIFFUSION_PIPELINE_PLACEHOLDER = DiffusionPipeline
_PRETRAINED_MODEL_PLACEHOLDER = PreTrainedModel


def disable_progress_bar() -> None:
    logging_module = import_module("transformers.utils.logging")
    cast(Any, getattr(logging_module, "disable_progress_bar"))()


def enable_progress_bar() -> None:
    logging_module = import_module("transformers.utils.logging")
    cast(Any, getattr(logging_module, "enable_progress_bar"))()


def _optional_type(module_name: str, type_name: str) -> type[Any] | None:
    if module_name in modules:
        module = modules[module_name]
        return cast(type[Any], getattr(module, type_name))
    return None


def _pretrained_model_type() -> type[Any] | None:
    if PreTrainedModel is not _PRETRAINED_MODEL_PLACEHOLDER:
        return cast(type[Any], PreTrainedModel)
    return _optional_type("transformers", "PreTrainedModel")


def _pretrained_tokenizer_types() -> tuple[type[Any], type[Any]] | None:
    tokenizer_type = _optional_type("transformers", "PreTrainedTokenizer")
    tokenizer_fast_type = _optional_type(
        "transformers", "PreTrainedTokenizerFast"
    )
    if tokenizer_type is None or tokenizer_fast_type is None:
        return None
    return tokenizer_type, tokenizer_fast_type


def _diffusion_pipeline_type() -> type[Any] | None:
    if DiffusionPipeline is not _DIFFUSION_PIPELINE_PLACEHOLDER:
        return cast(type[Any], DiffusionPipeline)
    return _optional_type("diffusers", "DiffusionPipeline")


def _set_transformers_progress_bar(enabled: bool) -> None:
    if "transformers" not in modules and find_spec("transformers") is None:
        return
    if enabled:
        enable_progress_bar()
    else:
        disable_progress_bar()


class Engine(ABC):
    _device: str
    _logger: Logger
    _model_id: str | None
    _settings: EngineSettings
    _transformers_logging_logger: Logger | None
    _transformers_logging_level: int | None
    _loaded_model: bool = False
    _loaded_tokenizer: bool = False
    _tokenizer: Any = None
    _model: Any = None
    _config: ModelConfig | SentenceTransformerModelConfig | None = None
    _tokenizer_config: TokenizerConfig | None = None
    _parameter_types: set[str] | None = None
    _parameter_count: int | None = None
    _exit_stack: AsyncExitStack = AsyncExitStack()
    _pending_exit_task: asyncio.Task[None] | None = None

    DTYPE_SIZES: dict[str, int] = {
        "bool": 1,
        "bfloat16": 2,
        "float16": 2,
        "float32": 4,
        "float64": 8,
        "int8": 1,
        "int16": 2,
        "i32": 4,
        "i64": 8,
        "ui8": 1,
    }

    _WEIGHTS: Final[dict[str, Literal["auto"] | str | dtype]] = {
        "bool": "bool",
        "bf16": "bfloat16",
        "f16": "float16",
        "fp16": "float16",
        "f32": "float32",
        "fp32": "float32",
        "f64": "float64",
        "fp64": "float64",
        "i8": "int8",
        "i16": "int16",
        "i32": "int32",
        "i64": "int64",
        "ui8": "uint8",
        "auto": "auto",
    }

    @staticmethod
    def _get_tp_plan(
        parallel: ParallelStrategy | dict[str, ParallelStrategy] | None,
    ) -> str | dict[str, str] | None:
        if parallel is None:
            return None
        if isinstance(parallel, dict):
            return {k: v.value for k, v in parallel.items()}
        return str(parallel.value)

    @staticmethod
    def _get_distributed_config(
        distributed_config: dict[str, object] | None,
    ) -> dict[str, object] | None:
        if distributed_config is None:
            return None
        config: dict[str, object] = {"enable_expert_parallel": False}
        config.update(distributed_config)
        return config

    @staticmethod
    def weight(weight_type: WeightType) -> Literal["auto"] | dtype:
        weight_name = Engine._WEIGHTS.get(weight_type, "auto")
        if weight_name == "auto":
            return "auto"
        if not isinstance(weight_name, str):
            return weight_name
        resolved = cast(dtype, getattr(import_module("torch"), weight_name))
        Engine._WEIGHTS[weight_type] = resolved
        return resolved

    @staticmethod
    def _get_config_dtype(config: object) -> dtype:
        missing = object()
        config_dtype = getattr(config, "dtype", missing)
        if config_dtype is not missing:
            return cast(dtype, config_dtype or Engine.weight("f32"))

        if hasattr(config, "torch_dtype"):
            return cast(
                dtype, getattr(config, "torch_dtype") or Engine.weight("f32")
            )

        return cast(dtype, Engine.weight("f32"))

    def _decode_token(self, token_id: int | None) -> str | None:
        return (
            cast(str, self._tokenizer.decode(token_id))
            if self._tokenizer is not None and token_id is not None
            else None
        )

    def __init__(
        self,
        model_id: str | None,
        settings: EngineSettings | None = None,
        logger: Logger = getLogger(__name__),
    ):
        self._logger = logger
        self._model_id = model_id
        self._settings = settings if settings else EngineSettings()
        self._device = (
            self._settings.device
            if self._settings.device
            else Engine.get_default_device()
        )
        uses_tokenizer = (
            self.uses_tokenizer()
            if callable(self.uses_tokenizer)
            else self.uses_tokenizer
        )
        self._transformers_logging_logger = None
        self._transformers_logging_level = None
        if (
            self._settings.change_transformers_logging_level
            and uses_tokenizer
            and find_spec("transformers") is not None
        ):
            self._transformers_logging_logger = cast(
                Logger, getattr(transformers_logging, "get_logger")()
            )
            self._transformers_logging_level = (
                self._transformers_logging_logger.level
            )

        auto_load_tokenizer = (
            uses_tokenizer and self._settings.auto_load_tokenizer
        )
        if self._settings.auto_load_model or auto_load_tokenizer:
            self._load(
                load_tokenizer=auto_load_tokenizer,
                tokenizer_name_or_path=(
                    self._settings.tokenizer_name_or_path
                    if auto_load_tokenizer
                    else None
                ),
            )

    @property
    def uses_tokenizer(self) -> bool:
        return False

    @property
    def config(
        self,
    ) -> ModelConfig | SentenceTransformerModelConfig | None:
        return self._config

    @property
    def model(
        self,
    ) -> Any:
        return self._model

    @property
    def model_type(self) -> str:
        return type(self).__name__

    @property
    def model_id(self) -> str | None:
        return self._model_id

    @property
    def parameter_count(self) -> int | None:
        return self._parameter_count

    @property
    def parameter_types(self) -> set[str] | None:
        return self._parameter_types

    @property
    def tokenizer_config(self) -> TokenizerConfig | None:
        return self._tokenizer_config

    @property
    def tokenizer(
        self,
    ) -> PreTrainedTokenizer | PreTrainedTokenizerFast | None:
        return cast(
            PreTrainedTokenizer | PreTrainedTokenizerFast | None,
            self._tokenizer,
        )

    async def __call__(self, *args: object, **kwargs: object) -> object:
        raise NotImplementedError()

    @abstractmethod
    def _load_model(
        self,
    ) -> Any:
        raise NotImplementedError()

    def _accepts_loaded_model(self, model: object) -> bool:
        """Return whether a subclass accepts a native model object."""
        return False

    def is_runnable(self, device: str | None = None) -> bool | None:
        if (
            self._parameter_types is None
            or not self._parameter_types
            or self._parameter_count is None
            or self._parameter_count <= 0
        ):
            return None

        if not device:
            device = Engine.get_default_device()

        available = Engine._get_device_memory(device)
        if not available:
            return False

        dtype = next(iter(self._parameter_types))
        bytes_per_param = self.DTYPE_SIZES.get(dtype, 4)
        required = self._parameter_count * bytes_per_param
        return required <= available

    def _load_tokenizer_with_tokens(
        self, tokenizer_name_or_path: str | None, use_fast: bool = True
    ) -> PreTrainedTokenizer | PreTrainedTokenizerFast:
        uses_tokenizer = (
            self.uses_tokenizer()
            if callable(self.uses_tokenizer)
            else self.uses_tokenizer
        )
        raise (
            TokenizerNotSupportedException()
            if not uses_tokenizer
            else NotImplementedError()
        )

    def __enter__(self) -> "Engine":
        _l = self._log
        if (
            self._transformers_logging_logger
            and self._transformers_logging_level != ERROR
        ):
            cast(Any, getattr(transformers_logging, "set_verbosity_error"))()
            _l(
                "Changed transformers logging level from %s to %s",
                self._transformers_logging_level,
                self._transformers_logging_logger.level,
            )
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: Any | None,
    ) -> Literal[False]:
        _l = self._log
        if (
            self._transformers_logging_logger
            and self._transformers_logging_level
            and self._transformers_logging_level != ERROR
        ):
            self._transformers_logging_logger.setLevel(
                self._transformers_logging_level
            )
            _l(
                "Restored transformers logging level to %s",
                self._transformers_logging_level,
            )
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            asyncio.run(self._exit_stack.aclose())
        else:
            self._pending_exit_task = loop.create_task(
                self._exit_stack.aclose()
            )
        return False

    async def wait_closed(self) -> None:
        """Wait for any asynchronous close task scheduled by __exit__."""
        if self._pending_exit_task is None:
            return
        await self._pending_exit_task
        self._pending_exit_task = None

    def _load(
        self,
        *args: object,
        load_tokenizer: bool,
        tokenizer_name_or_path: str | None,
    ) -> None:
        if (
            self._settings.auto_load_model
            and hasattr(self, "_loaded_model")
            and self._loaded_model
        ):
            raise ModelAlreadyLoadedException()
        elif (
            load_tokenizer
            and hasattr(self, "_loaded_tokenizer")
            and self._loaded_tokenizer
        ):
            raise TokenizerAlreadyLoadedException()

        _l = self._log

        uses_tokenizer = (
            self.uses_tokenizer()
            if callable(self.uses_tokenizer)
            else self.uses_tokenizer
        )
        should_change_transformers_progress = (
            self._settings.disable_loading_progress_bar and uses_tokenizer
        )
        if should_change_transformers_progress:
            _set_transformers_progress_bar(enabled=False)

        if load_tokenizer and self._model_id:
            _l(
                "Loading tokenizer %s",
                tokenizer_name_or_path or self._model_id,
            )
            self._tokenizer = self._load_tokenizer_with_tokens(
                tokenizer_name_or_path or self._model_id, use_fast=True
            )
            tokenizer_types = _pretrained_tokenizer_types()
            assert tokenizer_types is not None
            assert isinstance(
                self._tokenizer, tokenizer_types[0]
            ) or isinstance(self._tokenizer, tokenizer_types[1]), (
                "Unexpected pretrained tokenizer type: "
                + f"{type(self._tokenizer)}"
            )
            _l("Loaded tokenizer %s", self._tokenizer.name_or_path)

            self._loaded_tokenizer = True

        is_mlx = False
        is_sentence_transformer = False

        if self._settings.auto_load_model:
            _l(
                "Loading pretrained model %s from cache %s",
                self._model_id or str(self._model),
                self._settings.cache_dir,
            )
            self._model = self._load_model()

            is_vendor = isinstance(self._model, TextGenerationVendor)
            pretrained_model_type = (
                None if is_vendor else _pretrained_model_type()
            )
            is_pretrained_model = (
                pretrained_model_type is not None
                and isinstance(self._model, pretrained_model_type)
            )
            accepts_loaded_model = False
            if not is_pretrained_model and not is_vendor:
                accepts_loaded_model = self._accepts_loaded_model(self._model)
            diffusion_pipeline_type = (
                None
                if is_vendor or is_pretrained_model or accepts_loaded_model
                else _diffusion_pipeline_type()
            )
            is_diffusion_pipeline = bool(
                diffusion_pipeline_type is not None
                and isinstance(self._model, diffusion_pipeline_type)
            )

            if (
                not is_pretrained_model
                and not is_vendor
                and not is_diffusion_pipeline
                and not accepts_loaded_model
            ):
                is_mlx = Engine._is_mlx_model(self._model)

                sentence_transformer_type = (
                    None
                    if is_mlx
                    else _optional_type(
                        "sentence_transformers", "SentenceTransformer"
                    )
                )
                if sentence_transformer_type is not None:
                    is_sentence_transformer = isinstance(
                        self._model, sentence_transformer_type
                    )

            assert (
                is_pretrained_model
                or is_vendor
                or is_diffusion_pipeline
                or is_mlx
                or is_sentence_transformer
                or accepts_loaded_model
            ), f"Unexpected pretrained model type: {type(self._model)}"

            _l(
                "Loaded pretrained model %s from cache %s",
                self._model_id,
                self._settings.cache_dir,
            )

            if self._settings.enable_eval:
                _l("Setting model %s in eval mode", self._model_id)
                self._model.eval()

            if self._tokenizer and (
                self._settings.tokens or self._settings.special_tokens
            ):
                total_tokens = len(self._tokenizer)
                _l(
                    "Resizing embedding matrix to %s tokens for model %s",
                    total_tokens,
                    self._model_id,
                )
                self._model.resize_token_embeddings(total_tokens)
                _l(
                    "Resized embedding matrix to %s tokens for model %s",
                    total_tokens,
                    self._model_id,
                )

            self._loaded_model = True

        parameters = (
            getattr(self._model, "parameters", None)
            if not is_mlx and self._model
            else None
        )
        self._parameter_types = (
            {str(param.dtype).replace("torch.", "") for param in parameters()}
            if callable(parameters)
            else None
        )
        self._parameter_count = (
            sum(p.numel() for p in parameters())
            if callable(parameters)
            else None
        )

        if self._model and not self._config:
            config: ModelConfig | SentenceTransformerModelConfig | None = None
            mc = (
                self._model.config
                if hasattr(self._model, "config")
                else (
                    self._model[0].auto_model.config
                    if is_sentence_transformer
                    else None
                )
            )

            if mc:
                bos_token_id = getattr(mc, "bos_token_id", None)
                eos_token_id = getattr(mc, "eos_token_id", None)
                pad_token_id = getattr(mc, "pad_token_id", None)
                sep_token_id = getattr(mc, "sep_token_id", None)
                config = ModelConfig(
                    architectures=getattr(mc, "architectures", None),
                    attribute_map=cast(
                        dict[str, str], getattr(mc, "attribute_map", {})
                    ),
                    bos_token_id=bos_token_id,
                    bos_token=self._decode_token(bos_token_id),
                    decoder_start_token_id=getattr(
                        mc, "decoder_start_token_id", None
                    ),
                    eos_token_id=eos_token_id,
                    eos_token=self._decode_token(eos_token_id),
                    finetuning_task=getattr(mc, "finetuning_task", None),
                    hidden_size=(
                        mc.hidden_size if hasattr(mc, "hidden_size") else None
                    ),
                    hidden_sizes=(
                        mc.hidden_sizes
                        if hasattr(mc, "hidden_sizes")
                        else None
                    ),
                    keys_to_ignore_at_inference=cast(
                        list[str],
                        (
                            mc.keys_to_ignore_at_inference
                            if hasattr(mc, "keys_to_ignore_at_inference")
                            else []
                        ),
                    ),
                    loss_type=(
                        mc.loss_type if hasattr(mc, "loss_type") else None
                    ),
                    max_position_embeddings=(
                        mc.max_position_embeddings
                        if hasattr(mc, "max_position_embeddings")
                        else None
                    ),
                    model_type=getattr(mc, "model_type", None),
                    num_attention_heads=(
                        mc.num_attention_heads
                        if hasattr(mc, "num_attention_heads")
                        else None
                    ),
                    num_hidden_layers=(
                        mc.num_hidden_layers
                        if hasattr(mc, "num_hidden_layers")
                        else None
                    ),
                    num_labels=getattr(mc, "num_labels", None),
                    output_attentions=cast(
                        bool, getattr(mc, "output_attentions", False)
                    ),
                    output_hidden_states=cast(
                        bool, getattr(mc, "output_hidden_states", False)
                    ),
                    pad_token_id=pad_token_id,
                    pad_token=self._decode_token(pad_token_id),
                    prefix=getattr(mc, "prefix", None),
                    sep_token_id=sep_token_id,
                    sep_token=self._decode_token(sep_token_id),
                    state_size=(
                        len(self._model.state_dict().keys())
                        if hasattr(self._model, "state_dict")
                        and self._model.state_dict
                        else 0
                    ),
                    task_specific_params=getattr(
                        mc, "task_specific_params", None
                    ),
                    torch_dtype=Engine._get_config_dtype(mc),
                    vocab_size=(
                        mc.vocab_size if hasattr(mc, "vocab_size") else None
                    ),
                    tokenizer_class=getattr(mc, "tokenizer_class", None),
                )

            if is_sentence_transformer and config:
                config = SentenceTransformerModelConfig(
                    backend=self._model.backend,
                    similarity_function=self._model.similarity_fn_name,
                    truncate_dimension=self._model.truncate_dim,
                    transformer_model_config=cast(ModelConfig, config),
                )

            self._config = config

        if self._tokenizer and not self._tokenizer_config:
            tokenizer_types = _pretrained_tokenizer_types()
            tokenizer_fast_type = (
                tokenizer_types[1] if tokenizer_types is not None else None
            )
            self._tokenizer_config = TokenizerConfig(
                name_or_path=self._tokenizer.name_or_path,
                tokens=self._settings.tokens,
                special_tokens=self._tokenizer.all_special_tokens,
                tokenizer_model_max_length=getattr(
                    self._tokenizer, "model_max_length", 0
                ),
                fast=(
                    tokenizer_fast_type is not None
                    and isinstance(self._tokenizer, tokenizer_fast_type)
                ),
            )

        if should_change_transformers_progress:
            _set_transformers_progress_bar(enabled=True)

    @staticmethod
    def get_default_device() -> str:
        if find_spec("torch") is None:
            return "cpu"
        try:
            cuda_available = bool(cuda.is_available())
        except ModuleNotFoundError:
            return "cpu"
        if cuda_available:
            return "cuda"
        return "mps" if bool(mps.is_available()) else "cpu"

    @staticmethod
    def _has_module(name: str) -> bool:
        try:
            return find_spec(name) is not None
        except ModuleNotFoundError:
            return False

    @staticmethod
    def _is_mlx_model(model: object) -> bool:
        return any(
            cls.__name__ == "Module"
            and (
                cls.__module__ == "mlx.nn"
                or cls.__module__.startswith("mlx.nn.")
            )
            for cls in type(model).__mro__
        )

    @staticmethod
    def _get_device_memory(device: str) -> int:
        """Return available memory for device in bytes."""
        if device.startswith("cuda") or device == "cuda":
            if find_spec("torch") is None:
                return 0
            if not bool(cuda.is_available()):
                return 0
            index = (
                int(device.split(":", 1)[1])
                if ":" in device
                else int(cuda.current_device())
            )
            return int(cuda.get_device_properties(index).total_memory)

        from psutil import virtual_memory

        if (
            device == "mps"
            and find_spec("torch") is not None
            and bool(mps.is_available())
        ):
            return int(virtual_memory().total)
        return int(virtual_memory().total)

    def _log(self, message: str, *args: object) -> None:
        self._logger.debug(
            f"<Engine %s (%s)> {message}",
            type(self).__name__,
            self._model_id,
            *args,
        )
