from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import mlx.core as mx
import mlx.nn as nn


WEIGHT_FILES = {
    "bf16": "model-bf16.safetensors",
    "16bit": "model-bf16.safetensors",
    "8bit": "model-8bit.safetensors",
    "4bit": "model-4bit.safetensors",
}


class EuroBertConfig:
    def __init__(self, **kwargs: Any):
        self.vocab_size = kwargs.get("vocab_size", 128257)
        self.hidden_size = kwargs.get("hidden_size", 768)
        self.intermediate_size = kwargs.get("intermediate_size", 3072)
        self.num_hidden_layers = kwargs.get("num_hidden_layers", 12)
        self.num_attention_heads = kwargs.get("num_attention_heads", 12)
        self.num_key_value_heads = kwargs.get(
            "num_key_value_heads", self.num_attention_heads
        )
        self.head_dim = kwargs.get(
            "head_dim", self.hidden_size // self.num_attention_heads
        )
        self.hidden_act = kwargs.get("hidden_act", "silu")
        self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-5)
        self.rope_theta = kwargs.get("rope_theta", 250000.0)
        self.attention_bias = kwargs.get("attention_bias", False)
        self.mlp_bias = kwargs.get("mlp_bias", False)
        self.pad_token_id = kwargs.get("pad_token_id", 128001)
        self.max_position_embeddings = kwargs.get("max_position_embeddings", 8192)
        self.num_labels = kwargs.get("num_labels", len(kwargs.get("id2label", {})) or 2)
        self.raw = kwargs

    @classmethod
    def from_json(cls, path: str | Path) -> "EuroBertConfig":
        with open(path, "r", encoding="utf-8") as f:
            return cls(**json.load(f))


def _silu(x: mx.array) -> mx.array:
    return x * mx.sigmoid(x)


def _rotate_half(x: mx.array) -> mx.array:
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return mx.concatenate([-x2, x1], axis=-1)


def _apply_rotary_pos_emb(
    q: mx.array, k: mx.array, cos: mx.array, sin: mx.array
) -> Tuple[mx.array, mx.array]:
    cos = cos[:, None, :, :]
    sin = sin[:, None, :, :]
    return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin)


class EuroBertRotaryEmbedding(nn.Module):
    def __init__(self, config: EuroBertConfig):
        super().__init__()
        self.head_dim = config.head_dim
        self.rope_theta = config.rope_theta

    def __call__(self, position_ids: mx.array, dtype: mx.Dtype) -> Tuple[mx.array, mx.array]:
        steps = mx.arange(0, self.head_dim, 2).astype(mx.float32)
        inv_freq = 1.0 / (self.rope_theta ** (steps / self.head_dim))
        pos = position_ids.astype(mx.float32)
        freqs = pos[..., None] * inv_freq[None, None, :]
        emb = mx.concatenate([freqs, freqs], axis=-1)
        return mx.cos(emb).astype(dtype), mx.sin(emb).astype(dtype)


class EuroBertAttention(nn.Module):
    def __init__(self, config: EuroBertConfig):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.head_dim = config.head_dim
        self.scaling = self.head_dim**-0.5

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

    def _shape(self, x: mx.array, heads: int) -> mx.array:
        batch, seq_len, _ = x.shape
        x = x.reshape(batch, seq_len, heads, self.head_dim)
        return mx.transpose(x, (0, 2, 1, 3))

    def __call__(
        self,
        hidden_states: mx.array,
        position_embeddings: Tuple[mx.array, mx.array],
        attention_mask: Optional[mx.array],
    ) -> mx.array:
        batch, seq_len, _ = hidden_states.shape
        q = self._shape(self.q_proj(hidden_states), self.num_heads)
        k = self._shape(self.k_proj(hidden_states), self.num_key_value_heads)
        v = self._shape(self.v_proj(hidden_states), self.num_key_value_heads)

        cos, sin = position_embeddings
        q, k = _apply_rotary_pos_emb(q, k, cos, sin)

        if self.num_key_value_groups != 1:
            k = mx.repeat(k, self.num_key_value_groups, axis=1)
            v = mx.repeat(v, self.num_key_value_groups, axis=1)

        scores = (q @ mx.transpose(k, (0, 1, 3, 2))).astype(mx.float32)
        scores = scores * self.scaling
        if attention_mask is not None:
            scores = scores + attention_mask
        probs = mx.softmax(scores, axis=-1).astype(q.dtype)
        out = probs @ v
        out = mx.transpose(out, (0, 2, 1, 3)).reshape(batch, seq_len, -1)
        return self.o_proj(out)


