from __future__ import annotations

import base64
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import torch
from safetensors import safe_open
from torch import nn
from torch.nn import functional as F


@dataclass
class ModelConfig:
    model_name: str = "Fabric 1.5"
    architecture: str = "fabric"
    vocab_size: int = 65536
    hidden_size: int = 1536
    intermediate_size: int = 4096
    num_layers: int = 24
    num_query_heads: int = 24
    num_kv_heads: int = 6
    head_dim: int = 64
    sequence_length: int = 32768
    local_attention_window: int = 2048
    memory_chunk_size: int = 512
    summaries_per_chunk: int = 4
    rope_theta: float = 1000000.0
    rms_norm_eps: float = 1e-6
    tie_word_embeddings: bool = True
    attention_backend: str = "auto"
    attention_chunk_size: int = 1024
    activation_checkpointing: bool = False
    chunked_cross_entropy: bool = True
    loss_chunk_size: int = 1024


def _decode_structure(value: Any, tensors: dict[str, torch.Tensor]) -> Any:
    if not isinstance(value, dict) or "__kind__" not in value:
        return value
    kind = value["__kind__"]
    if kind == "tensor":
        return tensors[value["key"]]
    if kind == "dict":
        return {
            _decode_structure(key, tensors): _decode_structure(item, tensors)
            for key, item in value["items"]
        }
    if kind == "tuple":
        return tuple(_decode_structure(item, tensors) for item in value["items"])
    if kind == "list":
        return [_decode_structure(item, tensors) for item in value["items"]]
    if kind == "ndarray":
        return np.asarray(value["items"], dtype=np.dtype(value["dtype"])).reshape(value["shape"])
    if kind == "path":
        return Path(value["value"])
    if kind == "bytes":
        return base64.b64decode(value["value"])
    raise ValueError(f"unknown checkpoint structure kind: {kind}")


def load_checkpoint(path: str | Path, map_location: str | torch.device = "cpu") -> dict[str, Any]:
    with safe_open(path, framework="pt", device=str(map_location)) as handle:
        metadata = handle.metadata()
        if metadata.get("format") != "fabric_complete_checkpoint":
            raise ValueError("file is not a Fabric complete checkpoint")
        tensors = {key: handle.get_tensor(key) for key in handle.keys()}
        structure = json.loads(metadata["structure"])
    state = _decode_structure(structure, tensors)
    if not isinstance(state, dict):
        raise ValueError("checkpoint root must be a dictionary")
    return state


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

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


class SwiGLU(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


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


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, q: torch.Tensor, k: torch.Tensor, position_ids: torch.Tensor):
        angles = position_ids.float().unsqueeze(-1) * self.inv_freq.float()
        emb = torch.cat((angles, angles), dim=-1)
        cos = emb.cos().to(q.dtype).unsqueeze(1)
        sin = emb.sin().to(q.dtype).unsqueeze(1)
        return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin


def build_local_causal_mask(query_length: int, key_length: int, window: int, device, query_offset: int = 0):
    query_positions = torch.arange(query_offset, query_offset + query_length, device=device)
    key_positions = torch.arange(key_length, device=device)
    return (key_positions[None, :] <= query_positions[:, None]) & (
        key_positions[None, :] > query_positions[:, None] - window
    )


def repeat_kv(x: torch.Tensor, groups: int) -> torch.Tensor:
    if groups == 1:
        return x
    batch, kv_heads, length, dim = x.shape
    return x[:, :, None, :, :].expand(batch, kv_heads, groups, length, dim).reshape(
        batch, kv_heads * groups, length, dim
    )


def reference_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, allowed_mask: torch.Tensor):
    scores = torch.matmul(q.float(), k.float().transpose(-1, -2)) / math.sqrt(q.shape[-1])
    scores = scores.masked_fill(~allowed_mask, torch.finfo(scores.dtype).min)
    probabilities = torch.softmax(scores, dim=-1)
    probabilities = torch.where(allowed_mask.any(dim=-1, keepdim=True), probabilities, 0.0)
    return torch.matmul(probabilities.to(v.dtype), v)


