# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""Extended import-smoke + API surface checks for unsloth + unsloth-zoo
modules under the existing CUDA spoof harness.

Where `tests/vllm_compat/test_unsloth_zoo_imports.py` covers the
narrow "must import on a vllm-less runner" claim for 5 modules,
this file walks the FULL set of modules our public surface depends
on. Catches:

  - module-level imports that break on a fresh transformers / peft /
    bnb release (the symbol pinned at import time is gone)
  - feature flags / gates that flip under the spoof (e.g. _IS_MLX
    silently activating on a non-Mac CI box)
  - public API surface drift: sorted `dir()` of each FastModel class
    is dumped and asserted-stable across runs (a removed kwarg here
    is a notebook regression we want to catch)

CPU-only. Inherits the same _zoo_aggressive_cuda_spoof harness as
test_unsloth_zoo_imports.py.
"""

from __future__ import annotations

import importlib
import importlib.machinery
import importlib.util
import inspect
import os
import sys
import types
from pathlib import Path

import pytest


# Apply the spoof BEFORE any unsloth-touching import.
_SPOOF_DIR = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_SPOOF_DIR))
import _zoo_aggressive_cuda_spoof as _spoof  # noqa: E402

_spoof.apply()


# Stub modules the unsloth import path may probe but that aren't
# installed on a CPU-only runner. Mirrors test_unsloth_zoo_imports.py.
def _stub_module(name: str, attrs: dict | None = None) -> None:
    """Stub a missing optional dep. Sets __spec__ so importlib.util's
    `find_spec(name)` doesn't raise `ValueError: __spec__ is None`,
    which torch / transformers / torchcodec callers hit otherwise."""
    if name in sys.modules:
        return
    m = types.ModuleType(name)
    # Minimal viable spec so importlib treats the stub as a real module.
    m.__spec__ = importlib.machinery.ModuleSpec(name = name, loader = None, origin = "<test stub>")
    for k, v in (attrs or {}).items():
        setattr(m, k, v)
    sys.modules[name] = m


_stub_module(
    "pynvml",
    {
        "nvmlInit": lambda: None,
        "nvmlShutdown": lambda: None,
        "nvmlDeviceGetCount": lambda: 1,
        "nvmlDeviceGetHandleByIndex": lambda i: object(),
        "nvmlDeviceGetMemoryInfo": lambda h: type(
            "_M",
            (),
            {"total": 80 * 1024**3, "free": 70 * 1024**3, "used": 10 * 1024**3},
        )(),
    },
)
_stub_module("torchcodec")


@pytest.fixture(autouse = True)
def _torch_distributed_safe(monkeypatch):
    """unsloth_zoo modules occasionally probe torch.distributed."""
    try:
        import torch.distributed as dist

        monkeypatch.setattr(dist, "is_available", lambda: True, raising = False)
        monkeypatch.setattr(dist, "is_initialized", lambda: False, raising = False)
        monkeypatch.setattr(dist, "get_world_size", lambda *a, **k: 1, raising = False)
        monkeypatch.setattr(dist, "get_rank", lambda *a, **k: 0, raising = False)
    except Exception:
        pass


def _has_unsloth_zoo() -> bool:
    return importlib.util.find_spec("unsloth_zoo") is not None


def _has_unsloth() -> bool:
    return importlib.util.find_spec("unsloth") is not None


# Extended unsloth-zoo module list. Modules with no top-level vllm/CUDA
# import are expected to load cleanly on a CPU spoof runner.


_ZOO_VLLM_FREE_MODULES = [
    "unsloth_zoo.compiler",
    "unsloth_zoo.compiler_replacements",
    "unsloth_zoo.dataset_utils",
    "unsloth_zoo.device_type",
    "unsloth_zoo.empty_model",
    "unsloth_zoo.gradient_checkpointing",
    "unsloth_zoo.hf_utils",
    "unsloth_zoo.llama_cpp",
    "unsloth_zoo.logging_utils",
    "unsloth_zoo.loss_utils",
    "unsloth_zoo.patching_utils",
    "unsloth_zoo.patch_torch_functions",
    "unsloth_zoo.peft_utils",
    "unsloth_zoo.rl_replacements",
    "unsloth_zoo.saving_utils",
    "unsloth_zoo.tiled_mlp",
    "unsloth_zoo.tokenizer_utils",
    "unsloth_zoo.training_utils",
    "unsloth_zoo.utils",
    "unsloth_zoo.vision_utils",
]


@pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed")
@pytest.mark.parametrize("modname", _ZOO_VLLM_FREE_MODULES)
def test_unsloth_zoo_module_imports_under_spoof(modname: str):
    """Each unsloth_zoo module must import cleanly on a CPU-only spoof runner.
    Catches transformers/peft/bnb symbol drift that fails at import time."""
    # Force fresh resolution: drop stale partial-import state from a prior failure
    sys.modules.pop(modname, None)
    try:
        importlib.import_module(modname)
    except Exception as e:
        pytest.fail(
            f"{modname} failed to import under CUDA spoof: " f"{type(e).__name__}: {str(e)[:300]}"
        )


# Spoof correctness: _IS_MLX must remain False on a non-Mac runner
# and _IS_CUDA / DEVICE_TYPE must reflect the spoofed CUDA layer.


@pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
def test_unsloth_is_mlx_false_under_spoof():
    """The CUDA spoof should not flip the MLX flag on a Linux/Windows CI
    box (real Apple Silicon is the ONLY environment _IS_MLX activates)."""
    sys.modules.pop("unsloth", None)
    import unsloth

    assert unsloth._IS_MLX is False, (
        f"_IS_MLX activated on a non-Apple-Silicon runner under CUDA spoof; "
        f"the MLX gate logic in unsloth/__init__.py is too lax"
    )


# unsloth.models.* — core RL + sentence-transformer surfaces, loaded
# transitively by `from unsloth import FastLanguageModel`.


_UNSLOTH_CORE_MODULES = [
    "unsloth.models.rl",
    "unsloth.models.rl_replacements",
    "unsloth.models.sentence_transformer",
    "unsloth.models._utils",
    "unsloth.models.loader",
    "unsloth.models.loader_utils",
    "unsloth.models.mapper",
]


@pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
@pytest.mark.parametrize("modname", _UNSLOTH_CORE_MODULES)
def test_unsloth_core_module_imports_under_spoof(modname: str):
    """Core unsloth modules must import on a CPU-only runner under the CUDA
    spoof; module-top symbol drift crashes here before any user-visible call.

    Bootstraps via `import unsloth` first, since most sub-modules need the
    package's _gpu_init side effects (else an import-guard fires)."""
    try:
        import unsloth  # noqa: F401  -- triggers _gpu_init side effects
    except Exception as e:
        pytest.skip(f"`import unsloth` failed under spoof: {e}")
    sys.modules.pop(modname, None)
    try:
        importlib.import_module(modname)
    except OSError as e:
        # `OSError: could not get source code` happens when an editable
        # install + frozen sub-import combine; that's an environment
        # quirk, not a symbol-drift bug. Skip rather than false-fail.
        pytest.skip(f"{modname} env issue: {e!s}")
    except Exception as e:
        pytest.fail(
            f"{modname} failed to import under CUDA spoof: " f"{type(e).__name__}: {str(e)[:300]}"
        )


