from ..entities import (
    Input,
    Token,
    TransformerEngineSettings,
)
from ..model.engine import Engine

from abc import ABC, abstractmethod
from importlib import import_module
from logging import Logger, getLogger
from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast

if TYPE_CHECKING:
    from torch import Tensor
    from transformers import (
        AutoTokenizer,
        PreTrainedTokenizer,
        PreTrainedTokenizerFast,
    )
    from transformers.tokenization_utils_base import BatchEncoding
else:
    Tensor: TypeAlias = Any

    class BatchEncoding:  # noqa: D101
        pass

    class PreTrainedTokenizer:  # noqa: D101
        pass

    class PreTrainedTokenizerFast:  # noqa: D101
        pass

    class _LazyExternal:
        __test__ = False

        def __init__(self, module_name: str, name: str) -> None:
            self._module_name = module_name
            self._name = name

        def from_pretrained(self, *args: object, **kwargs: object) -> Any:
            module = import_module(self._module_name)
            target = getattr(module, self._name)
            return target.from_pretrained(*args, **kwargs)

        def __getattr__(self, name: str) -> Any:
            if name == "_is_coroutine" or (
                name.startswith("__") and name.endswith("__")
            ):
                raise AttributeError(name)
            module = import_module(self._module_name)
            target = getattr(module, self._name)
            return getattr(target, name)

    AutoTokenizer = _LazyExternal("transformers", "AutoTokenizer")


def _batch_encoding_type() -> type[Any]:
    module = import_module("transformers.tokenization_utils_base")
    return cast(type[Any], getattr(module, "BatchEncoding"))


def _tensor_type() -> type[Any]:
    return cast(type[Any], getattr(import_module("torch"), "Tensor"))


class TransformerModel(Engine, ABC):
    @property
    def supports_sample_generation(self) -> bool:
        return False

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

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

    @abstractmethod
    def _tokenize_input(
        self,
        input: Input,
        context: str | None = None,
        tensor_format: Literal["pt"] = "pt",
        **kwargs: object,
    ) -> dict[str, Tensor] | BatchEncoding | Tensor:
        raise NotImplementedError()

    def __init__(
        self,
        model_id: str,
        settings: TransformerEngineSettings | None = None,
        logger: Logger = getLogger(__name__),
    ) -> None:
        super().__init__(
            model_id, settings or TransformerEngineSettings(), logger
        )

    def save_tokenizer(self, path: str) -> list[str]:
        assert path and self._tokenizer
        _l = self._log
        _l(f"Saving tokenizer {self._tokenizer.name_or_path} to {path}")
        paths = self._tokenizer.save_pretrained(path)
        _l(
            f"Saved tokenizer {self._tokenizer.name_or_path} to {path}:"
            f" {paths}"
        )
        return list(paths)

    def tokenize(
        self, text: str, tokenizer_name_or_path: str | None = None
    ) -> list[Token]:
        _l = self._log
        if (
            not hasattr(self, "_loaded_tokenizer")
            or not self._loaded_tokenizer
        ):
            cast(
                Any,
                self,
            ).load(
                load_model=False,
                load_tokenizer=True,
                tokenizer_name_or_path=tokenizer_name_or_path,
            )

        _l(f'Tokenizing text "{text}"')
        token_ids = self._tokenizer.encode(text, add_special_tokens=True)
        _l(f'Tokenized text "{text}" into {len(token_ids)} tokens')

        return [
            Token(
                id=token_id,
                token=self._tokenizer.decode(
                    token_id, skip_special_tokens=False
                ),
                probability=1,
            )
            for i, token_id in enumerate(token_ids)
        ]

    def input_token_count(
        self,
        input: Input,
        system_prompt: str | None = None,
        developer_prompt: str | None = None,
        *,
        instructions: str | None = None,
    ) -> int:
        _l = self._log
        assert self._tokenizer, (
            f"Model {self._model} can't be executed "
            + "without a tokenizer loaded first"
        )
        inputs = self._tokenize_input(
            input,
            system_prompt=system_prompt,
            developer_prompt=developer_prompt,
            context=None,
            instructions=instructions,
        )
        if not isinstance(inputs, (dict, _batch_encoding_type())):
            return 0

        input_ids = inputs.get("input_ids")
        tensor_type = _tensor_type()
        if isinstance(input_ids, tensor_type):
            if input_ids.numel() == 0:
                return 0
            first_row = input_ids[0]
            return (
                int(first_row.shape[0])
                if isinstance(first_row, tensor_type)
                else 0
            )

        if (
            isinstance(input_ids, list)
            and input_ids
            and isinstance(input_ids[0], list)
        ):
            return len(input_ids[0])

        if not isinstance(input_ids, list):
            return 0

        return len(input_ids)

    def _load_tokenizer(
        self, tokenizer_name_or_path: str | None, use_fast: bool
    ) -> PreTrainedTokenizer | PreTrainedTokenizerFast:
        return cast(
            PreTrainedTokenizer | PreTrainedTokenizerFast,
            cast(Any, AutoTokenizer).from_pretrained(
                tokenizer_name_or_path or self._model_id,
                use_fast=use_fast,
                subfolder=self._settings.tokenizer_subfolder or "",
            ),
        )

    def _load_tokenizer_with_tokens(
        self, tokenizer_name_or_path: str | None, use_fast: bool = True
    ) -> PreTrainedTokenizer | PreTrainedTokenizerFast:
        _l = self._log
        tokenizer = self._load_tokenizer(tokenizer_name_or_path, use_fast)
        # Cleanup is intended for WordPiece and can corrupt BPE output.
        tokenizer.clean_up_tokenization_spaces = False
        if self._settings.tokens:
            _l(
                f"Adding {len(self._settings.tokens)} tokens to tokenizer "
                f"{tokenizer.name_or_path}: {self._settings.tokens}"
            )
            added_tokens = tokenizer.add_tokens(self._settings.tokens)
            _l(
                f"Added {added_tokens} tokens to tokenizer "
                f"{tokenizer.name_or_path}"
            )

        if self._settings.special_tokens:
            added_token_type = getattr(
                import_module("tokenizers"), "AddedToken"
            )
            _l(
                f"Adding {len(self._settings.special_tokens)} special tokens "
                f"to tokenizer {tokenizer.name_or_path}: "
                f"{self._settings.special_tokens}"
            )
            added_tokens = tokenizer.add_special_tokens(
                {
                    "additional_special_tokens": [
                        added_token_type(
                            token,
                            # Defines whether this token should strip all
                            # potential
                            # whitespaces on its left side
                            lstrip=False,
                            # Defines whether this token should match against
                            # the normalized version of the input text
                            normalized=False,
                            # Defines whether this token should strip all
                            # potential whitespaces on its right side
                            rstrip=False,
                            # Defines whether this token should only match
                            # single words
                            single_word=False,
                            # Defines whether this token should be skipped when
                            # decoding
                            special=False,
                        )
                        for token in self._settings.special_tokens
                    ]
                }
            )
            _l(
                f"Added {added_tokens} special tokens to tokenizer "
                f"{tokenizer.name_or_path}"
            )

        return tokenizer
