"""HuggingFace Transformers connector for AVP.

Wraps a HuggingFace PreTrainedModel + tokenizer to provide the EngineConnector
interface for latent communication.

Requires torch and transformers — uses lazy imports.
"""

import logging
from typing import Any, Dict, List, Optional, Tuple, Union

from ..context import AVPContext
from ..errors import EngineNotAvailableError, IncompatibleModelsError, RealignmentError
from ..handshake import compute_model_hash, extract_model_identity
from ..realign import (
    apply_realignment,
    compute_target_norm,
    get_or_compute_realignment,
    needs_realignment,
    project_to_embedding_space,
)
from ..types import ModelIdentity
from .base import EngineConnector, _render_prompt, _tokenize_prompt

logger = logging.getLogger(__name__)


def _require_deps():
    """Check that torch and transformers are available."""
    try:
        import torch
    except ImportError:
        raise EngineNotAvailableError(
            "huggingface (requires torch). Install with: pip install avp[hf]"
        )
    try:
        import transformers
    except ImportError:
        raise EngineNotAvailableError(
            "huggingface (requires transformers). Install with: pip install avp[hf]"
        )
    return torch, transformers


def _past_length(past_kv: Any) -> int:
    """Get sequence length from past_key_values."""
    if past_kv is None:
        return 0
    try:
        from transformers.cache_utils import Cache
        if isinstance(past_kv, Cache):
            return past_kv.get_seq_length()
    except ImportError:
        pass
    if isinstance(past_kv, (tuple, list)) and len(past_kv) > 0:
        first = past_kv[0]
        if isinstance(first, (tuple, list)) and len(first) > 0:
            return first[0].shape[-2]
    return 0


def _to_messages(prompt: Union[str, List[Dict[str, str]]]) -> List[Dict[str, str]]:
    """Normalize prompt to a list of chat message dicts."""
    if isinstance(prompt, str):
        return [{"role": "user", "content": prompt}]
    return prompt