class GQAAttention(nn.Module):
    def __init__(self, config: ModelConfig, window: int | None = None) -> None:
        super().__init__()
        self.num_query_heads = config.num_query_heads
        self.num_kv_heads = config.num_kv_heads
        self.head_dim = config.head_dim
        self.groups = config.num_query_heads // config.num_kv_heads
        self.window = window or config.sequence_length
        self.backend = "sdpa" if config.attention_backend == "flash_attn" else config.attention_backend
        self.attention_chunk_size = config.attention_chunk_size
        self.q_proj = nn.Linear(config.hidden_size, config.num_query_heads * config.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, config.num_kv_heads * config.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, config.num_kv_heads * config.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_query_heads * config.head_dim, config.hidden_size, bias=False)
        self.rope = RotaryEmbedding(config.head_dim, config.rope_theta)

    def forward(self, x: torch.Tensor, position_ids: torch.Tensor | None = None) -> torch.Tensor:
        batch, length, _ = x.shape
        if position_ids is None:
            position_ids = torch.arange(length, device=x.device).expand(batch, -1)
        q = self.q_proj(x).view(batch, length, self.num_query_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch, length, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch, length, self.num_kv_heads, self.head_dim).transpose(1, 2)
        q, k = self.rope(q, k, position_ids)
        k = repeat_kv(k, self.groups)
        v = repeat_kv(v, self.groups)
        use_sdpa = self.backend in {"auto", "sdpa"} and hasattr(F, "scaled_dot_product_attention")
        if use_sdpa:
            outputs = []
            for start in range(0, length, self.attention_chunk_size):
                end = min(start + self.attention_chunk_size, length)
                key_start = max(0, start - self.window + 1)
                key_end = end
                allowed = build_local_causal_mask(
                    end - start,
                    key_end - key_start,
                    self.window,
                    x.device,
                    query_offset=start - key_start,
                )[None, None]
                outputs.append(
                    F.scaled_dot_product_attention(
                        q[:, :, start:end],
                        k[:, :, key_start:key_end],
                        v[:, :, key_start:key_end],
                        attn_mask=allowed,
                        dropout_p=0.0,
                    )
                )
            output = torch.cat(outputs, dim=2)
        else:
            allowed = build_local_causal_mask(length, length, self.window, x.device)[None, None]
            output = reference_attention(q, k, v, allowed)
        output = output.transpose(1, 2).contiguous().view(batch, length, -1)
        return self.o_proj(output)


def build_completed_chunk_mask(sequence_length: int, num_chunks: int, summaries_per_chunk: int, chunk_size: int, device):
    query_chunk = torch.arange(sequence_length, device=device) // chunk_size
    summary_chunk = torch.arange(num_chunks, device=device).repeat_interleave(summaries_per_chunk)
    return summary_chunk[None, :] < query_chunk[:, None]


