from __future__ import annotations

import importlib
import importlib.util
import os
from configs import cfg

def _false() -> bool:
    return False

def describe_exception_chain(exc: Exception) -> str:
    messages: list[str] = []
    seen: set[int] = set()
    current: BaseException | None = exc

    while current is not None and id(current) not in seen:
        seen.add(id(current))
        message = f"{type(current).__name__}: {current}"
        if message not in messages:
            messages.append(message)
        current = current.__cause__ or current.__context__

    return " -> ".join(messages)

def disable_flash_attn_for_transformers() -> None:
    try:
        import transformers.utils as tf_utils

        tf_utils.is_flash_attn_2_available = _false
        if hasattr(tf_utils, "is_flash_attn_3_available"):
            tf_utils.is_flash_attn_3_available = _false
    except Exception:
        pass

    try:
        from transformers.utils import import_utils as tf_import_utils

        tf_import_utils.is_flash_attn_2_available = _false
        if hasattr(tf_import_utils, "is_flash_attn_3_available"):
            tf_import_utils.is_flash_attn_3_available = _false
    except Exception:
        pass

    try:
        import transformers.modeling_utils as modeling_utils

        if hasattr(modeling_utils, "is_flash_attn_2_available"):
            modeling_utils.is_flash_attn_2_available = _false
        if hasattr(modeling_utils, "is_flash_attn_3_available"):
            modeling_utils.is_flash_attn_3_available = _false
    except Exception:
        pass

    try:
        import transformers.modeling_flash_attention_utils as flash_utils

        flash_utils.is_flash_attn_2_available = _false
        if hasattr(flash_utils, "is_flash_attn_3_available"):
            flash_utils.is_flash_attn_3_available = _false
    except Exception:
        pass

def resolve_attention_backend(logger) -> str:
    forced_backend = os.environ.get("QUINTUS_ATTENTION_BACKEND")
    if forced_backend:
        logger.info(f"  [ATTENTION] Forced backend via QUINTUS_ATTENTION_BACKEND={forced_backend!r}.")
        return forced_backend
    
    try:
        from transformers.utils import is_flash_attn_3_available
        if is_flash_attn_3_available():
            logger.info("  [ATTENTION] Using flash_attention_3.")
            return "flash_attention_3"
    except Exception:
        pass

    try:
        importlib.import_module("flash_attn")
        logger.info("  [ATTENTION] Using flash_attention_2.")
        return "flash_attention_2"
    except Exception as exc:
        if importlib.util.find_spec("flash_attn") is not None:
            disable_flash_attn_for_transformers()
            logger.warning(
                "flash-attn appears installed but failed to import (%s: %s); "
                "masking flash-attn from Transformers and falling back to sdpa.",
                type(exc).__name__,
                exc,
            )
        else:
            logger.info("  [ATTENTION] Using PyTorch SDPA.")
        return "sdpa"

def _requires_remote_code_opt_in(exc: Exception) -> bool:
    message = str(exc).lower()
    return (
        "trust_remote_code" in message
        or "requires you to execute the configuration file" in message
        or "requires remote code" in message
    )

def format_model_load_error(subject: str, exc: Exception) -> str:
    if not cfg.model.allow_remote_code and _requires_remote_code_opt_in(exc):
        return (
            f"{subject} failed because the selected model/tokenizer requires remote code, "
            "but Quintus is configured with allow_remote_code=false. Review the upstream "
            "repository and rerun with QUINTUS_ALLOW_REMOTE_CODE=1 only if you explicitly "
            "trust that code."
        )
    return f"{subject} failed: {describe_exception_chain(exc)}"