from __future__ import annotations

from typing import Optional

import torch
from torch import nn
from transformers.models.modernbert.modeling_modernbert import (
    ModernBertForMaskedLM,
    ModernBertModel,
    ModernBertPredictionHead,
    ModernBertPreTrainedModel,
)


class ModernGenaModel(ModernBertModel):
    """Keep unpadded Flash outputs rank-2 when called by the MLM wrapper."""

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        sliding_window_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        indices: Optional[torch.Tensor] = None,
        cu_seqlens: Optional[torch.Tensor] = None,
        max_seqlen: Optional[int] = None,
        batch_size: Optional[int] = None,
        seq_len: Optional[int] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        use_return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        output = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            sliding_window_mask=sliding_window_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            indices=indices,
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
            batch_size=batch_size,
            seq_len=seq_len,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=use_return_dict,
        )

        # ModernBertModel 4.57.6 adds a singleton batch dimension to an
        # already-unpadded Flash output whenever hidden states are requested.
        # ModernBertForMaskedLM expects (total_nonpad, hidden_size) here and
        # performs the final repadding itself.
        needs_squeeze = (
            self.config._attn_implementation == "flash_attention_2"
            and indices is not None
            and output[0].dim() == 3
            and output[0].shape[0] == 1
        )
        if needs_squeeze:
            if use_return_dict:
                output.last_hidden_state = output.last_hidden_state.squeeze(0)
            else:
                output = (output[0].squeeze(0),) + output[1:]
        return output


class ModernGenaForMaskedLM(ModernBertForMaskedLM):
    """ModernBERT MLM with correct Flash hidden-state/logit repadding."""

    def __init__(self, config):
        ModernBertPreTrainedModel.__init__(self, config)
        self.config = config
        self.model = ModernGenaModel(config)
        self.head = ModernBertPredictionHead(config)
        self.decoder = nn.Linear(
            config.hidden_size,
            config.vocab_size,
            bias=config.decoder_bias,
        )
        self.sparse_prediction = config.sparse_prediction
        self.sparse_pred_ignore_index = config.sparse_pred_ignore_index
        self.post_init()


__all__ = ["ModernGenaForMaskedLM", "ModernGenaModel"]