"""``boot_inspect`` — wrap an ``inspect_ai`` provider in a RemoteBridge via InspectDriver."""
from __future__ import annotations

import logging
import warnings
from typing import Any, Optional

import torch

from transformer_lens.factories.architecture_adapter_factory import (
    ArchitectureAdapterFactory,
)
from transformer_lens.model_bridge.remote_bridge import RemoteBridge
from transformer_lens.model_bridge.sources._bridge_builder import (
    build_bridge_config_from_hf,
    configure_tokenizer,
    skip_tokenizer_for_modality,
)
from transformer_lens.model_bridge.sources._hf_format import (
    determine_architecture_from_hf_config,
)
from transformer_lens.utilities.hf_utils import get_hf_token

from . import profiles
from .driver import InspectDriver

# Providers that expose the structural self-check + capture wire format the InspectDriver
# consumes. boot_inspect queries supported_kinds on these; others route via for_provider.
_TL_BRIDGE_PROVIDERS = {"tl_bridge", "tl_bridge_vllm"}


def boot_inspect(
    model_name: str,
    tokenizer: Optional[Any] = None,
    dtype: Optional[torch.dtype] = None,
    provider: str = "tl_bridge",
    **inspect_kwargs: Any,
) -> RemoteBridge:
    """Boot a model via an ``inspect_ai`` provider and wrap it in a :class:`RemoteBridge`.

    The driver is provider-agnostic: ``provider`` defaults to our own HF-backed
    ``tl_bridge`` provider (residual/attn/mlp capture + full affine interventions +
    full-sequence logits); ``"vllm-lens"`` targets a running vllm-lens vLLM provider
    (residual-only, additive-steering-only) — wire-aligned with its documented format,
    but not yet verified against a live provider.

    Fireable hooks (``tl_bridge``, TransformerBridge-native names): ``blocks.{i}.hook_in``
    (resid_pre) / ``ln2.hook_in`` (resid_mid) /
    ``hook_out`` (resid_post) / ``attn.hook_out`` / ``mlp.hook_out``, plus the head-split
    attention hooks where the structural probe finds them: ``attn.hook_q/k/v`` (pre-RoPE
    projection outputs; separate-projection archs only — fused qkv gates them),
    ``attn.hook_z`` (out-projection input), and ``attn.hook_pattern`` (post-softmax,
    capture-only, eager attention required). The provider runs a structural self-check per
    model and gates any boundary it can't serve faithfully: ``resid_mid`` for
    parallel-residual or norm-variant blocks, ``attn_out``/``mlp_out`` when their submodule
    isn't locatable (it warns when it gates one). ``embed``, ``ln_final``, and
    ``attn.hook_attn_scores`` are always non-fireable — use ``boot_transformers()`` for
    those.

    For parity with ``boot_transformers`` the provider loads with the same dtype (fp32 by
    default) and eager attention. Full-sequence logits ride on ``return_logits=True`` (the
    default); pass ``return_logits=False`` to skip the (seq × d_vocab) payload for pure
    activation capture (``run_with_cache`` keeps them since it returns logits).
    """
    from inspect_ai.model import get_model
    from transformers import AutoConfig, AutoTokenizer

    from . import (  # noqa: F401 — import registers @modelapi
        transformers_provider as _provider,
    )

    hf_token = get_hf_token()
    hf_config = AutoConfig.from_pretrained(model_name, token=hf_token)
    # Shared resolution (not architectures[0]) handles architectures=None configs via
    # model_type and rejects unsupported archs before the provider loads weights.
    architecture = determine_architecture_from_hf_config(hf_config)
    # Default fp32 to match boot_transformers (which loads/casts fp32 regardless of the
    # config's native dtype); an explicit dtype still wins.
    resolved_dtype = dtype if dtype is not None else torch.float32

    bridge_config = build_bridge_config_from_hf(hf_config, architecture, model_name, resolved_dtype)
    adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_config)
    if tokenizer is None and not skip_tokenizer_for_modality(adapter.cfg):
        tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
    if tokenizer is not None:
        # Match boot_transformers' tokenizer setup so to_tokens(str) is token-identical.
        tokenizer = configure_tokenizer(tokenizer, adapter.cfg)

    if provider == "tl_bridge":
        # The provider's raw HF forward must match boot_transformers' load: same dtype,
        # eager attention (TL forces eager — SDPA/flash diverge and accumulate with depth),
        # and auth/remote-code so gated/custom models load at all.
        inspect_kwargs["model_kwargs"] = _provider_model_kwargs(
            dict(inspect_kwargs.get("model_kwargs", {})), adapter, resolved_dtype, hf_token
        )
    elif provider == "tl_bridge_vllm":
        # Otherwise the provider defaults to the HF-config dtype and bridge_config.dtype
        # lies about what the engine actually loaded.
        inspect_kwargs["dtype"] = resolved_dtype

    # memoize=False: inspect_ai caches get_model by name, which would (a) return a stale
    # model ignoring a changed dtype/kwargs on re-boot and (b) keep weights resident past
    # close(). Each boot must honor its own args and own its model's lifecycle.
    model = get_model(f"{provider}/{model_name}", memoize=False, **inspect_kwargs)
    # Both TL-bridge providers (HF + vLLM) restrict their profile to the boundaries the
    # provider's structural self-check / overlay found this model can serve; warn only
    # if it gated something. Third-party providers (e.g. vllm-lens) route via for_provider.
    if provider in _TL_BRIDGE_PROVIDERS:
        api = getattr(model, "api", None)
        kinds = None
        note = ""
        # Default True for back-compat; vLLM provider sets False so RemoteBridge.forward
        # rejects loss/both (otherwise loss over -inf earlier positions silently NaNs).
        psl = True
        if api is not None:
            kinds = api.supported_kinds() if hasattr(api, "supported_kinds") else None
            note = api.capability_note() if hasattr(api, "capability_note") else ""
            psl = bool(getattr(api, "provides_sequence_logits", True))
        profile = profiles.TLBridgeProfile(supported_kinds=kinds, provides_sequence_logits=psl)
        if note:
            warnings.warn(note, UserWarning, stacklevel=2)
    else:
        profile = profiles.for_provider(provider)
    driver = InspectDriver(model=model, adapter=adapter, tokenizer=tokenizer, profile=profile)
    bridge = RemoteBridge(adapter=adapter, tokenizer=tokenizer, driver=driver)
    _log_hook_summary(model_name, architecture, provider, driver)
    return bridge


def _provider_model_kwargs(
    model_kwargs: dict[str, Any], adapter: Any, dtype: torch.dtype, hf_token: Optional[str]
) -> dict[str, Any]:
    """Load kwargs for the HF provider that mirror boot_transformers, so the provider's
    raw forward matches the bridge. Caller-supplied keys win (setdefault)."""
    model_kwargs.setdefault("torch_dtype", dtype)
    # boot_transformers forces eager unless the adapter pins an implementation.
    model_kwargs.setdefault(
        "attn_implementation", getattr(adapter.cfg, "attn_implementation", None) or "eager"
    )
    if hf_token:
        model_kwargs.setdefault("token", hf_token)
    return model_kwargs


def _log_hook_summary(
    model_name: str, architecture: str, provider: str, driver: InspectDriver
) -> None:
    log = logging.getLogger("transformer_lens.inspect")
    fireable = sorted(driver.supported_hook_points)
    log.info(
        "Inspect source on %s (%s) via provider %r serves %d fireable hook(s).",
        model_name,
        architecture,
        provider,
        len(fireable),
    )


__all__ = ["boot_inspect"]