class ChunkSummarizer(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.chunk_size = config.memory_chunk_size
        self.num_summaries = config.summaries_per_chunk
        self.hidden_size = config.hidden_size
        self.queries = nn.Parameter(torch.empty(self.num_summaries, self.hidden_size))
        nn.init.normal_(self.queries, std=0.02)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, length, hidden = x.shape
        num_chunks = (length + self.chunk_size - 1) // self.chunk_size
        padded_length = num_chunks * self.chunk_size
        if padded_length != length:
            x = torch.cat((x, x.new_zeros(batch, padded_length - length, hidden)), dim=1)
        chunks = x.view(batch, num_chunks, self.chunk_size, hidden)
        scores = torch.einsum("mh,bnch->bnmc", self.queries.float(), chunks.float()) / math.sqrt(hidden)
        if padded_length != length:
            valid = torch.arange(padded_length, device=x.device).view(num_chunks, self.chunk_size) < length
            scores = scores.masked_fill(~valid[None, :, None, :], torch.finfo(scores.dtype).min)
        weights = torch.softmax(scores, dim=-1).to(chunks.dtype)
        return torch.einsum("bnmc,bnch->bnmh", weights, chunks)


class MemoryAttention(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.num_query_heads = config.num_query_heads
        self.num_kv_heads = config.num_kv_heads
        self.head_dim = config.head_dim
        self.groups = config.num_query_heads // config.num_kv_heads
        self.chunk_size = config.memory_chunk_size
        self.num_summaries = config.summaries_per_chunk
        self.q_proj = nn.Linear(config.hidden_size, config.num_query_heads * config.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, config.num_kv_heads * config.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, config.num_kv_heads * config.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_query_heads * config.head_dim, config.hidden_size, bias=False)

    def forward(self, x: torch.Tensor, summaries: torch.Tensor) -> torch.Tensor:
        batch, length, _ = x.shape
        num_chunks = summaries.shape[1]
        flat = summaries.reshape(batch, num_chunks * self.num_summaries, -1)
        q = self.q_proj(x).view(batch, length, self.num_query_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(flat).view(batch, -1, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(flat).view(batch, -1, self.num_kv_heads, self.head_dim).transpose(1, 2)
        k, v = repeat_kv(k, self.groups), repeat_kv(v, self.groups)
        scores = torch.matmul(q.float(), k.float().transpose(-1, -2)) / math.sqrt(self.head_dim)
        allowed = build_completed_chunk_mask(length, num_chunks, self.num_summaries, self.chunk_size, x.device)[None, None]
        scores = scores.masked_fill(~allowed, torch.finfo(scores.dtype).min)
        probabilities = torch.softmax(scores, dim=-1)
        probabilities = torch.where(allowed.any(dim=-1, keepdim=True), probabilities, 0.0)
        output = torch.matmul(probabilities.to(v.dtype), v)
        output = output.transpose(1, 2).contiguous().view(batch, length, -1)
        return self.o_proj(output)


class LocalBlock(nn.Module):
    def __init__(self, config: ModelConfig, window: int | None = None) -> None:
        super().__init__()
        self.attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attention = GQAAttention(config, window or config.local_attention_window)
        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = SwiGLU(config.hidden_size, config.intermediate_size)

    def forward(self, x: torch.Tensor, position_ids: torch.Tensor | None = None) -> torch.Tensor:
        x = x + self.attention(self.attention_norm(x), position_ids)
        return x + self.mlp(self.mlp_norm(x))


class FabricMemoryBlock(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.local_attention = GQAAttention(config, config.local_attention_window)
        self.summarizer = ChunkSummarizer(config)
        self.memory_attention = MemoryAttention(config)
        self.gate = nn.Linear(config.hidden_size, 1, bias=True)
        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = SwiGLU(config.hidden_size, config.intermediate_size)

    def forward(self, x: torch.Tensor, position_ids: torch.Tensor | None = None) -> torch.Tensor:
        normalized = self.attention_norm(x)
        local = self.local_attention(normalized, position_ids)
        summaries = self.summarizer(normalized)
        memory = self.memory_attention(normalized, summaries)
        gate = torch.sigmoid(self.gate(normalized))
        x = x + gate * local + (1.0 - gate) * memory
        return x + self.mlp(self.mlp_norm(x))


@dataclass
class CausalLMOutput:
    logits: torch.Tensor | None
    loss: torch.Tensor | None = None


class FabricCoreForCausalLM(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        if config.architecture == "fabric":
            layers = [FabricMemoryBlock(config) if i % 3 == 2 else LocalBlock(config) for i in range(config.num_layers)]
        else:
            window = config.sequence_length if config.architecture == "full" else config.local_attention_window
            layers = [LocalBlock(config, window) for _ in range(config.num_layers)]
        self.layers = nn.ModuleList(layers)
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight

    def forward(self, input_ids: torch.Tensor, labels: torch.Tensor | None = None) -> CausalLMOutput:
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence]")
        if input_ids.shape[1] > self.config.sequence_length:
            raise ValueError("input sequence exceeds configured sequence_length")
        position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).expand(input_ids.shape[0], -1)
        hidden = self.embed_tokens(input_ids)
        for layer in self.layers:
            hidden = layer(hidden, position_ids)
        normalized = self.norm(hidden)
        logits = self.lm_head(normalized).float()
        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, logits.shape[-1]),
                labels[:, 1:].reshape(-1),
                ignore_index=-100,
            )
        return CausalLMOutput(logits=logits, loss=loss)