# Public API surface for FastLanguageModel / FastVisionModel / FastModel under
# spoof: surface must be non-empty and the notebook-relied methods present.


@pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
def test_fast_model_class_surface_under_spoof():
    sys.modules.pop("unsloth", None)
    import unsloth

    found_at_least_one = False
    for cls_name in ("FastLanguageModel", "FastVisionModel", "FastModel"):
        cls = getattr(unsloth, cls_name, None)
        if cls is None:
            continue
        found_at_least_one = True
        public = sorted(n for n in dir(cls) if not n.startswith("_"))
        # Notebooks rely on these methods; loss of any one is a regression
        for method in ("from_pretrained", "get_peft_model"):
            assert method in public, (
                f"unsloth.{cls_name}.{method} missing under spoof; "
                f"every Colab notebook calling it breaks"
            )
    assert found_at_least_one, (
        f"none of FastLanguageModel/FastVisionModel/FastModel reachable "
        f"on `unsloth` package root"
    )


# RL surface drill-down: GRPO/SFT/DPO classes must be reachable and the
# source-rewriter dispatch table populated. Catches rl_replacements importing
# cleanly while RL_FUNCTIONS / RL_REPLACEMENTS is silently empty.


@pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
def test_unsloth_rl_replacements_dispatch_populated():
    try:
        import unsloth  # noqa: F401  -- _gpu_init bootstrap
    except Exception as e:
        pytest.skip(f"`import unsloth` failed under spoof: {e}")
    sys.modules.pop("unsloth.models.rl_replacements", None)
    try:
        rl = importlib.import_module("unsloth.models.rl_replacements")
    except OSError as e:
        pytest.skip(f"env issue importing rl_replacements: {e!s}")
    funcs = getattr(rl, "RL_FUNCTIONS", None)
    if funcs is None:
        pytest.skip("RL_FUNCTIONS attribute not present (architecture changed; check)")
    assert isinstance(funcs, dict), f"RL_FUNCTIONS expected dict, got {type(funcs).__name__}"
    # The trainer types unsloth-zoo dispatches against MUST be keys.
    for key in ("grpo_trainer", "sft_trainer", "dpo_trainer"):
        assert key in funcs, (
            f"RL_FUNCTIONS missing dispatch key '{key}'; "
            f"unsloth_zoo source rewrites silently no-op"
        )
        assert (
            isinstance(funcs[key], list) and len(funcs[key]) > 0
        ), f"RL_FUNCTIONS[{key!r}] is empty list; rewrites no-op"


