"""HuggingFace/PyTorch shard backend for real node inference."""

from __future__ import annotations

import base64
from dataclasses import dataclass
import json
from pathlib import Path
from typing import Any, Literal

Quantization = Literal["auto", "bfloat16", "int8", "nf4"]


class ModelBackendError(RuntimeError):
    """Base class for real model backend startup and execution failures."""


class MissingModelDependencyError(ModelBackendError):
    """Raised when optional model dependencies are not installed."""


class InsufficientVRAMError(ModelBackendError):
    """Raised when a requested shard cannot fit in available CUDA memory."""


class PartialModelLoadUnsupported(ModelBackendError):
    """Raised when a shard cannot be materialized from a local snapshot subset."""


@dataclass(frozen=True)
class TensorPayload:
    body: bytes
    shape: list[int]
    attention_mask_header: str | None
    position_ids_header: str | None


def validate_quantization(value: str) -> Quantization:
    if value not in {"auto", "bfloat16", "int8", "nf4"}:
        raise ValueError("quantization must be one of: auto, bfloat16, int8, nf4")
    return value  # type: ignore[return-value]


def build_quantization_config(quantization: Quantization) -> Any | None:
    """Return a transformers BitsAndBytesConfig for quantized weights."""
    if quantization in {"auto", "bfloat16"}:
        return None
    try:
        import torch
        from transformers import BitsAndBytesConfig
    except ModuleNotFoundError as exc:
        raise MissingModelDependencyError(
            "transformers and torch are required for int8/nf4 quantization"
        ) from exc

    if quantization == "int8":
        return BitsAndBytesConfig(load_in_8bit=True)
    return BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )


