"""
Legacy model implementation of EAGLE-3. This is deprecated and will be removed in the
future.

Currently only used for converting workflow.
"""

import os
from typing import ClassVar, Literal

import torch
from torch import nn
from transformers import PretrainedConfig, PreTrainedModel
from transformers.models.llama.modeling_llama import (
    LlamaMLP,
    LlamaRMSNorm,
    apply_rotary_pos_emb,
    repeat_kv,
)

from speculators import SpeculatorModel
from speculators.models.eagle3 import Eagle3SpeculatorConfig

__all__ = ["Eagle3Speculator"]


class Eagle3Attention(nn.Module):
    """
    Eagle-3 attention module that processes concatenated embeddings and hidden states.

    Modified from standard Llama attention to accept 2x hidden_size input
    for Q/K/V projections while maintaining standard output size.
    """

    def __init__(self, config: PretrainedConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx

        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.hidden_size = config.hidden_size
        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads

        input_size = 2 * self.hidden_size
        self.q_proj = nn.Linear(
            input_size, self.num_heads * self.head_dim, bias=config.attention_bias
        )
        self.k_proj = nn.Linear(
            input_size,
            self.num_key_value_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.v_proj = nn.Linear(
            input_size,
            self.num_key_value_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.o_proj = nn.Linear(
            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs,  # noqa: ARG002
    ) -> tuple:
        """
        Forward pass for Eagle-3 attention.
        Taken from Llama Attention but modified to accept 2x hidden_size input.

        :param hidden_states: Input tensor of shape [batch, seq_len, 2*hidden_size]
        :param attention_mask: Optional attention mask
        :param position_ids: Optional position IDs for rotary embeddings
        :param past_key_value: Optional cached key-value pairs
        :param output_attentions: Whether to return attention weights
        :param use_cache: Whether to cache key-value pairs
        :param position_embeddings: Optional precomputed rotary embeddings
        :return: Tuple of (hidden_states, [attention_weights], [past_key_value])
        """
        bsz, q_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(
            bsz, q_len, self.num_heads, self.head_dim
        ).transpose(1, 2)
        key_states = key_states.view(
            bsz, q_len, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)
        value_states = value_states.view(
            bsz, q_len, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)

        if position_embeddings is not None:
            cos, sin = position_embeddings
            query_states, key_states = apply_rotary_pos_emb(
                query_states, key_states, cos, sin, position_ids
            )

        past_key_value_out = None
        if past_key_value is not None:
            past_key = past_key_value[0]
            past_value = past_key_value[1]
            key_states = torch.cat([past_key, key_states], dim=2)
            value_states = torch.cat([past_value, value_states], dim=2)

        if use_cache:
            past_key_value_out = (key_states, value_states)

        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / (
            self.head_dim**0.5
        )

        if attention_mask is not None:
            attn_weights = attn_weights + attention_mask

        attn_weights = nn.functional.softmax(
            attn_weights, dim=-1, dtype=torch.float32
        ).to(query_states.dtype)

        attn_output = torch.matmul(attn_weights, value_states)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(bsz, q_len, self.hidden_size)

        attn_output = self.o_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, past_key_value_out


class Eagle3DecoderLayer(nn.Module):
    """
    Eagle-3 decoder layer that processes concatenated embeddings and hidden states.

    Accepts 2x hidden_size input from concatenated embeddings and fused hidden states.
    Uses Eagle3Attention for the self-attention computation.
    """

    def __init__(
        self,
        config: PretrainedConfig,
        layer_idx: int,
        norm_before_residual: bool = False,
    ):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.norm_before_residual = norm_before_residual

        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.hidden_norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = LlamaRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

        self.self_attn = Eagle3Attention(config, layer_idx)

        self.mlp = LlamaMLP(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
        output_attentions: bool | None = False,
        use_cache: bool | None = False,
        cache_position: torch.LongTensor | None = None,  # noqa: ARG002
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs,  # noqa: ARG002
    ) -> tuple:
        """
        Process concatenated embeddings and hidden states through modified decoder
        layer.

        :param hidden_states: Input tensor of shape [batch, seq_len, 2*hidden_size]
        :return: Tuple of layer outputs
        """
        embeds = hidden_states[:, :, : self.hidden_size]
        hidden = hidden_states[:, :, self.hidden_size : 2 * self.hidden_size]

        if self.norm_before_residual:
            hidden = self.hidden_norm(hidden)
            residual = hidden
        else:
            residual = hidden
            hidden = self.hidden_norm(hidden)

        embeds = self.input_layernorm(embeds)

        attn_input = torch.cat([embeds, hidden], dim=-1)

        attn_output, attn_weights, past_key_value_out = self.self_attn(
            hidden_states=attn_input,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            position_embeddings=position_embeddings,
        )

        hidden_states = residual + attn_output

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (attn_weights,)  # type: ignore[assignment]

        if use_cache:
            outputs += (past_key_value_out,)  # type: ignore[assignment]

        return outputs


class Eagle3Speculator(SpeculatorModel):
    """
    EAGLE-3 speculator with vocabulary mapping and multi-layer fusion.

    EAGLE-3 processes concatenated hidden states from multiple verifier layers
    through a fusion layer, then combines with embeddings for a custom decoder
    layer that accepts 2x hidden_size input.
    """

    config_class: ClassVar[type[Eagle3SpeculatorConfig]] = Eagle3SpeculatorConfig  # type: ignore[misc]
    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [  # type: ignore[misc]
        "verifier*",
    ]
    _keys_to_ignore_on_save: ClassVar[list[str]] = []  # type: ignore[misc,assignment]

    def __init__(
        self,
        config: Eagle3SpeculatorConfig,
        verifier: str | os.PathLike | PreTrainedModel | None = None,
        verifier_attachment_mode: Literal["detached", "full", "train_only"]
        | None = None,
        reduce_vocab_size: bool = True,
        has_drafter_embedding: bool = True,
    ):
        """
        Initialize Eagle3 speculator.

        :param config: Eagle3SpeculatorConfig instance
        :param verifier: Optional verifier model
        :param verifier_attachment_mode: How to attach the verifier
        :param reduce_vocab_size: Whether to reduce vocabulary size with mapping
        :param has_drafter_embedding: Whether drafter embedding weights are provided
        """
        if not isinstance(config, Eagle3SpeculatorConfig):
            raise ValueError(
                f"config must be Eagle3SpeculatorConfig, got {type(config)}"
            )

        self.config: Eagle3SpeculatorConfig = config

        self.hidden_size = config.transformer_layer_config.hidden_size
        self.draft_vocab_size = config.draft_vocab_size
        self.target_vocab_size = config.target_vocab_size

        # Use target_hidden_size if specified, otherwise use draft model's hidden_size
        self.target_hidden_size = (
            config.target_hidden_size
            if config.target_hidden_size is not None
            else self.hidden_size
        )

        super().__init__(
            config=config,
            verifier=verifier,
            verifier_attachment_mode=verifier_attachment_mode,
        )

        if has_drafter_embedding:
            self.embed_tokens = nn.Embedding(
                self.target_vocab_size,
                self.hidden_size,
                padding_idx=config.transformer_layer_config.pad_token_id
                if hasattr(config.transformer_layer_config, "pad_token_id")
                else None,
            )

        self.fc = nn.Linear(
            3 * self.target_hidden_size,  # Use target model's hidden size
            self.hidden_size,
            bias=False,
        )

        self.layers = nn.ModuleList(
            [
                Eagle3DecoderLayer(
                    config.transformer_layer_config,
                    layer_idx=0,
                    norm_before_residual=config.norm_before_residual,
                )
            ]
        )

        self.norm = LlamaRMSNorm(
            self.hidden_size,
            eps=config.transformer_layer_config.rms_norm_eps,
        )

        self.lm_head = nn.Linear(
            self.hidden_size,
            self.draft_vocab_size,
            bias=False,
        )
        if reduce_vocab_size:
            self.register_buffer(  # type: ignore[attr-defined]
                "d2t",
                torch.zeros(self.draft_vocab_size, dtype=torch.long),
            )
            self.register_buffer(  # type: ignore[attr-defined]
                "t2d",
                torch.zeros(self.target_vocab_size, dtype=torch.bool),
            )

            # Type hints for buffers
            self.d2t: torch.Tensor
            self.t2d: torch.Tensor
        self.post_init()  # type: ignore[attr-defined]

    def tie_weights(self):
        """
        Override tie_weights to prevent vocabulary corruption in transformers 4.54.1+

        Eagle3 intentionally uses different vocabulary sizes:
        - Input embeddings (embed_tokens): 128256 (full vocabulary)
        - Output embeddings (lm_head): 32000 (draft vocabulary)

        The default tie_weights() tries to make them identical, breaking Eagle3.
        This override preserves the intentional vocabulary size difference.
        """
        # Don't call super().tie_weights() - this prevents vocabulary corruption
        # that occurs when _tie_or_clone_weights replaces lm_head.weight with
        # embed_tokens.weight

    def forward(
        self,
        input_ids: torch.LongTensor,
        hidden_states: torch.FloatTensor,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: tuple[tuple[torch.FloatTensor]] | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,  # noqa: ARG002
        return_dict: bool | None = None,
    ) -> torch.FloatTensor:
        """
        Forward pass for EAGLE-3 speculation.

        :param input_ids: Input token IDs from draft vocabulary
        :param hidden_states: Concatenated hidden states from 3 verifier layers
            [B, L, 3*target_H] where target_H is the target model's hidden size
        :param attention_mask: Optional attention mask
        :param position_ids: Optional position IDs
        :param past_key_values: Optional cached key-values
        :param use_cache: Whether to cache key-values
        :param output_attentions: Return attention weights
        :param output_hidden_states: Return hidden states
        :param return_dict: Return dict output
        :return: Model outputs with draft vocabulary logits
        """
        raise NotImplementedError("Eagle3Speculator.forward is not implemented yet.")
