import importlib
import inspect
import json
import os
import shutil
import subprocess
import sys
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import numpy as np
from transformers import PreTrainedTokenizer
from transformers.utils import cached_file

def _prepare_tokenmonster_runtime(tokenmonster_module) -> None:
    candidate_roots = [
        os.environ.get("STENTOR_TOKENMONSTER_DIR"),
        "/root/.cache" if os.path.isdir("/root/.cache") else None,
        os.getcwd(),
        os.environ.get("HOME"),
        os.environ.get("TMPDIR"),
        tempfile.gettempdir(),
    ]
    runtime_dir = None
    for root in candidate_roots:
        if not root:
            continue
        try:
            os.makedirs(root, exist_ok=True)
        except Exception:
            continue
        if not (os.access(root, os.W_OK) and os.access(root, os.X_OK)):
            continue
        runtime_dir = os.path.join(root, "stentor_tokenmonster_runtime")
        break
    if runtime_dir is None:
        raise SystemExit(
            "Unable to find a writable/executable directory for TokenMonster runtime files."
        )
    rank_suffix = os.environ.get("LOCAL_RANK") or os.environ.get("RANK")
    if rank_suffix is not None and str(rank_suffix).strip() != "":
        runtime_dir = os.path.join(runtime_dir, f"rank_{rank_suffix}")
    os.makedirs(runtime_dir, exist_ok=True)
    if hasattr(tokenmonster_module, "set_local_directory"):
        tokenmonster_module.set_local_directory(runtime_dir)

    server_name = "tokenmonsterserver"
    os_name = None
    if hasattr(tokenmonster_module, "_get_binary_filename"):
        try:
            os_name, server_name = tokenmonster_module._get_binary_filename()
        except Exception:
            os_name, server_name = None, "tokenmonsterserver"
    server_path = os.path.join(runtime_dir, server_name)
    if (not os.path.exists(server_path)) and os_name:
        try:
            cached_server = hf_hub_download(
                repo_id=DEFAULT_TOKENMONSTER_REPO,
                filename=f"binaries/{os_name}/{server_name}",
                token=resolve_hf_token(),
            )
            shutil.copy2(cached_server, server_path)
        except Exception:
            pass
    if os.path.exists(server_path):
        try:
            os.chmod(server_path, 0o755)
        except Exception:
            pass
        if not os.access(server_path, os.X_OK):
            raise SystemExit(
                f"TokenMonster server binary is not executable after chmod: {server_path}"
            )


def _repair_tokenmonster_server(tokenmonster_module) -> None:
    _prepare_tokenmonster_runtime(tokenmonster_module)
    vocab_cls = getattr(tokenmonster_module, "Vocab", None)
    if vocab_cls is None:
        return

    runtime_dir = getattr(vocab_cls, "_dir", None)
    server_name = None
    os_name = None
    if hasattr(tokenmonster_module, "_get_binary_filename"):
        try:
            os_name, server_name = tokenmonster_module._get_binary_filename()
        except Exception:
            os_name, server_name = None, None
    if not runtime_dir or not server_name:
        return

    server_path = os.path.join(runtime_dir, server_name)
    try:
        if os.path.exists(server_path):
            os.remove(server_path)
    except Exception:
        pass

    if os_name:
        try:
            cached_server = hf_hub_download(
                repo_id=DEFAULT_TOKENMONSTER_REPO,
                filename=f"binaries/{os_name}/{server_name}",
                token=resolve_hf_token(),
            )
            shutil.copy2(cached_server, server_path)
        except Exception:
            pass

    if os.path.exists(server_path):
        try:
            os.chmod(server_path, 0o755)
        except Exception:
            pass


def _ensure_tokenmonster_available():
    global tokenmonster
    if tokenmonster is not None:
        _prepare_tokenmonster_runtime(tokenmonster)
        return tokenmonster
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "tokenmonster"])
        tokenmonster = importlib.import_module("tokenmonster")
        _prepare_tokenmonster_runtime(tokenmonster)
        return tokenmonster
    except Exception as exc:
        raise SystemExit(
            "tokenmonster is required to use TokenMonster vocab files. "
            "Automatic install failed. Install with `pip install tokenmonster`."
        ) from exc