class TorchModelShard:
    """Executable subset of a HuggingFace causal language model."""

    def __init__(
        self,
        model_id: str,
        shard_start: int,
        shard_end: int,
        quantization: Quantization = "auto",
        cache_dir: Path | None = None,
    ) -> None:
        if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
            raise ValueError("shard_start must be <= shard_end and non-negative")
        self.model_id = model_id
        self.shard_start = shard_start
        self.shard_end = shard_end
        self.quantization = quantization

        try:
            import torch
            from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
        except ModuleNotFoundError as exc:
            raise MissingModelDependencyError(
                "real model backend requires torch, transformers, safetensors, accelerate, and bitsandbytes"
            ) from exc

        self.torch = torch
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id
        quant_config, dtype, uses_quantized_weights = _model_load_plan(
            AutoConfig,
            load_source,
            quantization,
            torch,
            None if load_source != model_id else cache_dir,
        )
        try:
            total_layers_hint = _total_layers_for_local_snapshot(AutoConfig, load_source)
            if _should_partial_materialize_shard(
                load_source,
                shard_start,
                shard_end,
                total_layers_hint=total_layers_hint,
                uses_quantized_weights=uses_quantized_weights,
            ):
                self.model = _load_partial_model_from_snapshot(
                    AutoConfig,
                    AutoModelForCausalLM,
                    torch,
                    load_source,
                    shard_start,
                    shard_end,
                    dtype,
                    self.device,
                )
            else:
                load_kwargs = {
                    "device_map": "auto" if uses_quantized_weights else None,
                    "dtype": dtype,
                    "low_cpu_mem_usage": True,
                    "cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None,
                }
                if quant_config is not None:
                    load_kwargs["quantization_config"] = quant_config
                self.model = AutoModelForCausalLM.from_pretrained(
                    load_source,
                    **load_kwargs,
                )
                if not uses_quantized_weights:
                    self.model.to(self.device)
        except Exception as exc:
            if _looks_like_oom(exc):
                raise InsufficientVRAMError(
                    f"insufficient VRAM to load {model_id} layers {shard_start}:{shard_end} "
                    f"with {quantization} quantization; choose a smaller shard or lower quantization"
                ) from exc
            raise

        self.model.eval()
        self.tokenizer = AutoTokenizer.from_pretrained(
            load_source,
            cache_dir=str(cache_dir) if cache_dir is not None and load_source == model_id else None,
        )
        self.layers = _model_layers(self.model)
        self.total_layers = len(self.layers)
        # shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention.
        if shard_end >= self.total_layers:
            raise ValueError(
                f"shard_end {shard_end} exceeds last layer index {self.total_layers - 1}"
            )
        self.is_head = shard_start == 0
        self.is_tail = shard_end >= self.total_layers - 1
        self.hidden_size = int(
            getattr(self.model.config, "hidden_size", 0)
            or getattr(self.model.config, "n_embd", 0)
        )
        self._embed_tokens = _embed_tokens(self.model) if self.is_head else None
        self._position_embeddings = _position_embeddings(self.model)
        self._norm = _final_norm(self.model) if self.is_tail else None
        self._lm_head = getattr(self.model, "lm_head", None) if self.is_tail else None

    def encode_prompt(self, prompt: str) -> TensorPayload:
        if not self.is_head or self._embed_tokens is None:
            raise ModelBackendError("text prompts can only be accepted by the head shard")
        encoded = self.tokenizer(prompt, return_tensors="pt")
        input_ids = encoded["input_ids"].to(self.device)
        attention_mask = encoded.get("attention_mask")
        if attention_mask is None:
            attention_mask = self.torch.ones_like(input_ids)
        attention_mask = attention_mask.to(self.device)
        position_ids = _position_ids(attention_mask, self.torch)
        hidden_states = self._embed_tokens(input_ids)
        if self._position_embeddings is not None:
            hidden_states = hidden_states + self._position_embeddings(position_ids)
        hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
        return self._payload(hidden_states, attention_mask, position_ids)

    def forward_bytes(
        self,
        body: bytes,
        shape: list[int],
        attention_mask_header: str | None,
        position_ids_header: str | None,
        start_layer: int | None = None,
    ) -> TensorPayload | str:
        hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
            self.device
        )
        attention_mask = _tensor_from_int64_header(
            attention_mask_header, self.torch, self.device
        )
        position_ids = _tensor_from_int64_header(
            position_ids_header, self.torch, self.device
        )
        hidden_states = self._run_layers(
            hidden_states, attention_mask, position_ids, start_layer=start_layer
        )
        if self.is_tail:
            return self.decode_tail(hidden_states)
        return self._payload(hidden_states, attention_mask, position_ids)

    def decode_tail(self, hidden_states: Any) -> str:
        if self._norm is not None:
            hidden_states = self._norm(hidden_states)
        if self._lm_head is None:
            raise ModelBackendError("tail shard has no lm_head")
        logits = self._lm_head(hidden_states)
        token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
        return self.tokenizer.decode([token_id], skip_special_tokens=True)

    def generate_text(
        self,
        messages: list[dict],
        max_new_tokens: int = 5120,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ) -> str:
        """Autoregressive generation using HF generate() — single-node (head+tail) mode."""
        if not self.is_head or not self.is_tail:
            raise ModelBackendError("local generation requires a full head+tail shard")
        encoded = self._encode_messages(messages)
        input_ids = encoded["input_ids"].to(self.device)
        attention_mask = encoded.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)
        pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
        do_sample = temperature != 1.0 or top_p != 1.0
        with self.torch.inference_mode():
            generated = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=max(1, int(max_new_tokens)),
                do_sample=do_sample,
                temperature=temperature if do_sample else None,
                top_p=top_p if do_sample else None,
                pad_token_id=pad_token_id,
            )
        new_tokens = generated[0, input_ids.shape[-1]:]
        return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

    def generate_text_streaming(
        self,
        messages: list[dict],
        max_new_tokens: int = 5000,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ):
        """Yield decoded token strings one at a time using HF TextIteratorStreamer."""
        if not self.is_head or not self.is_tail:
            raise ModelBackendError("streaming generation requires a full head+tail shard")
        import threading
        try:
            from transformers import TextIteratorStreamer  # type: ignore[import]
        except ImportError:
            yield self.generate_text(messages, max_new_tokens, temperature, top_p)
            return

        encoded = self._encode_messages(messages)
        input_ids = encoded["input_ids"].to(self.device)
        attention_mask = encoded.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)
        pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
        do_sample = temperature != 1.0 or top_p != 1.0

        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
        gen_kwargs = dict(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max(1, int(max_new_tokens)),
            do_sample=do_sample,
            temperature=temperature if do_sample else None,
            top_p=top_p if do_sample else None,
            pad_token_id=pad_token_id,
            streamer=streamer,
        )
        t = threading.Thread(target=self.model.generate, kwargs=gen_kwargs, daemon=True)
        t.start()
        for token_text in streamer:
            yield token_text
        t.join()

    def count_prompt_tokens(self, messages: list[dict]) -> int:
        """Return tokenizer-backed prompt token count for OpenAI usage metadata."""
        encoded = self._encode_messages(messages)
        input_ids = encoded["input_ids"]
        return int(input_ids.shape[-1])

    def count_text_tokens(self, text: str) -> int:
        """Return tokenizer-backed completion token count for OpenAI usage metadata."""
        try:
            encoded = self.tokenizer(
                text,
                return_tensors="pt",
                add_special_tokens=False,
            )
        except TypeError:
            encoded = self.tokenizer(text, return_tensors="pt")
        return int(encoded["input_ids"].shape[-1])

    def _encode_messages(self, messages: list[dict]) -> dict:
        """Format messages with chat template (if available) and tokenize."""
        if hasattr(self.tokenizer, "apply_chat_template"):
            try:
                prompt_str = self.tokenizer.apply_chat_template(
                    messages,
                    add_generation_prompt=True,
                    tokenize=False,
                )
                return dict(self.tokenizer(prompt_str, return_tensors="pt"))
            except Exception:
                pass
        prompt = " ".join(
            str(m.get("content", ""))
            for m in messages
            if isinstance(m, dict) and m.get("role") == "user"
        )
        return dict(self.tokenizer(prompt, return_tensors="pt"))

    def _run_layers(
        self,
        hidden_states: Any,
        attention_mask: Any,
        position_ids: Any,
        start_layer: int | None = None,
    ) -> Any:
        # start_layer overrides shard_start for overlapping-shard routing
        # (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
        # indexing outside the loaded weights.
        effective_start = (
            max(self.shard_start, start_layer)
            if start_layer is not None
            else self.shard_start
        )
        position_embeddings = _rotary_position_embeddings(
            self.model,
            hidden_states,
            position_ids,
        )
        layer_attention_mask = _decoder_attention_mask(
            attention_mask,
            hidden_states,
            self.torch,
        )
        with self.torch.inference_mode():
            for layer in self.layers[effective_start:self.shard_end + 1]:
                hidden_states = _call_layer(
                    layer,
                    hidden_states,
                    layer_attention_mask,
                    position_ids,
                    position_embeddings,
                )
            return hidden_states.to(self.torch.bfloat16)

    def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload:
        hidden_states = hidden_states.to(self.torch.bfloat16).contiguous()
        return TensorPayload(
            body=_tensor_to_bytes(hidden_states),
            shape=list(hidden_states.shape),
            attention_mask_header=_int_tensor_header(attention_mask)
            if attention_mask is not None
            else None,
            position_ids_header=_int_tensor_header(position_ids)
            if position_ids is not None
            else None,
        )


