import math
from dataclasses import dataclass
from typing import Optional

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

from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

from .configuration_binaryllm import BinaryLLMConfig

try:
    import flash_attn_v100_cuda
    _FLASH_V100_AVAILABLE = True
except Exception:
    flash_attn_v100_cuda = None
    _FLASH_V100_AVAILABLE = False


class PositionalEncoding(nn.Module):
    """
    Sinusoidal positional encoding, stocké en fp32,
    puis casté au dtype de x à chaque forward.
    """

    def __init__(self, d_model: int, max_len: int) -> None:
        super().__init__()
        pe = torch.zeros(max_len, d_model, dtype=torch.float32)
        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float32)
            * (-torch.log(torch.tensor(10000.0)) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        t = x.size(1)
        pe = self.pe[:, :t, :].to(device=x.device, dtype=x.dtype)
        return x + pe


@dataclass
class _InnerCfg:
    block_size: int
    embed_dim: int
    vocab_size: int
    num_heads: int
    num_layers: int
    ff_hidden_dim: int
    dropout: float
    layernorm_dim: Optional[int] = None
    head_dim: Optional[int] = None
    attn_backend: str = "auto"


class FlashSelfAttentionPortable(nn.Module):
    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        causal: bool = True,
        backend: str = "auto",
    ) -> None:
        super().__init__()

        if embed_dim % num_heads != 0:
            raise ValueError(
                f"embed_dim ({embed_dim}) doit être divisible par num_heads ({num_heads})"
            )

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.dropout = float(dropout)
        self.causal = bool(causal)
        self.backend = str(backend)
        self.softmax_scale = self.head_dim ** -0.5

        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=True)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)

    def _shape_qkv(
        self,
        x: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.dtype]:
        bsz, seqlen, _ = x.shape
        residual_dtype = x.dtype

        proj_dtype = self.q_proj.weight.dtype
        if x.dtype != proj_dtype:
            x = x.to(proj_dtype)

        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        q = q.view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k = k.view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

        return q, k, v, residual_dtype

    def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
        bsz, nheads, seqlen, head_dim = x.shape
        return x.transpose(1, 2).contiguous().view(bsz, seqlen, nheads * head_dim)

    def _can_use_v100_kernel(self, q: torch.Tensor, padding_mask: Optional[torch.Tensor]) -> bool:
        if not _FLASH_V100_AVAILABLE:
            return False

        if not q.is_cuda:
            return False

        if padding_mask is not None and bool(padding_mask.any().item()):
            return False

        cc = torch.cuda.get_device_capability(q.device)
        if cc != (7, 0):
            return False

        hd = q.size(-1)
        if hd % 2 != 0:
            return False
        if hd % 8 != 0:
            return False
        if hd > 256:
            return False

        return True

    def _flash_attn_v100(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
    ) -> torch.Tensor:
        if q.dtype != torch.float16:
            q = q.to(torch.float16)
        if k.dtype != torch.float16:
            k = k.to(torch.float16)
        if v.dtype != torch.float16:
            v = v.to(torch.float16)

        result = flash_attn_v100_cuda.fwd(
            q,
            k,
            v,
            None,
            None,
            0.0,
            self.softmax_scale,
            self.causal,
            -1,
            -1,
            0.0,
            False,
            None,
        )

        out = result[0]
        return out

    def _sdpa_attn(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        bsz, nheads, tq, _ = q.shape
        tk = k.size(-2)

        attn_mask = None
        if padding_mask is not None:
            key_mask = padding_mask[:, None, None, :].to(device=q.device, dtype=torch.bool)
            key_mask = key_mask.expand(bsz, nheads, tq, tk)
            attn_mask = ~key_mask

        dropout_p = self.dropout if self.training else 0.0

        with torch.backends.cuda.sdp_kernel(
            enable_flash=True,
            enable_mem_efficient=True,
            enable_math=True,
        ):
            out = F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=attn_mask,
                dropout_p=dropout_p,
                is_causal=self.causal if attn_mask is None else False,
                scale=self.softmax_scale,
            )

        return out

    def _eager_attn(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * self.softmax_scale

        if self.causal:
            tq = q.size(-2)
            tk = k.size(-2)
            causal_mask = torch.triu(
                torch.ones(tq, tk, device=scores.device, dtype=torch.bool),
                diagonal=1,
            )
            scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))

        if padding_mask is not None:
            key_mask = padding_mask[:, None, None, :].to(device=scores.device, dtype=torch.bool)
            scores = scores.masked_fill(key_mask, float("-inf"))

        probs = torch.softmax(scores, dim=-1)

        if self.training and self.dropout > 0.0:
            probs = F.dropout(probs, p=self.dropout)

        out = torch.matmul(probs, v.float())
        return out.to(q.dtype)

    def forward(
        self,
        x: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        q, k, v, residual_dtype = self._shape_qkv(x)

        if padding_mask is not None:
            padding_mask = padding_mask.to(device=x.device, dtype=torch.bool)

        backend = self.backend

        if backend == "v100":
            if not self._can_use_v100_kernel(q, padding_mask):
                raise RuntimeError(
                    "backend='v100' demandé mais indisponible "
                    "(flash_attn_v100_cuda absent, GPU non sm70/V100, padding présent, "
                    "ou head_dim incompatible)."
                )
            out = self._flash_attn_v100(q, k, v)

        elif backend == "sdpa":
            out = self._sdpa_attn(q, k, v, padding_mask=padding_mask)

        elif backend == "eager":
            out = self._eager_attn(q, k, v, padding_mask=padding_mask)

        elif backend == "auto":
            if self._can_use_v100_kernel(q, padding_mask):
                out = self._flash_attn_v100(q, k, v)
            else:
                out = self._sdpa_attn(q, k, v, padding_mask=padding_mask)

        else:
            raise ValueError(f"backend d'attention non supporté: {backend}")

        out = self._merge_heads(out)

        out_proj_dtype = self.out_proj.weight.dtype
        if out.dtype != out_proj_dtype:
            out = out.to(out_proj_dtype)

        out = self.out_proj(out)

        if out.dtype != residual_dtype:
            out = out.to(residual_dtype)

        return out


class FlashTransformerEncoderLayerPortable(nn.Module):
    def __init__(
        self,
        d_model: int,
        nhead: int,
        dim_feedforward: int,
        dropout: float = 0.1,
        activation: str = "gelu",
        batch_first: bool = True,
        attn_backend: str = "auto",
    ) -> None:
        super().__init__()

        if not batch_first:
            raise ValueError("Cette implémentation supporte batch_first=True uniquement.")

        self.self_attn = FlashSelfAttentionPortable(
            embed_dim=d_model,
            num_heads=nhead,
            dropout=dropout,
            causal=True,
            backend=attn_backend,
        )

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        self.dropout = nn.Dropout(dropout)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        if activation == "gelu":
            self.activation = F.gelu
        elif activation == "relu":
            self.activation = F.relu
        else:
            raise ValueError(f"activation non supportée: {activation}")

    def _sa_block(
        self,
        x: torch.Tensor,
        src_key_padding_mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        x = self.self_attn(x, padding_mask=src_key_padding_mask)
        x = self.dropout1(x)
        return x

    def _ff_block(self, x: torch.Tensor) -> torch.Tensor:
        ff_dtype = self.linear1.weight.dtype
        x_ff = x if x.dtype == ff_dtype else x.to(ff_dtype)

        x_ff = self.linear1(x_ff)
        x_ff = self.activation(x_ff)
        x_ff = self.dropout(x_ff)
        x_ff = self.linear2(x_ff)
        x_ff = self.dropout2(x_ff)

        if x_ff.dtype != x.dtype:
            x_ff = x_ff.to(x.dtype)

        return x_ff

    def forward(
        self,
        src: torch.Tensor,
        src_mask: Optional[torch.Tensor] = None,
        src_key_padding_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        x = src
        x = self.norm1(x + self._sa_block(x, src_key_padding_mask))
        x = self.norm2(x + self._ff_block(x))
        return x


class FlashTransformerEncoderPortable(nn.Module):
    def __init__(
        self,
        encoder_layer: FlashTransformerEncoderLayerPortable,
        num_layers: int,
        attn_backend: str = "auto",
    ) -> None:
        super().__init__()

        d_model = encoder_layer.norm1.normalized_shape[0]
        nhead = encoder_layer.self_attn.num_heads
        dim_feedforward = encoder_layer.linear1.out_features
        dropout = encoder_layer.dropout.p

        self.layers = nn.ModuleList(
            [
                FlashTransformerEncoderLayerPortable(
                    d_model=d_model,
                    nhead=nhead,
                    dim_feedforward=dim_feedforward,
                    dropout=dropout,
                    activation="gelu",
                    batch_first=True,
                    attn_backend=attn_backend,
                )
                for _ in range(num_layers)
            ]
        )

    def forward(
        self,
        src: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        src_key_padding_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        x = src
        for layer in self.layers:
            x = layer(x, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
        return x


class TinyTransformerLM(nn.Module):
    def __init__(self, cfg: _InnerCfg) -> None:
        super().__init__()
        self.cfg = cfg

        vocab_size = cfg.vocab_size
        self.tok_embed = nn.Embedding(vocab_size, cfg.embed_dim)
        self.pos_encoding = PositionalEncoding(cfg.embed_dim, cfg.block_size)

        encoder_layer = FlashTransformerEncoderLayerPortable(
            d_model=cfg.embed_dim,
            nhead=cfg.num_heads,
            dim_feedforward=cfg.ff_hidden_dim,
            dropout=cfg.dropout,
            activation="gelu",
            batch_first=True,
            attn_backend=cfg.attn_backend,
        )
        self.encoder = FlashTransformerEncoderPortable(
            encoder_layer,
            num_layers=cfg.num_layers,
            attn_backend=cfg.attn_backend,
        )

        ln_dim = cfg.layernorm_dim or cfg.embed_dim
        head_dim = cfg.head_dim or ln_dim

        self.pre_ln_proj: Optional[nn.Linear] = None
        if ln_dim != cfg.embed_dim:
            self.pre_ln_proj = nn.Linear(cfg.embed_dim, ln_dim)

        self.ln = nn.LayerNorm(ln_dim)

        self.head_pre: Optional[nn.Linear] = None
        if head_dim != ln_dim:
            self.head_pre = nn.Linear(ln_dim, head_dim)

        self.head = nn.Linear(head_dim, vocab_size, bias=False)

        if self.pre_ln_proj is None and self.head_pre is None and head_dim == cfg.embed_dim:
            self.head.weight = self.tok_embed.weight

        causal = torch.triu(torch.ones(cfg.block_size, cfg.block_size, dtype=torch.bool), diagonal=1)
        self.register_buffer("causal_mask", causal, persistent=False)

    def forward(self, tokens: torch.Tensor, padding_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        x = self.tok_embed(tokens)
        x = self.pos_encoding(x)

        seq_len = tokens.size(1)
        attn_mask = self.causal_mask[:seq_len, :seq_len].to(device=tokens.device)

        if padding_mask is not None:
            padding_mask = padding_mask[:, :seq_len].to(device=tokens.device, dtype=torch.bool)

        x = self.encoder(x, mask=attn_mask, src_key_padding_mask=padding_mask)

        if self.pre_ln_proj is not None:
            proj_dtype = self.pre_ln_proj.weight.dtype
            if x.dtype != proj_dtype:
                x = x.to(proj_dtype)
            x = self.pre_ln_proj(x)

        ln_dtype = self.ln.weight.dtype
        if x.dtype != ln_dtype:
            x = x.to(ln_dtype)
        x = self.ln(x)

        if self.head_pre is not None:
            head_pre_dtype = self.head_pre.weight.dtype
            if x.dtype != head_pre_dtype:
                x = x.to(head_pre_dtype)
            x = self.head_pre(x)

        head_dtype = self.head.weight.dtype
        if x.dtype != head_dtype:
            x = x.to(head_dtype)

        return self.head(x)


class BinaryLLMForCausalLM(PreTrainedModel):
    config_class = BinaryLLMConfig
    main_input_name = "input_ids"

    def __init__(self, config: BinaryLLMConfig):
        super().__init__(config)

        attn_backend = getattr(config, "attn_backend", "auto")

        inner = _InnerCfg(
            block_size=int(config.max_position_embeddings),
            embed_dim=int(config.hidden_size),
            vocab_size=int(config.vocab_size),
            num_heads=int(config.num_attention_heads),
            num_layers=int(config.num_hidden_layers),
            ff_hidden_dim=int(config.intermediate_size),
            dropout=float(getattr(config, "dropout", 0.0)),
            layernorm_dim=None,
            head_dim=None,
            attn_backend=str(attn_backend),
        )
        self.model = TinyTransformerLM(inner)

        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.model.tok_embed

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.model.tok_embed = value

    def get_output_embeddings(self) -> nn.Module:
        return self.model.head

    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
        self.model.head = new_embeddings

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs,
    ):
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
        }

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> CausalLMOutput:
        padding_mask = None
        if attention_mask is not None:
            padding_mask = ~attention_mask.to(torch.bool)

        logits = self.model(input_ids, padding_mask=padding_mask)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        return CausalLMOutput(loss=loss, logits=logits)