def _tokenmonster_remote_module_source() -> str:
    imports = [
        "import importlib",
        "import inspect",
        "import json",
        "import os",
        "import shutil",
        "import subprocess",
        "import sys",
        "from typing import Dict, Iterable, List, Optional, Sequence, Tuple",
        "",
        "import numpy as np",
        "from transformers import PreTrainedTokenizer",
        "from transformers.utils import cached_file",
    ]
    try:
        prepare_source = inspect.getsource(_prepare_tokenmonster_runtime)
        repair_source = inspect.getsource(_repair_tokenmonster_server)
        helper_source = inspect.getsource(_ensure_tokenmonster_available)
        remote_source = inspect.getsource(_tokenmonster_remote_module_source)
        class_source = inspect.getsource(TokenMonsterTokenizer)
    except (OSError, TypeError):
        source_text = _load_current_source_text()
        prepare_source = _extract_source_block(
            source_text,
            "def _prepare_tokenmonster_runtime(tokenmonster_module) -> None:",
            ["def _repair_tokenmonster_server(", "def print_hardware_diagnostic(", "def _ensure_tokenmonster_available("],
        )
        repair_source = _extract_source_block(
            source_text,
            "def _repair_tokenmonster_server(tokenmonster_module) -> None:",
            ["def print_hardware_diagnostic(", "def _ensure_tokenmonster_available("],
        )
        helper_source = _extract_source_block(
            source_text,
            "def _ensure_tokenmonster_available():",
            ["def _extract_source_block(", "def _tokenmonster_remote_module_source("],
        )
        extract_source = _extract_source_block(
            source_text,
            "def _extract_source_block(",
            ["def _load_current_source_text(", "def _tokenmonster_remote_module_source("],
        )
        load_source = _extract_source_block(
            source_text,
            "def _load_current_source_text(",
            ["def _tokenmonster_remote_module_source(", "class TokenMonsterTokenizer("],
        )
        remote_source = _extract_source_block(
            source_text,
            "def _tokenmonster_remote_module_source() -> str:",
            ["class TokenMonsterTokenizer("],
        )
        class_source = _extract_source_block(
            source_text,
            "class TokenMonsterTokenizer(PreTrainedTokenizer):",
            ["@dataclass", "\n\ndef parse_args("],
        )
        return "\n\n".join(
            [
                "\n".join(imports),
                prepare_source.strip(),
                repair_source.strip(),
                helper_source.strip(),
                extract_source.strip(),
                load_source.strip(),
                remote_source.strip(),
                class_source.strip(),
            ]
        ) + "\n"
    return "\n\n".join(
        [
            "\n".join(imports),
            prepare_source,
            repair_source,
            helper_source,
            remote_source,
            class_source,
        ]
    ) + "\n"


class TokenMonsterTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "tokenmonster.vocab"}
    model_input_names = ["input_ids", "attention_mask"]

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
        if "vocab_file" not in kwargs:
            resolved_vocab = None
            if os.path.isdir(pretrained_model_name_or_path):
                candidate = os.path.join(
                    pretrained_model_name_or_path,
                    cls.vocab_files_names["vocab_file"],
                )
                if os.path.exists(candidate):
                    resolved_vocab = candidate
            if resolved_vocab is None:
                cached_kwargs = {
                    key: kwargs[key]
                    for key in (
                        "cache_dir",
                        "force_download",
                        "local_files_only",
                        "proxies",
                        "resume_download",
                        "revision",
                        "subfolder",
                        "token",
                    )
                    if key in kwargs
                }
                try:
                    resolved_vocab = cached_file(
                        pretrained_model_name_or_path,
                        cls.vocab_files_names["vocab_file"],
                        **cached_kwargs,
                    )
                except Exception:
                    resolved_vocab = None
            if resolved_vocab is not None:
                kwargs["vocab_file"] = resolved_vocab
        return super().from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

    def __init__(self, vocab_file: str, model_max_length: int = 1_000_000, **kwargs) -> None:
        name_or_path = kwargs.get("name_or_path")
        if (
            vocab_file
            and not os.path.isabs(vocab_file)
            and not os.path.exists(vocab_file)
            and name_or_path
            and os.path.isdir(name_or_path)
        ):
            candidate = os.path.join(name_or_path, vocab_file)
            if os.path.exists(candidate):
                vocab_file = candidate
        self.vocab_file_path = os.path.abspath(vocab_file)
        self._tm = _ensure_tokenmonster_available()
        try:
            self._vocab = self._tm.load(self.vocab_file_path)
        except RuntimeError as exc:
            if "enable execute permissions" not in str(exc):
                raise
            _repair_tokenmonster_server(self._tm)
            self._vocab = self._tm.load(self.vocab_file_path)
        self._token_to_id: Dict[str, int] = {}
        self._id_to_token: Dict[int, str] = {}
        self._vocab_size_with_gaps = 0
        self._refresh_dictionary()

        kwargs.setdefault("bos_token", "<s>")
        kwargs.setdefault("eos_token", "</s>")
        kwargs.setdefault("unk_token", "<unk>")
        kwargs.setdefault("pad_token", "<pad>")
        kwargs.setdefault("model_max_length", int(model_max_length))

        super().__init__(vocab_file=self.vocab_file_path, **kwargs)
        if not hasattr(self, "additional_special_tokens"):
            self.additional_special_tokens = []
        elif self.additional_special_tokens is None:
            self.additional_special_tokens = []

    @property
    def vocab_size(self) -> int:
        return self._vocab_size_with_gaps

    def __len__(self) -> int:
        return self._vocab_size_with_gaps

    def get_vocab(self) -> Dict[str, int]:
        return dict(self._token_to_id)

    def prepare_for_tokenization(self, text, is_split_into_words: bool = False, **kwargs):
        return (text, kwargs)

    def _refresh_dictionary(self) -> None:
        token_to_id: Dict[str, int] = {}
        id_to_token: Dict[int, str] = {}
        max_token_id = -1
        entries = self._vocab.get_dictionary()
        if isinstance(entries, dict):
            entries = entries.values()
        for entry in entries:
            token = entry["token"]
            if isinstance(token, bytes):
                token = token.decode("utf-8", errors="replace")
            token_id = int(entry["id"])
            token_to_id[str(token)] = token_id
            id_to_token[token_id] = str(token)
            if token_id > max_token_id:
                max_token_id = token_id
        self._token_to_id = token_to_id
        self._id_to_token = id_to_token
        self._vocab_size_with_gaps = max(0, max_token_id + 1)

    def _clear_added_token_state(self) -> None:
        # Intentionally left as a no-op. Clearing HF's added-token state here causes
        # additional/special tokens to disappear from saved checkpoints.
        return

    @staticmethod
    def _normalize_token_value(token) -> Optional[str]:
        if token is None:
            return None
        if hasattr(token, "content"):
            return str(token.content)
        return str(token)

    @staticmethod
    def _normalize_token_ids(token_ids) -> List[int]:
        if isinstance(token_ids, np.ndarray):
            return [int(x) for x in token_ids.tolist()]
        if isinstance(token_ids, (list, tuple)):
            return [int(x) for x in token_ids]
        return [int(token_ids)]

    def _add_backend_token(self, token: Optional[str], special: bool) -> bool:
        if not token:
            return False
        if token in self._token_to_id:
            return False
        if special:
            self._vocab.add_special_token(token)
        else:
            self._vocab.add_token(token)
        self._refresh_dictionary()
        return token in self._token_to_id

    def _tokenize(self, text: str) -> List[str]:
        ids = self._normalize_token_ids(self._vocab.tokenize(text))
        unk_token = self.unk_token or "<unk>"
        return [self._id_to_token.get(int(i), unk_token) for i in ids]

    def _convert_token_to_id(self, token: str) -> Optional[int]:
        if token is None:
            return None
        token_id = self._token_to_id.get(token)
        if token_id is not None:
            return int(token_id)
        if self.unk_token is not None:
            return self._token_to_id.get(self.unk_token)
        return None

    def _convert_id_to_token(self, index: int) -> str:
        token = self._id_to_token.get(int(index))
        if token is not None:
            return token
        return self.unk_token or "<unk>"

    def convert_tokens_to_string(self, tokens: Sequence[str]) -> str:
        token_ids: List[int] = []
        for token in tokens:
            token_id = self._convert_token_to_id(token)
            if token_id is not None:
                token_ids.append(int(token_id))
        if not token_ids:
            return ""
        return self._vocab.decode(token_ids)

    def build_inputs_with_special_tokens(
        self,
        token_ids_0: List[int],
        token_ids_1: Optional[List[int]] = None,
    ) -> List[int]:
        output = list(token_ids_0)
        if self.bos_token_id is not None:
            output = [self.bos_token_id] + output
        if self.eos_token_id is not None:
            output = output + [self.eos_token_id]
        if token_ids_1 is None:
            return output
        second = list(token_ids_1)
        if self.eos_token_id is not None:
            second = second + [self.eos_token_id]
        return output + second

    def get_special_tokens_mask(
        self,
        token_ids_0: List[int],
        token_ids_1: Optional[List[int]] = None,
        already_has_special_tokens: bool = False,
    ) -> List[int]:
        special_ids = set(self.all_special_ids)
        if already_has_special_tokens:
            combined = list(token_ids_0)
            return [1 if token in special_ids else 0 for token in combined]
        mask = [0] * len(token_ids_0)
        if self.bos_token_id is not None:
            mask = [1] + mask
        if self.eos_token_id is not None:
            mask = mask + [1]
        if token_ids_1 is None:
            return mask
        second_mask = [0] * len(token_ids_1)
        if self.eos_token_id is not None:
            second_mask = second_mask + [1]
        return mask + second_mask

    def create_token_type_ids_from_sequences(
        self,
        token_ids_0: List[int],
        token_ids_1: Optional[List[int]] = None,
    ) -> List[int]:
        length = len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1))
        return [0] * length

    def add_tokens(self, new_tokens: Iterable[str], special_tokens: bool = False) -> int:
        if isinstance(new_tokens, str):
            new_tokens = [new_tokens]
        added = 0
        for token in new_tokens:
            if self._add_backend_token(self._normalize_token_value(token), special_tokens):
                added += 1
        return added

    def add_special_tokens(
        self,
        special_tokens_dict: Dict[str, object],
        replace_additional_special_tokens: bool = True,
    ) -> int:
        added = 0
        single_token_keys = (
            "unk_token",
            "bos_token",
            "eos_token",
            "pad_token",
            "sep_token",
            "cls_token",
            "mask_token",
        )
        for key in single_token_keys:
            token = self._normalize_token_value(special_tokens_dict.get(key))
            if token is None:
                continue
            if self._add_backend_token(token, special=True):
                added += 1
            setattr(self, key, token)

        extras = special_tokens_dict.get("additional_special_tokens")
        if extras is not None:
            normalized_extras: List[str] = []
            for token in extras:
                normalized = self._normalize_token_value(token)
                if normalized is None:
                    continue
                if self._add_backend_token(normalized, special=True):
                    added += 1
                if normalized not in normalized_extras:
                    normalized_extras.append(normalized)
            if replace_additional_special_tokens:
                self.additional_special_tokens = normalized_extras
            else:
                merged = list(getattr(self, "additional_special_tokens", []) or [])
                for token in normalized_extras:
                    if token not in merged:
                        merged.append(token)
                self.additional_special_tokens = merged
        return added

    def _decode(
        self,
        token_ids,
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: Optional[bool] = None,
        spaces_between_special_tokens: bool = False,
        **kwargs,
    ) -> str:
        ids = self._normalize_token_ids(token_ids)
        if skip_special_tokens:
            special_ids = set(self.all_special_ids)
            ids = [token_id for token_id in ids if token_id not in special_ids]
        if not ids:
            return ""
        return self._vocab.decode(ids)

    def save_vocabulary(
        self,
        save_directory: str,
        filename_prefix: Optional[str] = None,
    ) -> Tuple[str]:
        os.makedirs(save_directory, exist_ok=True)
        prefix = f"{filename_prefix}-" if filename_prefix else ""
        vocab_path = os.path.join(save_directory, f"{prefix}tokenmonster.vocab")
        vocab_json_path = os.path.join(save_directory, f"{prefix}vocab.json")
        self._vocab.save(vocab_path)
        with open(vocab_json_path, "w", encoding="utf-8") as handle:
            json.dump(self.get_vocab(), handle, ensure_ascii=False, indent=2, sort_keys=True)
        return (vocab_path,)

    def save_pretrained(
        self,
        save_directory: str,
        legacy_format: Optional[bool] = None,
        filename_prefix: Optional[str] = None,
        push_to_hub: bool = False,
        **kwargs,
    ):
        saved_files = super().save_pretrained(
            save_directory,
            legacy_format=legacy_format,
            filename_prefix=filename_prefix,
            push_to_hub=push_to_hub,
            **kwargs,
        )

        module_out = os.path.join(save_directory, "tokenmonster_hf.py")
        try:
            module_path = inspect.getsourcefile(type(self))
        except (OSError, TypeError):
            module_path = None
        if module_path and os.path.isfile(module_path) and os.path.basename(module_path) == "tokenmonster_hf.py":
            shutil.copy(module_path, module_out)
        else:
            with open(module_out, "w", encoding="utf-8") as handle:
                handle.write(_tokenmonster_remote_module_source())

        prefix = f"{filename_prefix}-" if filename_prefix else ""
        config_path = os.path.join(save_directory, f"{prefix}tokenizer_config.json")
        config = {}
        if os.path.exists(config_path):
            with open(config_path, "r", encoding="utf-8") as handle:
                config = json.load(handle)
        config.update(
            {
                "tokenizer_class": "TokenMonsterTokenizer",
                "auto_map": {
                    "AutoTokenizer": ["tokenmonster_hf.TokenMonsterTokenizer", None]
                },
                "vocab_file": self.vocab_files_names["vocab_file"],
                "bos_token": self.bos_token,
                "eos_token": self.eos_token,
                "unk_token": self.unk_token,
                "pad_token": self.pad_token,
            }
        )
        with open(config_path, "w", encoding="utf-8") as handle:
            json.dump(config, handle, ensure_ascii=False, indent=2, sort_keys=True)
        added_tokens_path = os.path.join(save_directory, f"{prefix}added_tokens.json")
        added_tokens = {}
        for token in list(getattr(self, "additional_special_tokens", []) or []):
            token_id = self._convert_token_to_id(token)
            if token_id is not None:
                added_tokens[str(token)] = int(token_id)
        if added_tokens:
            with open(added_tokens_path, "w", encoding="utf-8") as handle:
                json.dump(added_tokens, handle, ensure_ascii=False, indent=2, sort_keys=True)
        return saved_files