def load_torch_shard(
    model_id: str,
    shard_start: int,
    shard_end: int,
    quantization: Quantization = "auto",
    cache_dir: Path | None = None,
) -> TorchModelShard:
    return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)


def _total_layers_for_local_snapshot(auto_config: Any, load_source: str) -> int | None:
    snapshot_dir = Path(load_source)
    if not (snapshot_dir / "config.json").exists():
        return None
    from .model_catalog import layers_from_config

    try:
        cfg = auto_config.from_pretrained(str(snapshot_dir))
    except Exception:
        return None
    return layers_from_config(cfg)


def _should_partial_materialize_shard(
    load_source: str,
    shard_start: int,
    shard_end: int,
    *,
    total_layers_hint: int | None,
    uses_quantized_weights: bool,
) -> bool:
    if uses_quantized_weights:
        return False
    snapshot_dir = Path(load_source)
    if not snapshot_dir.exists() or not (snapshot_dir / "config.json").exists():
        return False
    if not (snapshot_dir / "model.safetensors.index.json").exists():
        return False
    if total_layers_hint is None:
        return False
    return not (shard_start == 0 and shard_end >= total_layers_hint - 1)


def _load_partial_model_from_snapshot(
    auto_config: Any,
    auto_model_for_causal_lm: Any,
    torch: Any,
    load_source: str,
    shard_start: int,
    shard_end: int,
    dtype: Any,
    device: Any,
    *,
    init_empty_weights_fn: Any | None = None,
    set_tensor_fn: Any | None = None,
    safe_open_fn: Any | None = None,
) -> Any:
    from .model_catalog import layers_from_config
    from .safetensors_selection import (
        INDEX_FILENAME,
        select_tensor_names_for_layers_from_index,
    )

    if init_empty_weights_fn is None:
        from accelerate import init_empty_weights as init_empty_weights_fn
    if set_tensor_fn is None:
        from accelerate.utils import set_module_tensor_to_device as set_tensor_fn
    if safe_open_fn is None:
        from safetensors import safe_open as safe_open_fn

    snapshot_dir = Path(load_source)
    cfg = auto_config.from_pretrained(str(snapshot_dir))
    total_layers = layers_from_config(cfg)
    if total_layers is None:
        raise PartialModelLoadUnsupported(
            f"could not determine num_hidden_layers for local snapshot {snapshot_dir}"
        )
    if shard_end >= total_layers:
        raise ValueError(
            f"shard_end {shard_end} exceeds last layer index {total_layers - 1}"
        )

    index_path = snapshot_dir / INDEX_FILENAME
    try:
        index = json.loads(index_path.read_text(encoding="utf-8"))
    except FileNotFoundError as exc:
        raise PartialModelLoadUnsupported(
            f"missing SafeTensors index for partial load: {index_path}"
        ) from exc
    weight_map = index.get("weight_map")
    if not isinstance(weight_map, dict):
        raise PartialModelLoadUnsupported(f"{INDEX_FILENAME} must contain a weight_map object")

    tensor_names = select_tensor_names_for_layers_from_index(
        weight_map,
        shard_start,
        shard_end,
        total_layers=total_layers,
    )
    if not tensor_names:
        raise PartialModelLoadUnsupported(
            f"no checkpoint tensors matched layers {shard_start}-{shard_end} in {snapshot_dir}"
        )

    with init_empty_weights_fn():
        model = auto_model_for_causal_lm.from_config(cfg, torch_dtype=dtype)
    tie_weights = getattr(model, "tie_weights", None)
    if callable(tie_weights):
        tie_weights()

    tensors_by_file: dict[str, list[str]] = {}
    for tensor_name in sorted(tensor_names):
        rel_file = weight_map.get(tensor_name)
        if not isinstance(rel_file, str):
            continue
        tensors_by_file.setdefault(rel_file, []).append(tensor_name)

    for rel_file, names in tensors_by_file.items():
        checkpoint_file = snapshot_dir / rel_file
        if not checkpoint_file.exists():
            raise PartialModelLoadUnsupported(
                f"checkpoint file advertised in {INDEX_FILENAME} is missing: {checkpoint_file}"
            )
        with safe_open_fn(str(checkpoint_file), framework="pt", device="cpu") as handle:
            for tensor_name in names:
                set_tensor_fn(
                    model,
                    tensor_name,
                    device,
                    value=handle.get_tensor(tensor_name),
                    dtype=dtype,
                )

    for module in _active_modules_for_shard(model, shard_start, shard_end):
        if hasattr(module, "to"):
            module.to(device)
    return model