class EuroBertMLP(nn.Module):
    def __init__(self, config: EuroBertConfig):
        super().__init__()
        self.gate_proj = nn.Linear(
            config.hidden_size, config.intermediate_size, bias=config.mlp_bias
        )
        self.up_proj = nn.Linear(
            config.hidden_size, config.intermediate_size, bias=config.mlp_bias
        )
        self.down_proj = nn.Linear(
            config.intermediate_size, config.hidden_size, bias=config.mlp_bias
        )

    def __call__(self, x: mx.array) -> mx.array:
        return self.down_proj(_silu(self.gate_proj(x)) * self.up_proj(x))


class EuroBertDecoderLayer(nn.Module):
    def __init__(self, config: EuroBertConfig):
        super().__init__()
        self.self_attn = EuroBertAttention(config)
        self.mlp = EuroBertMLP(config)
        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = nn.RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    def __call__(
        self,
        hidden_states: mx.array,
        attention_mask: Optional[mx.array],
        position_embeddings: Tuple[mx.array, mx.array],
    ) -> mx.array:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(
            hidden_states, position_embeddings, attention_mask
        )
        hidden_states = residual + hidden_states

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


class EuroBertModel(nn.Module):
    def __init__(self, config: EuroBertConfig):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = [
            EuroBertDecoderLayer(config) for _ in range(config.num_hidden_layers)
        ]
        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = EuroBertRotaryEmbedding(config)

    def _attention_mask(self, attention_mask: Optional[mx.array]) -> Optional[mx.array]:
        if attention_mask is None:
            return None
        keep = attention_mask.astype(mx.bool_)
        mask = mx.where(keep[:, None, None, :], 0.0, mx.finfo(mx.float32).min)
        return mask.astype(mx.float32)

    def __call__(
        self,
        input_ids: mx.array,
        attention_mask: Optional[mx.array] = None,
        position_ids: Optional[mx.array] = None,
    ) -> mx.array:
        hidden_states = self.embed_tokens(input_ids)
        batch, seq_len = input_ids.shape
        if position_ids is None:
            position_ids = mx.broadcast_to(mx.arange(seq_len)[None, :], (batch, seq_len))

        mask = self._attention_mask(attention_mask)
        position_embeddings = self.rotary_emb(position_ids, hidden_states.dtype)
        for layer in self.layers:
            hidden_states = layer(hidden_states, mask, position_embeddings)
        return self.norm(hidden_states)


class EuroBertForTokenClassification(nn.Module):
    def __init__(self, config: EuroBertConfig):
        super().__init__()
        self.config = config
        self.model = EuroBertModel(config)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

    def __call__(
        self,
        input_ids: mx.array,
        attention_mask: Optional[mx.array] = None,
        position_ids: Optional[mx.array] = None,
    ) -> mx.array:
        hidden_states = self.model(input_ids, attention_mask, position_ids)
        return self.classifier(hidden_states)


def build_model(
    config: EuroBertConfig | Dict[str, Any],
    quantization: Optional[Dict[str, Any]] = None,
) -> EuroBertForTokenClassification:
    if isinstance(config, dict):
        config = EuroBertConfig(**config)
    model = EuroBertForTokenClassification(config)
    if quantization:
        nn.quantize(
            model,
            group_size=quantization.get("group_size", 64),
            bits=quantization["bits"],
            mode=quantization.get("mode", "affine"),
        )
    return model


def load_model(
    model_path: str | Path,
    variant: str = "bf16",
) -> EuroBertForTokenClassification:
    model_path = Path(model_path)
    with open(model_path / "config.json", "r", encoding="utf-8") as f:
        raw_config = json.load(f)
    with open(model_path / "mlx_config.json", "r", encoding="utf-8") as f:
        mlx_config = json.load(f)

    variant = variant.lower()
    if variant not in WEIGHT_FILES:
        raise ValueError(f"Unknown variant {variant!r}; expected one of {sorted(WEIGHT_FILES)}")
    weight_name = WEIGHT_FILES[variant]
    quantization = mlx_config["variants"].get(weight_name, {}).get("quantization")
    model = build_model(raw_config, quantization=quantization)
    model.load_weights(str(model_path / weight_name))
    mx.eval(model.parameters())
    return model