import math
import sys
from typing import Optional, Tuple, Union

# Monkeypatch safetensors to handle None/missing metadata which crashes transformers
try:
    import safetensors
    original_safe_open = safetensors.safe_open

    class SafeOpenWrapper:
        def __init__(self, original_obj):
            self.original_obj = original_obj

        def __enter__(self):
            self.original_obj.__enter__()
            return self

        def __exit__(self, exc_type, exc_val, exc_tb):
            return self.original_obj.__exit__(exc_type, exc_val, exc_tb)

        def metadata(self):
            meta = self.original_obj.metadata()
            if meta is None:
                return {"format": "pt"}
            return meta

        def __getattr__(self, name):
            return getattr(self.original_obj, name)

    def patched_safe_open(*args, **kwargs):
        f = original_safe_open(*args, **kwargs)
        return SafeOpenWrapper(f)

    safetensors.safe_open = patched_safe_open

    if "transformers.modeling_utils" in sys.modules:
        import transformers.modeling_utils
        transformers.modeling_utils.safe_open = patched_safe_open
except Exception:
    pass


import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import PreTrainedModel
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

from .configuration_tiny import TinyConfig


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., :x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)


def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    assert dim % 2 == 0
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(end)
    freqs = torch.outer(t, freqs).float()
    cos = torch.cos(freqs)
    sin = torch.sin(freqs)
    cos = torch.cat([cos, cos], dim=-1)
    sin = torch.cat([sin, sin], dim=-1)
    return cos, sin


def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    T = x.shape[2]
    cos_t = cos[:T, :].unsqueeze(0).unsqueeze(1)
    sin_t = sin[:T, :].unsqueeze(0).unsqueeze(1)
    return (x * cos_t) + (rotate_half(x) * sin_t)


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.pow(2).mean(-1, keepdim=True)
        return x * torch.rsqrt(variance + self.eps) * self.weight


class FeedForward(nn.Module):
    def __init__(self, config: TinyConfig):
        super().__init__()
        self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.w2 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.w3 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.dropout = nn.Dropout(config.hidden_dropout) if config.hidden_dropout > 0.0 else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.silu(self.w1(x)) * self.w2(x)
        out = self.w3(out)
        if self.dropout is not None:
            out = self.dropout(out)
        return out


class Attention(nn.Module):
    def __init__(self, config: TinyConfig):
        super().__init__()
        self.n_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.head_dim = config.hidden_size // config.num_attention_heads
        
        assert self.n_heads * self.head_dim == self.hidden_size
        
        self.wq = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.wk = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.wv = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.wo = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.dropout_p = config.attention_dropout

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        B, T, C = x.shape
        q = self.wq(x)
        k = self.wk(x)
        v = self.wv(x)
        
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        
        q = apply_rotary_emb(q, cos, sin)
        k = apply_rotary_emb(k, cos, sin)
        
        dropout_p = self.dropout_p if self.training else 0.0
        
        if attention_mask is not None:
            if torch.all(attention_mask == 1):
                attn_mask = None
                is_causal = True
            else:
                causal_mask = torch.tril(torch.ones((T, T), dtype=torch.bool, device=x.device))
                padding_mask = attention_mask.to(torch.bool).unsqueeze(1).unsqueeze(2) # shape: (B, 1, 1, T)
                attn_mask = causal_mask.unsqueeze(0).unsqueeze(1) & padding_mask # shape: (B, 1, T, T)
                is_causal = False
        else:
            attn_mask = None
            is_causal = True
            
        out = F.scaled_dot_product_attention(
            q, k, v, 
            attn_mask=attn_mask, 
            dropout_p=dropout_p, 
            is_causal=is_causal
        )
        
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.wo(out)


class TransformerBlock(nn.Module):
    def __init__(self, config: TinyConfig):
        super().__init__()
        self.attention = Attention(config)
        self.feed_forward = FeedForward(config)
        self.attention_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        x = x + self.attention(self.attention_norm(x), cos, sin, attention_mask)
        x = x + self.feed_forward(self.ffn_norm(x))
        return x


class TinyPreTrainedModel(PreTrainedModel):
    config_class = TinyConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["TransformerBlock"]

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, (TinyModel, TinyForCausalLM)):
            module.gradient_checkpointing = value


class TinyModel(TinyPreTrainedModel):
    def __init__(self, config: TinyConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.tok_embeddings = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [TransformerBlock(config) for _ in range(config.num_hidden_layers)]
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        cos, sin = precompute_freqs_cis(
            dim=config.hidden_size // config.num_attention_heads,
            end=config.max_position_embeddings * 2,
            theta=config.rope_theta,
        )
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.tok_embeddings

    def set_input_embeddings(self, value):
        self.tok_embeddings = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            hidden_states = self.tok_embeddings(input_ids)
        elif inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            raise ValueError("You must specify either input_ids or inputs_embeds")

        T = hidden_states.shape[1]
        cos = self.cos[:T]
        sin = self.sin[:T]

        all_hidden_states = () if output_hidden_states else None
        
        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:
                hidden_states = self._gradient_checkpointing_func(
                    layer.__call__,
                    hidden_states,
                    cos,
                    sin,
                    attention_mask,
                )
            else:
                hidden_states = layer(
                    hidden_states,
                    cos,
                    sin,
                    attention_mask,
                )

        hidden_states = self.norm(hidden_states)

        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        if not return_dict:
            return (hidden_states,)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            past_key_values=None,
        )


class TinyForCausalLM(TinyPreTrainedModel):
    _tied_weights_keys = ["tok_embeddings.weight"]

    def __init__(self, config: TinyConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.tok_embeddings = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [TransformerBlock(config) for _ in range(config.num_hidden_layers)]
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        
        self.output = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        cos, sin = precompute_freqs_cis(
            dim=config.hidden_size // config.num_attention_heads,
            end=config.max_position_embeddings * 2,
            theta=config.rope_theta,
        )
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

        if config.tie_word_embeddings:
            self.output.weight = self.tok_embeddings.weight

        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.tok_embeddings

    def set_input_embeddings(self, value):
        self.tok_embeddings = value
        if self.config.tie_word_embeddings:
            self.output.weight = value.weight

    def get_output_embeddings(self):
        return self.output

    def set_output_embeddings(self, new_embeddings):
        self.output = new_embeddings

    def tie_weights(self):
        if self.config.tie_word_embeddings:
            self.tok_embeddings.weight = self.output.weight

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            hidden_states = self.tok_embeddings(input_ids)
        elif inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            raise ValueError("You must specify either input_ids or inputs_embeds")

        T = hidden_states.shape[1]
        cos = self.cos[:T]
        sin = self.sin[:T]

        all_hidden_states = () if output_hidden_states else None
        
        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:
                hidden_states = self._gradient_checkpointing_func(
                    layer.__call__,
                    hidden_states,
                    cos,
                    sin,
                    attention_mask,
                )
            else:
                hidden_states = layer(
                    hidden_states,
                    cos,
                    sin,
                    attention_mask,
                )

        hidden_states = self.norm(hidden_states)
        logits = self.output(hidden_states)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))

        if not return_dict:
            output = (logits,)
            if output_hidden_states:
                output = output + (all_hidden_states,)
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=all_hidden_states,
            attentions=None,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        **kwargs,
    ) -> dict:
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs["attention_mask"] = attention_mask
        model_inputs["past_key_values"] = past_key_values
        return model_inputs