"""Gemma3 decoder with SLIP cross-attention blocks at layers >= split_layer.

Refactored from HealthSensorSLM-Bench/model_factory/multimodal_gemma.py. The
key change is how sensor context reaches the cross-attention blocks: instead
of caching ``self.vis_x`` as module state (which is unsafe when the same
decoder is called multiple times per step with different micro-batch shapes,
as CODI does), each call threads ``vis_x`` through a ``contextvars.ContextVar``
scoped by a context manager. Wrapped decoder layers read the ContextVar on
every forward — so the context is always fresh and batch-size-correct.

The ``.to(torch.bfloat16)`` hard-cast from the original is also removed; dtype
is taken from the surrounding model so autocast / fp32 runs work unchanged.
"""

import os
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Optional

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModelForCausalLM

from .ts_transformer import CrossAttention


_SLIP_VIS_X: "ContextVar[Optional[torch.Tensor]]" = ContextVar("slip_vis_x", default=None)


@contextmanager
def slip_vis_x_ctx(vis_x: Optional[torch.Tensor]):
    """Bind ``vis_x`` as the cross-attention context for the enclosed forward.

    Safe to nest; the outer binding is restored on exit. Reset on exceptions.
    """
    token = _SLIP_VIS_X.set(vis_x)
    try:
        yield
    finally:
        _SLIP_VIS_X.reset(token)


class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x, *args, **kwargs):
        return self.fn(x, *args, **kwargs) + x


class Gemma3MultimodalLayer(nn.Module):
    """Wraps a Gemma3 decoder layer with a post-hoc cross-attention block.

    Reads the current sensor context from the module-level ``_SLIP_VIS_X``
    ContextVar rather than from ``self``, so the context stays accurate across
    consecutive forwards with changing batch shapes.
    """

    _xattn_debug_calls: int = 0
    _XATTN_DEBUG_MAX_CALLS: int = 60

    def __init__(self, original_layer, cross_attn_block):
        super().__init__()
        self.original_layer = original_layer
        self.cross_attn_block = cross_attn_block

    def __getattr__(self, name):
        try:
            return super().__getattr__(name)
        except AttributeError:
            try:
                original_layer = super().__getattr__("original_layer")
            except AttributeError:
                raise AttributeError(name)
            return getattr(original_layer, name)

    def forward(self, hidden_states, **kwargs):
        vis_x = _SLIP_VIS_X.get()
        assert vis_x is not None, (
            "SLIP cross-attention invoked without a sensor context. "
            "Wrap the model forward in `slip_vis_x_ctx(sensor_hidden)` first."
        )

        outputs = self.original_layer(hidden_states, **kwargs)
        if isinstance(outputs, tuple):
            hidden_states = outputs[0]
        else:
            hidden_states = outputs

        vis_x_cast = vis_x.to(dtype=hidden_states.dtype, device=hidden_states.device)

        cls = Gemma3MultimodalLayer
        _debug = (
            os.environ.get("TLR_TS_DEBUG") == "1"
            and cls._xattn_debug_calls < cls._XATTN_DEBUG_MAX_CALLS
        )
        if _debug:
            pre = hidden_states
            post = self.cross_attn_block(pre, context=vis_x_cast)
            pre_f = pre.detach().float()
            post_f = post.detach().float()
            base = pre_f.norm().clamp_min(1e-6)
            delta = (post_f - pre_f).norm()
            print(
                f"[XATTN] call={cls._xattn_debug_calls} "
                f"vis_x_norm={vis_x_cast.detach().float().norm().item():.3f} "
                f"pre_norm={pre_f.norm().item():.3f} "
                f"post_norm={post_f.norm().item():.3f} "
                f"delta_over_pre={(delta / base).item():.4e}",
                flush=True,
            )
            cls._xattn_debug_calls += 1
            hidden_states = post
        else:
            hidden_states = self.cross_attn_block(hidden_states, context=vis_x_cast)

        if isinstance(outputs, tuple):
            return (hidden_states,) + outputs[1:]
        return hidden_states


class Gemma3MultimodalModel(nn.Module):
    """Standard HF Gemma3 with SLIP cross-attention inserted at/after ``split_layer``."""

    def __init__(
        self,
        model_id: str = "google/gemma-3-270m",
        post_train: bool = False,
        split_layer: int = 12,
        num_heads: Optional[int] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        super().__init__()
        load_kwargs = dict(attn_implementation="sdpa", trust_remote_code=True)
        if dtype is not None:
            load_kwargs["dtype"] = dtype

        # post_train=False is the right default when loading SLIP via
        # SlipForCausalLM.from_pretrained(...): the SLIP safetensors contains the
        # full post-trained Gemma weights already, so downloading the vanilla
        # google/gemma-3-270m checkpoint just to overwrite it is wasted work (and
        # produces a misleading "lm_head.weight newly initialized" warning —
        # Gemma-3 ties lm_head to embed_tokens, so it's never in the checkpoint).
        # Flip to True only when you want a randomly-initialised cross-attention
        # head on top of actually-pretrained Gemma weights, not overwritten by an
        # outer load.
        if post_train:
            self.model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs)
        else:
            config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
            self.model = AutoModelForCausalLM.from_config(config, **load_kwargs)

        self.split_layer = split_layer
        hidden_size = self.model.config.hidden_size
        self.hidden_size = hidden_size
        if num_heads is None:
            num_heads = self.model.config.num_attention_heads

        for i in range(split_layer, len(self.model.model.layers)):
            cross_attn = CrossAttention(
                dim=hidden_size, context_dim=hidden_size, num_heads=num_heads, dropout_rate=0.1
            )
            original_layer = self.model.model.layers[i]
            self.model.model.layers[i] = Gemma3MultimodalLayer(original_layer, Residual(cross_attn))