def _model_load_plan(
    auto_config: Any,
    model_id: str,
    quantization: Quantization,
    torch: Any,
    cache_dir: Path | None = None,
) -> tuple[Any | None, Any, bool]:
    """Return (explicit quant config, dtype, uses quantized weights)."""
    if quantization != "auto":
        quant_config = build_quantization_config(quantization)
        return quant_config, torch.bfloat16, quant_config is not None

    cfg = auto_config.from_pretrained(
        model_id,
        cache_dir=str(cache_dir) if cache_dir is not None else None,
    )
    if _native_quantization_config(cfg) is not None:
        return None, _native_torch_dtype(cfg, torch), True
    return None, _native_torch_dtype(cfg, torch), False


def _config_candidates(cfg: Any) -> list[Any]:
    candidates = [cfg]
    get_text_config = getattr(cfg, "get_text_config", None)
    if callable(get_text_config):
        try:
            candidates.append(get_text_config())
        except Exception:
            pass
    text_config = getattr(cfg, "text_config", None)
    if text_config is not None:
        candidates.append(text_config)
    return candidates


def _native_quantization_config(cfg: Any) -> Any | None:
    for candidate in _config_candidates(cfg):
        quant_config = getattr(candidate, "quantization_config", None)
        if quant_config:
            return quant_config
    return None


def _native_torch_dtype(cfg: Any, torch: Any) -> Any:
    for candidate in _config_candidates(cfg):
        for attr in ("dtype", "torch_dtype"):
            dtype = getattr(candidate, attr, None)
            if dtype is None:
                continue
            if isinstance(dtype, str):
                dtype_name = dtype.removeprefix("torch.")
                dtype_value = getattr(torch, dtype_name, None)
                if dtype_value is not None:
                    return dtype_value
            else:
                return dtype
    return torch.bfloat16


def _model_layers(model: Any) -> Any:
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return model.model.layers
    if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
        return model.transformer.h
    raise ModelBackendError(
        "unsupported HuggingFace model architecture: no transformer layers found"
    )


def _embed_tokens(model: Any) -> Any:
    if hasattr(model, "model") and hasattr(model.model, "embed_tokens"):
        return model.model.embed_tokens
    if hasattr(model, "transformer") and hasattr(model.transformer, "wte"):
        return model.transformer.wte
    raise ModelBackendError(
        "unsupported HuggingFace model architecture: no token embeddings found"
    )


