"""Opt-in loader for model-local MLX-VLM architecture files.

MLX-VLM 0.6.4 does not yet consult ``vlm_model_file`` in a model config.
This module adds that one lookup in-process.  It never edits the installed
``mlx_vlm`` package, and it only executes model-local code when the caller has
explicitly passed ``trust_remote_code=True``.
"""

from __future__ import annotations

import importlib.util
import sys
from contextvars import ContextVar
from pathlib import Path
from types import ModuleType
from typing import Any


_ACTIVE_MODEL_PATH: ContextVar[Path | None] = ContextVar(
    "mlx_vlm_model_file_path", default=None
)
_ACTIVE_TRUST_REMOTE_CODE: ContextVar[bool] = ContextVar(
    "mlx_vlm_model_file_trust", default=False
)
_MODULE_CACHE: dict[Path, ModuleType] = {}


def _load_local_module(model_path: Path, filename: str) -> ModuleType:
    root = model_path.resolve()
    module_path = (root / filename).resolve()

    try:
        module_path.relative_to(root)
    except ValueError as exc:
        raise ValueError(
            f"vlm_model_file must stay inside the model directory: {filename!r}"
        ) from exc

    if module_path.suffix != ".py" or not module_path.is_file():
        raise FileNotFoundError(f"Model-local MLX-VLM runtime not found: {module_path}")

    cached = _MODULE_CACHE.get(module_path)
    if cached is not None:
        return cached

    module_name = f"mlx_vlm_remote_{abs(hash(str(module_path))):x}"
    spec = importlib.util.spec_from_file_location(module_name, module_path)
    if spec is None or spec.loader is None:
        raise ImportError(f"Could not import model-local runtime: {module_path}")

    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    try:
        spec.loader.exec_module(module)
    except Exception:
        sys.modules.pop(module_name, None)
        raise

    _MODULE_CACHE[module_path] = module
    return module


def install() -> None:
    """Install the model-file lookup into the current Python process."""

    import mlx_vlm.utils as utils

    if getattr(utils.get_model_and_args, "_model_file_loader_installed", False):
        return

    original_get_model_and_args = utils.get_model_and_args
    original_load_model = utils.load_model

    def get_model_and_args(config: dict, *args: Any, **kwargs: Any):
        filename = config.get("vlm_model_file")
        model_path = kwargs.get("model_path") or _ACTIVE_MODEL_PATH.get()

        # Calls made later by processor setup do not carry a model path in
        # MLX-VLM 0.6.4; let its native Qwen implementation handle those.
        if not filename or model_path is None:
            return original_get_model_and_args(config, *args, **kwargs)

        trusted = bool(
            kwargs.get("trust_remote_code", False)
            or _ACTIVE_TRUST_REMOTE_CODE.get()
        )
        if not trusted:
            raise PermissionError(
                "This checkpoint includes a model-local MLX-VLM runtime. "
                "Re-run with --trust-remote-code (or trust_remote_code=True)."
            )

        module = _load_local_module(Path(model_path), str(filename))
        return module, f"model-local:{filename}"

    def load_model(model_path: Path, lazy: bool = False, **kwargs: Any):
        path_token = _ACTIVE_MODEL_PATH.set(Path(model_path))
        trust_token = _ACTIVE_TRUST_REMOTE_CODE.set(
            bool(kwargs.get("trust_remote_code", False))
        )
        try:
            return original_load_model(model_path, lazy=lazy, **kwargs)
        finally:
            _ACTIVE_TRUST_REMOTE_CODE.reset(trust_token)
            _ACTIVE_MODEL_PATH.reset(path_token)

    get_model_and_args._model_file_loader_installed = True  # type: ignore[attr-defined]
    load_model._model_file_loader_installed = True  # type: ignore[attr-defined]
    utils.get_model_and_args = get_model_and_args
    utils.load_model = load_model


def load(path_or_hf_repo: str, **kwargs: Any):
    """Programmatic convenience wrapper around :func:`mlx_vlm.load`."""

    install()
    kwargs.setdefault("trust_remote_code", True)
    from mlx_vlm import load as mlx_vlm_load

    return mlx_vlm_load(path_or_hf_repo, **kwargs)