class HuggingFaceConnector(EngineConnector):
    """Engine connector for HuggingFace Transformers models.

    Accepts either a pre-loaded model+tokenizer or loads from model_id.

    Example:
        >>> from transformers import AutoModelForCausalLM, AutoTokenizer
        >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B")
        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
        >>> connector = HuggingFaceConnector(model=model, tokenizer=tokenizer)
    """

    def __init__(
        self,
        model: Optional[Any] = None,
        tokenizer: Optional[Any] = None,
        model_id: Optional[str] = None,
        device: Optional[str] = None,
    ):
        torch, transformers = _require_deps()

        if model is not None and tokenizer is not None:
            self.model = model
            self.tokenizer = tokenizer
        elif model_id is not None:
            self.model = transformers.AutoModelForCausalLM.from_pretrained(model_id)
            self.tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
        else:
            raise ValueError("Provide either (model, tokenizer) or model_id")

        if device:
            self.device = device
            self.model = self.model.to(device)
        else:
            self.device = str(next(self.model.parameters()).device)

        # Ensure pad token (following LatentMAS: try eos first, then add <pad>)
        if self.tokenizer.pad_token_id is None:
            if self.tokenizer.eos_token_id is not None:
                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
            else:
                self.tokenizer.add_special_tokens({"pad_token": "<pad>"})
                self.model.resize_token_embeddings(len(self.tokenizer))

        self._model_hash = compute_model_hash(self.model.config.to_dict())
        self._identity = extract_model_identity(self.model, tokenizer=self.tokenizer)
        self._w_realign = None
        self._target_norm = None
        self._avp_map_cache: Dict[str, Any] = {}  # source_hash → AVPMap

    @classmethod
    def from_pretrained(
        cls,
        model_id: str,
        device: Optional[str] = None,
        dtype: Optional[Any] = None,
        **model_kwargs: Any,
    ) -> "HuggingFaceConnector":
        """Load a model from HuggingFace Hub with auto-detected device and dtype.

        Args:
            model_id: HuggingFace model identifier (e.g. 'Qwen/Qwen2.5-7B-Instruct').
            device: Target device. Auto-detects: cuda > mps > cpu.
            dtype: Torch dtype. Auto-detects: bfloat16 (cuda+bf16) > float16 (cuda) > float32.
            **model_kwargs: Extra kwargs passed to AutoModelForCausalLM.from_pretrained().

        Returns:
            Configured HuggingFaceConnector with model in eval mode.
        """
        torch, transformers = _require_deps()

        # Auto-detect device
        if device is None:
            if torch.cuda.is_available():
                device = "cuda"
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                device = "mps"
            else:
                device = "cpu"

        # Auto-detect dtype
        if dtype is None:
            if device == "cuda" and torch.cuda.is_bf16_supported():
                dtype = torch.bfloat16
            elif device == "cuda":
                dtype = torch.float16
            else:
                dtype = torch.float32

        # If user specified device_map in kwargs, skip explicit .to(device)
        has_device_map = "device_map" in model_kwargs

        model = transformers.AutoModelForCausalLM.from_pretrained(
            model_id, dtype=dtype, **model_kwargs
        )
        model.eval()
        tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)

        if has_device_map:
            # device_map handles placement; infer device from parameters
            return cls(model=model, tokenizer=tokenizer)
        return cls(model=model, tokenizer=tokenizer, device=device)

    def get_model_identity(self) -> ModelIdentity:
        return self._identity

    def extract_hidden_state(
        self,
        input_ids: Any,
        attention_mask: Optional[Any] = None,
        past_key_values: Optional[Any] = None,
    ) -> Tuple[Any, Any, Any]:
        """Run forward pass and extract hidden states + KV-cache.

        Returns:
            Tuple of (last_hidden_state [B, D], all_hidden_states, past_key_values).
        """
        import torch

        if input_ids.dim() != 2:
            raise ValueError("input_ids must be 2D [batch, seq_len]")

        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, device=self.device)
        else:
            attention_mask = attention_mask.to(self.device)

        # Prepend past attention if using KV-cache
        if past_key_values is not None:
            past_len = _past_length(past_key_values)
            if past_len > 0:
                past_mask = torch.ones(
                    (attention_mask.shape[0], past_len),
                    dtype=attention_mask.dtype,
                    device=attention_mask.device,
                )
                attention_mask = torch.cat([past_mask, attention_mask], dim=-1)

        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids.to(self.device),
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                use_cache=True,
                output_hidden_states=True,
                return_dict=True,
            )

        last_hidden = outputs.hidden_states[-1][:, -1, :]  # [B, D]
        return last_hidden, outputs.hidden_states, outputs.past_key_values

    def inject_and_generate(
        self,
        inputs_embeds: Any,
        attention_mask: Optional[Any] = None,
        past_key_values: Optional[Any] = None,
        max_new_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.95,
    ) -> Tuple[List[str], Any]:
        """Generate text from injected embeddings."""
        import torch

        if attention_mask is None:
            batch_size = inputs_embeds.shape[0]
            total_len = inputs_embeds.shape[1]
            if past_key_values is not None:
                total_len += _past_length(past_key_values)
            attention_mask = torch.ones(
                (batch_size, total_len), dtype=torch.long, device=self.device
            )

        with torch.no_grad():
            outputs = self.model.generate(
                inputs_embeds=inputs_embeds.to(self.device),
                attention_mask=attention_mask.to(self.device),
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=self.tokenizer.pad_token_id,
                return_dict_in_generate=True,
                past_key_values=past_key_values,
            )

        # Decode generated tokens (skip prompt tokens).
        # model.generate() sequences contain [input_len + generated_len] —
        # past KV-cache positions are NOT included in the output.
        prompt_len = inputs_embeds.shape[1]
        texts = []
        for seq in outputs.sequences:
            generated_ids = seq[prompt_len:]
            text = self.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
            texts.append(text)

        return texts, getattr(outputs, "past_key_values", None)

    def get_embedding_weights(self) -> Tuple[Any, Any]:
        input_embeds = self.model.get_input_embeddings()
        output_embeds = self.model.get_output_embeddings()
        if output_embeds is None:
            output_embeds = getattr(self.model, "lm_head", None)
        return (
            input_embeds.weight if input_embeds else None,
            output_embeds.weight if output_embeds else None,
        )

    def tokenize(self, text: str) -> Any:
        return self.tokenizer(
            text,
            add_special_tokens=False,
            return_tensors="pt",
        )["input_ids"].to(self.device)

    def needs_realignment(self) -> bool:
        return needs_realignment(self.model)

    def _ensure_realignment(self) -> Tuple[Any, Any]:
        """Load or compute realignment matrix."""
        if self._w_realign is None:
            self._w_realign, self._target_norm = get_or_compute_realignment(
                self.model, self._model_hash, device=self.device
            )
        return self._w_realign, self._target_norm

    def _ensure_target_norm(self) -> Any:
        """Get target norm, computing if needed.

        Target norm is needed for ALL models (tied and untied) because hidden
        states from the last transformer layer have different norms than input
        embeddings. LatentMAS always normalizes, even without the projection.
        """
        if self._target_norm is None:
            self._target_norm = compute_target_norm(self.model, device=self.device)
        return self._target_norm

    def generate_latent_steps(
        self,
        input_ids: Any,
        latent_steps: int,
        attention_mask: Optional[Any] = None,
        past_key_values: Optional[Any] = None,
        collect_hidden_states: bool = False,
    ) -> Any:
        """Run the LatentMAS latent generation loop.

        Performs `latent_steps` forward passes, each time:
        1. Extract hidden state from last token
        2. Apply realignment projection (untied models) or just normalize (tied models)
        3. Feed aligned hidden state as inputs_embeds for next step
        4. Accumulate KV-cache

        Normalization to target_norm is ALWAYS applied, even for tied-weight
        models. This matches LatentMAS behavior: hidden states from the last
        transformer layer have different norms than input embeddings, and
        injecting un-normalized vectors gives the model out-of-distribution inputs.

        Ported from LatentMAS models.py:276-350.

        Args:
            input_ids: Token IDs [batch, seq_len].
            latent_steps: Number of latent thinking steps.
            attention_mask: Optional attention mask.
            past_key_values: Optional prior KV-cache.
            collect_hidden_states: If True, collect raw hidden states (before
                realignment) at each step and return them alongside the KV-cache.

        Returns:
            If collect_hidden_states is False: updated past_key_values.
            If collect_hidden_states is True: tuple of (past_key_values,
                hidden_states tensor of shape [1 + latent_steps, D]).
        """
        import torch

        if input_ids.dim() != 2:
            raise ValueError("input_ids must be 2D [batch, seq_len]")

        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, device=self.device)
        else:
            attention_mask = attention_mask.to(self.device)

        # Prepend past attention
        if past_key_values is not None:
            past_len = _past_length(past_key_values)
            if past_len > 0:
                past_mask = torch.ones(
                    (attention_mask.shape[0], past_len),
                    dtype=attention_mask.dtype,
                    device=attention_mask.device,
                )
                attention_mask = torch.cat([past_mask, attention_mask], dim=-1)

        # Initial forward pass
        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids.to(self.device),
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                use_cache=True,
                output_hidden_states=True,
                return_dict=True,
            )

        past = outputs.past_key_values
        last_hidden = outputs.hidden_states[-1][:, -1, :]  # [B, D]

        collected_hidden = []
        if collect_hidden_states:
            collected_hidden.append(last_hidden.clone())

        # Determine realignment strategy
        do_realign = self.needs_realignment()
        if do_realign:
            w_realign, target_norm = self._ensure_realignment()
            embed_weight = None
        else:
            # Tied models: project hidden → logits → softmax → soft embedding.
            # Simple normalization doesn't work because hidden state directions
            # have very low cosine similarity (~0.24) to embedding vectors.
            target_norm = None
            embed_weight = self.model.get_input_embeddings().weight

        # Latent loop
        import numpy as np

        for step in range(latent_steps):
            if do_realign:
                latent_np = apply_realignment(last_hidden, w_realign, target_norm)
            else:
                # Tied models: project through vocabulary to get valid embedding
                latent_np = project_to_embedding_space(
                    last_hidden, embed_weight, temperature=1.0
                )

            # Convert numpy projection result back to torch for model forward
            latent_vec = torch.from_numpy(
                np.ascontiguousarray(latent_np)
            ).to(device=self.device, dtype=self.model.dtype)
            latent_embed = latent_vec.unsqueeze(1)  # [B, 1, D]

            # Build attention mask for this step
            past_len = _past_length(past)
            latent_mask = torch.ones(
                (latent_embed.shape[0], past_len + 1),
                dtype=torch.long,
                device=self.device,
            )

            with torch.no_grad():
                outputs = self.model(
                    inputs_embeds=latent_embed,
                    attention_mask=latent_mask,
                    past_key_values=past,
                    use_cache=True,
                    output_hidden_states=True,
                    return_dict=True,
                )

            past = outputs.past_key_values
            last_hidden = outputs.hidden_states[-1][:, -1, :]

            if collect_hidden_states:
                collected_hidden.append(last_hidden.clone())

            # Early exit on NaN — continuing would only accumulate garbage KV entries
            if torch.isnan(last_hidden).any():
                logger.warning(
                    "NaN in hidden state at latent step %d/%d — stopping early",
                    step + 1, latent_steps,
                )
                break

        if collect_hidden_states:
            # Each element is [B, D] where B=1 in practice.
            # cat along dim=0 gives [1 + latent_steps, D] for B=1.
            return past, torch.cat(collected_hidden, dim=0)

        return past

    # --- High-level API ---

    @property
    def can_think(self) -> bool:
        return True

    def think(
        self,
        prompt: Union[str, List[Dict[str, str]]],
        steps: int = 20,
        context: Optional[AVPContext] = None,
        _diagnostics: Optional[Any] = None,
    ) -> AVPContext:
        """Generate latent context via thinking steps.

        Args:
            prompt: A string (wrapped as user message) or list of chat messages.
            steps: Number of latent thinking steps (must be >= 1).
            context: Optional AVPContext from a prior think() to continue from.
            _diagnostics: Internal. Pre-created TransferDiagnostics to populate.

        Returns:
            AVPContext with accumulated KV-cache.

        Raises:
            ValueError: If steps < 1.
            IncompatibleModelsError: If context is from a different model.
        """
        if steps < 1:
            raise ValueError(f"steps must be >= 1, got {steps}")

        messages = _to_messages(prompt)
        prompt_text = _render_prompt(self.tokenizer, messages)
        input_ids, attention_mask = _tokenize_prompt(
            self.tokenizer, prompt_text, self.device
        )

        past_kv = None
        accumulated_steps = 0
        if context is not None:
            if context.model_hash != self._model_hash:
                raise IncompatibleModelsError(
                    f"Context model_hash {context.model_hash!r} does not match "
                    f"this connector's model_hash {self._model_hash!r}. "
                    "Use the same model or serialize via to_bytes() for cross-model transfer."
                )
            past_kv = context.past_key_values
            accumulated_steps = context.num_steps

        past_kv, hidden_states = self.generate_latent_steps(
            input_ids, steps, attention_mask=attention_mask,
            past_key_values=past_kv, collect_hidden_states=True,
        )
        last_hidden = hidden_states[-1].unsqueeze(0)  # [1, D]

        # Always-on: check for NaN/Inf in hidden states
        import warnings as _warnings
        import torch
        has_nan = bool(torch.isnan(hidden_states).any())
        has_inf = bool(torch.isinf(hidden_states).any())
        if has_nan:
            _warnings.warn(
                "AVP: NaN detected in hidden states during think(). "
                "Downstream generation will likely produce garbage.",
                RuntimeWarning, stacklevel=2,
            )
        if has_inf:
            _warnings.warn(
                "AVP: Inf detected in hidden states during think(). "
                "Downstream generation will likely produce garbage.",
                RuntimeWarning, stacklevel=2,
            )

        if _diagnostics is not None:
            _diagnostics.has_nan = has_nan
            _diagnostics.has_inf = has_inf
            _diagnostics.transfer_mode = "latent"
            _diagnostics.prompt_tokens = input_ids.shape[-1]

        # Norm trajectory (when diagnostics are being collected)
        if _diagnostics is not None:
            norms = [float(hidden_states[i].norm()) for i in range(hidden_states.shape[0])]
            _diagnostics.norm_trajectory = norms

        return AVPContext(
            past_key_values=past_kv,
            model_hash=self._model_hash,
            num_steps=accumulated_steps + steps,
            seq_len=_past_length(past_kv),
            model_family=self._identity.model_family,
            hidden_dim=self._identity.hidden_dim,
            num_layers=self._identity.num_layers,
            last_hidden_state=last_hidden,
        )

    def generate(
        self,
        prompt: Union[str, List[Dict[str, str]]],
        context: Optional[AVPContext] = None,
        source: Optional["HuggingFaceConnector"] = None,
        cross_model: bool = False,
        max_new_tokens: int = 512,
        temperature: float = 0.7,
        top_p: float = 0.95,
        do_sample: bool = True,
        _diagnostics: Optional[Any] = None,
    ) -> str:
        """Generate text, optionally conditioned on latent context from think().

        Args:
            prompt: A string (wrapped as user message) or list of chat messages.
            context: Optional AVPContext to condition generation on.
            source: Optional source HuggingFaceConnector for cross-model projection.
                When provided and context is from a different model, automatically
                handles calibration, projection, and KV-cache priming.
                Requires ``cross_model=True``.
            cross_model: Must be True to enable cross-model projection.
                Cross-model (Rosetta Stone) is experimental. Default False.
            max_new_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.
            top_p: Nucleus sampling threshold.
            do_sample: Whether to use sampling (True) or greedy decoding (False).
            _diagnostics: Internal. Pre-created TransferDiagnostics to populate.

        Returns:
            Generated text string.

        Raises:
            IncompatibleModelsError: If context is from a different model
                and no source connector is provided.
        """
        import torch

        # Cross-model: auto-project when source connector is provided
        if (
            source is not None
            and context is not None
            and context.model_hash != self._model_hash
        ):
            if not cross_model:
                import warnings as _warnings
                _warnings.warn(
                    "Cross-model projection (Rosetta Stone) is experimental. "
                    "Falling back to text-only generation (ignoring cross-model context). "
                    "Pass cross_model=True to connector.generate() to enable latent transfer.",
                    UserWarning,
                    stacklevel=2,
                )
                context = None
                source = None
            else:
                context = self._apply_rosetta_projection(
                    source, context, _diagnostics=_diagnostics,
                )

        messages = _to_messages(prompt)
        prompt_text = _render_prompt(self.tokenizer, messages)
        input_ids, attention_mask = _tokenize_prompt(
            self.tokenizer, prompt_text, self.device
        )

        past_kv = None
        cache_position = None

        # prompt_len must be computed BEFORE extending attention_mask with
        # past KV-cache entries.  model.generate() returns sequences shaped
        # [input_ids_len + generated_len] — past positions are NOT included.
        prompt_len = attention_mask.sum(dim=1).tolist()[0]

        if context is not None:
            if context.model_hash != self._model_hash:
                raise IncompatibleModelsError(
                    f"Context model_hash {context.model_hash!r} does not match "
                    f"this connector's model_hash {self._model_hash!r}. "
                    "Pass source= with the source connector for cross-model projection."
                )
            past_kv = context.past_key_values
            past_len = _past_length(past_kv)

            cache_position = torch.arange(
                past_len,
                past_len + input_ids.shape[-1],
                dtype=torch.long,
                device=self.device,
            )
            if past_len > 0:
                past_mask = torch.ones(
                    (attention_mask.shape[0], past_len),
                    dtype=attention_mask.dtype,
                    device=attention_mask.device,
                )
                attention_mask = torch.cat([past_mask, attention_mask], dim=-1)

        with torch.no_grad():
            gen_kwargs = dict(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                do_sample=do_sample,
                pad_token_id=self.tokenizer.pad_token_id,
                return_dict_in_generate=True,
            )
            if do_sample:
                gen_kwargs["temperature"] = temperature
                gen_kwargs["top_p"] = top_p
            if past_kv is not None:
                gen_kwargs["past_key_values"] = past_kv
                gen_kwargs["cache_position"] = cache_position
            outputs = self.model.generate(**gen_kwargs)

        generated_ids = outputs.sequences[0, prompt_len:]
        text = self.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

        # Always-on: warn on empty output
        if not text:
            import warnings as _warnings
            _warnings.warn(
                f"AVP: generate() produced empty output (prompt_len={prompt_len}, "
                f"context={'yes' if context is not None else 'no'}). "
                "Check max_new_tokens and input prompt.",
                RuntimeWarning, stacklevel=2,
            )

        if _diagnostics is not None:
            _diagnostics.output_length = len(text)
            _diagnostics.output_empty = not text
            if _diagnostics.prompt_tokens == 0:
                _diagnostics.prompt_tokens = input_ids.shape[-1]
            if context is not None and _diagnostics.transfer_mode == "":
                _diagnostics.transfer_mode = "latent"

        return text

    # --- Cross-model rosetta projection ---

    def _apply_rosetta_projection(
        self,
        source: "HuggingFaceConnector",
        context: AVPContext,
        _diagnostics: Optional[Any] = None,
    ) -> AVPContext:
        """Project cross-model context to this model's space.

        Handles: AVPMap lookup/calibration, quality gate, projection, KV-cache priming.
        Returns a new AVPContext with model_hash == self._model_hash.
        """
        import torch

        if context.last_hidden_state is None:
            raise ValueError(
                "Cross-model generate() requires context with last_hidden_state. "
                "Use think() to produce the context (it captures the hidden state "
                "automatically)."
            )

        avp_map = self._get_or_calibrate_map(source)

        # Advisory quality gate — log warning but proceed
        prompt_tokens = max(0, context.seq_len - context.num_steps)
        from ..rosetta.quality import assess_transfer
        quality = assess_transfer(prompt_tokens=prompt_tokens)
        if not quality.recommend_latent:
            logger.warning(
                "Quality gate recommends JSON fallback for this transfer: %s",
                quality.reason,
            )

        # Determine projection method name
        method = getattr(avp_map, "method", None)
        method_name = method.name if hasattr(method, "name") else str(method or "")

        if _diagnostics is not None:
            _diagnostics.transfer_mode = "rosetta"
            _diagnostics.projection_method = method_name
            _diagnostics.quality_gate_passed = quality.recommend_latent
            _diagnostics.quality_gate_reason = quality.reason
            _diagnostics.prompt_tokens = prompt_tokens

        # Project source hidden state → target embedding space
        use_metrics = _diagnostics is not None
        result = source.project_hidden_for_cross_model(
            context.last_hidden_state, avp_map,
            return_metrics=use_metrics,
        )
        if use_metrics and isinstance(result, tuple):
            projected, proj_metrics = result
            if _diagnostics is not None and proj_metrics:
                _diagnostics.hidden_state_norm = float(proj_metrics["hidden_state_norm"].mean()) if "hidden_state_norm" in proj_metrics else None
                _diagnostics.nearest_cos_sim = float(proj_metrics["nearest_cos_sim"].mean()) if "nearest_cos_sim" in proj_metrics else None
        else:
            projected = result

        # Ensure [1, N, D_tgt] shape for KV-cache priming
        if projected.dim() == 1:
            projected = projected.unsqueeze(0)
        if projected.dim() == 2:
            projected = projected.unsqueeze(0)

        # Prime this model's KV-cache with projected embedding
        embed_input = projected.to(self.device).to(self.model.dtype)
        embed_mask = torch.ones(
            (1, embed_input.shape[1]), dtype=torch.long, device=self.device,
        )
        with torch.no_grad():
            prime_out = self.model(
                inputs_embeds=embed_input,
                attention_mask=embed_mask,
                use_cache=True,
                return_dict=True,
            )

        return AVPContext(
            past_key_values=prime_out.past_key_values,
            model_hash=self._model_hash,
            num_steps=context.num_steps,
            seq_len=embed_input.shape[1],
            model_family=self._identity.model_family,
            hidden_dim=self._identity.hidden_dim,
            num_layers=self._identity.num_layers,
        )

    def _get_or_calibrate_map(self, source: "HuggingFaceConnector") -> Any:
        """Get AVPMap for source→self, calibrating if needed.

        3-tier cache: memory → disk registry → calibrate (one-time).
        """
        # Tier 1: In-memory cache
        cache_key = source._model_hash
        if cache_key in self._avp_map_cache:
            return self._avp_map_cache[cache_key]

        # Tier 2: Disk registry (~/.avp/maps/)
        from ..rosetta.registry import load_map
        avp_map = load_map(
            source._model_hash, self._model_hash, device=self.device,
        )
        if avp_map is not None:
            self._avp_map_cache[cache_key] = avp_map
            return avp_map

        # Tier 3: Calibrate (auto-saves to disk)
        from ..rosetta.calibrate import calibrate
        avp_map = calibrate(
            source.model, self.model,
            source.tokenizer, self.tokenizer,
            device=self.device,
        )
        self._avp_map_cache[cache_key] = avp_map
        return avp_map

    def project_hidden_for_cross_model(
        self, hidden_state: Any, avp_map: Any, temperature: float = 1.0,
        return_metrics: bool = False,
    ) -> Any:
        """Project source hidden state to target embedding space via Rosetta Stone.

        Supports two method families:
        - vocab_mediated: Uses shared vocabulary as bridge (zero learned params).
          source lm_head -> softmax -> target input embeddings.
        - vocab_overlap: Cross-family variant — projects through overlapping
          tokens between different tokenizers.

        Args:
            hidden_state: Tensor of shape [..., D_src] from this (source) model.
            avp_map: AVPMap with w_map, target_norm, and optional bias.
            temperature: Softmax temperature for vocab-mediated/overlap projection.
                Lower = sharper (closer to argmax), higher = softer. Default 1.0.
                Tested 0.3–2.0 on GSM8K 2-agent rosetta (Qwen 7B → Llama 3B,
                n=200, A100): no significant difference between 0.5 and 1.0;
                temperatures above 1.0 degrade accuracy. Standard softmax (1.0)
                is the correct default.

        Returns:
            Projected tensor of shape [..., D_tgt] suitable for injection
            into the target model via inputs_embeds.
        """
        import torch as _torch

        def _np_to_torch(arr):
            """Convert numpy projection result back to torch tensor."""
            if isinstance(arr, _torch.Tensor):
                return arr
            return _torch.from_numpy(arr).to(
                device=hidden_state.device, dtype=hidden_state.dtype,
            )

        from ..types import ProjectionMethod
        if avp_map.method == ProjectionMethod.VOCAB_MEDIATED:
            from ..rosetta.project import vocabulary_mediated_projection
            source_lm_head = self.model.get_output_embeddings()
            if source_lm_head is None:
                source_lm_head = getattr(self.model, "lm_head", None)
            if source_lm_head is None or not hasattr(source_lm_head, "weight"):
                raise RealignmentError(
                    "Cannot get source output embeddings (lm_head) for "
                    "vocabulary-mediated projection."
                )
            result = vocabulary_mediated_projection(
                hidden_state,
                source_lm_head_weight=source_lm_head.weight,
                target_embed_weight=avp_map.w_map,
                temperature=temperature,
                target_norm=avp_map.target_norm,
                return_metrics=return_metrics,
            )
            if return_metrics:
                return _np_to_torch(result[0]), result[1]
            return _np_to_torch(result)
        elif avp_map.method == ProjectionMethod.VOCAB_OVERLAP:
            from ..rosetta.project import vocab_overlap_projection
            source_lm_head = self.model.get_output_embeddings()
            if source_lm_head is None:
                source_lm_head = getattr(self.model, "lm_head", None)
            if source_lm_head is None or not hasattr(source_lm_head, "weight"):
                raise RealignmentError(
                    "Cannot get source output embeddings (lm_head) for "
                    "vocab-overlap projection."
                )
            result = vocab_overlap_projection(
                hidden_state,
                source_lm_head_weight=source_lm_head.weight,
                shared_target_embed_weight=avp_map.w_map,
                src_indices=avp_map.src_indices,
                temperature=temperature,
                target_norm=avp_map.target_norm,
                return_metrics=return_metrics,
            )
            if return_metrics:
                return _np_to_torch(result[0]), result[1]
            return _np_to_torch(result)
        else:
            from ..rosetta.project import apply_cross_model_projection
            projected = apply_cross_model_projection(
                hidden_state, avp_map.w_map, avp_map.target_norm, avp_map.bias
            )
            if return_metrics:
                return _np_to_torch(projected), {}
            return _np_to_torch(projected)