def _position_embeddings(model: Any) -> Any | None:
    if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
        return model.transformer.wpe
    return None


def _rotary_embedding_module(model: Any) -> Any | None:
    if hasattr(model, "model") and hasattr(model.model, "rotary_emb"):
        return model.model.rotary_emb
    if hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"):
        return model.transformer.rotary_emb
    return None


def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]:
    active: list[Any] = []

    def add(module: Any | None) -> None:
        if module is None:
            return
        if any(existing is module for existing in active):
            return
        active.append(module)

    if shard_start == 0:
        add(_embed_tokens(model))
        add(_position_embeddings(model))
    add(_rotary_embedding_module(model))
    for layer in _model_layers(model)[shard_start:shard_end + 1]:
        add(layer)
    total_layers = len(_model_layers(model))
    if shard_end >= total_layers - 1:
        add(_final_norm(model))
        add(getattr(model, "lm_head", None))
    return active


def _final_norm(model: Any) -> Any | None:
    if hasattr(model, "model") and hasattr(model.model, "norm"):
        return model.model.norm
    if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"):
        return model.transformer.ln_f
    return None


def _position_ids(attention_mask: Any, torch: Any) -> Any:
    position_ids = attention_mask.long().cumsum(-1) - 1
    return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long)


def _decoder_attention_mask(attention_mask: Any, hidden_states: Any, torch: Any) -> Any:
    """Build a causal additive mask for decoder layers called outside model.forward."""
    if attention_mask is None:
        return None
    if len(getattr(attention_mask, "shape", ())) != 2:
        return attention_mask
    batch_size, seq_len = attention_mask.shape
    if seq_len <= 1:
        return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype)

    min_value = torch.finfo(hidden_states.dtype).min
    causal = torch.full(
        (seq_len, seq_len),
        min_value,
        dtype=hidden_states.dtype,
        device=hidden_states.device,
    )
    causal = torch.triu(causal, diagonal=1)
    causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone()

    padding = attention_mask.to(device=hidden_states.device)
    if not bool(padding.all()):
        causal = causal.masked_fill(padding[:, None, None, :] == 0, min_value)
    return causal


def _rotary_position_embeddings(model: Any, hidden_states: Any, position_ids: Any) -> Any | None:
    """Return model-level rotary embeddings required by newer HF decoder layers."""
    if position_ids is None:
        return None
    rotary = _rotary_embedding_module(model)
    if rotary is None:
        return None
    return rotary(hidden_states, position_ids)


def _call_layer(
    layer: Any,
    hidden_states: Any,
    attention_mask: Any,
    position_ids: Any,
    position_embeddings: Any | None = None,
) -> Any:
    attempts = (
        {
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "position_embeddings": position_embeddings,
            "use_cache": False,
        },
        {
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "use_cache": False,
        },
        {"attention_mask": attention_mask, "use_cache": False},
        {"use_cache": False},
        {},
    )
    last_exc: Exception | None = None
    for kwargs in attempts:
        filtered = {key: value for key, value in kwargs.items() if value is not None}
        try:
            output = layer(hidden_states, **filtered)
            return output[0] if isinstance(output, tuple) else output
        except TypeError as exc:
            last_exc = exc
    if last_exc is not None:
        raise last_exc
    return layer(hidden_states)[0]


def _tensor_to_bytes(tensor: Any) -> bytes:
    import torch

    return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()


def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any:
    tensor = torch.frombuffer(bytearray(body), dtype=torch.bfloat16)
    return tensor.reshape(shape)


def _int_tensor_header(tensor: Any) -> str:
    data = tensor.detach().cpu().long().contiguous()
    raw = data.numpy().tobytes()
    shape = ",".join(str(dim) for dim in data.shape)
    encoded = base64.b64encode(raw).decode("ascii")
    return f"{shape}:{encoded}"


def _tensor_from_int64_header(value: str | None, torch: Any, device: Any) -> Any | None:
    if not value:
        return None
    shape_text, encoded = value.split(":", 1)
    shape = [int(part) for part in shape_text.split(",") if part]
    raw = base64.b64decode(encoded.encode("ascii"))
    return torch.frombuffer(bytearray(raw), dtype=torch.int64).reshape(shape).to(device)


def _looks_like_oom(exc: BaseException) -> bool:
    current: BaseException | None = exc
    while current is not None:
        text = str(current).lower()
        if "out of memory" in text or "cuda error: out of memory" in text:
            return True
        current = current.__cause__ or current.__context__
    return False