# unsloth-zoo compiler test_apply_fused_lm_head — the fused-LM-head emit path
# (named test in compiler.py:1983); we just confirm it's callable.


@pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed")
def test_zoo_compiler_apply_fused_lm_head_callable():
    sys.modules.pop("unsloth_zoo.compiler", None)
    compiler = importlib.import_module("unsloth_zoo.compiler")
    fn = getattr(compiler, "test_apply_fused_lm_head", None)
    assert fn is not None and callable(fn), (
        f"unsloth_zoo.compiler.test_apply_fused_lm_head missing or non-callable; "
        f"the in-file CPU regression test is the only fused-LM-head coverage"
    )


# Spot-check signature stability of FastModel.from_pretrained — every notebook
# call site relies on these kwargs; a removal becomes silent positional drift.


@pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
def test_fast_model_from_pretrained_kwargs_under_spoof():
    sys.modules.pop("unsloth", None)
    import unsloth

    cls = getattr(unsloth, "FastLanguageModel", None) or getattr(unsloth, "FastModel", None)
    if cls is None:
        pytest.skip("FastLanguageModel/FastModel not exported")
    fn = getattr(cls, "from_pretrained", None)
    if fn is None:
        pytest.skip("from_pretrained not on class (might be classmethod stub)")
    try:
        params = list(inspect.signature(fn).parameters)
    except (TypeError, ValueError):
        pytest.skip("from_pretrained signature not introspectable")
    # Notebooks use these by name everywhere.
    for kwarg in ("model_name", "max_seq_length", "load_in_4bit"):
        assert kwarg in params, (
            f"FastLanguageModel.from_pretrained missing kwarg `{kwarg}`; "
            f"every Colab notebook breaks at the install cell"
        )