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

IM_START   = 151644
IM_END     = 151645
EOT        = 151643
THINK_END  = 151668

EXOCORE_IDENTITY = (
    "You are ExocoreV1, a highly capable AI built from scratch by Exocore. "
    "You are knowledgeable, analytical, direct, and precise. "
    "You reason carefully before answering. "
    "You identify yourself only as ExocoreV1."
)


def build_chat_prompt(messages, think=False, prior_thinking=None, search_context=None):
    identity = EXOCORE_IDENTITY
    if search_context:
        identity += f"\n\nResearch context:\n{search_context}"

    prompt = f"<|im_start|>system\n{identity}<|im_end|>\n"
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"

    if prior_thinking:
        prompt += f"<|im_start|>assistant\n<think>\n{prior_thinking}\n\n</think>\n"
    elif think:
        prompt += "<|im_start|>assistant\n<think>\n"
    else:
        prompt += "<|im_start|>assistant\n<think>\n\n</think>\n"
    return prompt


class ExocoreRMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        xf = x.float()
        norm = (xf.pow(2).mean(-1, keepdim=True) + self.eps).rsqrt()
        return (xf * norm).to(x.dtype) * self.weight


class ExocoreRotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq=40960, base=1000000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        t = torch.arange(max_seq).float()
        freqs = torch.outer(t, inv_freq)
        cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1)
        sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

    def apply_rope(self, q, k):
        T = q.shape[2]
        cos = self.rope_cos[:T].unsqueeze(0).unsqueeze(0)
        sin = self.rope_sin[:T].unsqueeze(0).unsqueeze(0)

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

        q_rot = q * cos + rotate_half(q) * sin
        k_rot = k * cos + rotate_half(k) * sin
        return q_rot.to(q.dtype), k_rot.to(k.dtype)


class ExocoreAttention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_heads    = cfg["num_attention_heads"]
        self.n_kv_heads = cfg["num_key_value_heads"]
        self.head_dim   = cfg["head_dim"]
        hidden          = cfg["hidden_size"]
        eps             = cfg["rms_norm_eps"]

        self.q_proj = nn.Linear(hidden, self.n_heads    * self.head_dim, bias=False)
        self.k_proj = nn.Linear(hidden, self.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(hidden, self.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_heads * self.head_dim, hidden,    bias=False)
        self.q_norm = ExocoreRMSNorm(self.head_dim, eps=eps)
        self.k_norm = ExocoreRMSNorm(self.head_dim, eps=eps)

    def forward(self, x, rope, mask=None):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads,    self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)

        q = self.q_norm(q)
        k = self.k_norm(k)
        q, k = rope.apply_rope(q, k)

        groups = self.n_heads // self.n_kv_heads
        k = k.repeat_interleave(groups, dim=1)
        v = v.repeat_interleave(groups, dim=1)

        scale = self.head_dim ** -0.5
        attn  = torch.matmul(q, k.transpose(-2, -1)) * scale
        if mask is not None:
            attn = attn + mask
        attn = F.softmax(attn.float(), dim=-1).to(q.dtype)

        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim)
        return self.o_proj(out)


class ExocoreMLP(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.gate_proj = nn.Linear(cfg["hidden_size"], cfg["intermediate_size"], bias=False)
        self.up_proj   = nn.Linear(cfg["hidden_size"], cfg["intermediate_size"], bias=False)
        self.down_proj = nn.Linear(cfg["intermediate_size"], cfg["hidden_size"], bias=False)

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


class ExocoreBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.input_layernorm          = ExocoreRMSNorm(cfg["hidden_size"], eps=cfg["rms_norm_eps"])
        self.post_attention_layernorm = ExocoreRMSNorm(cfg["hidden_size"], eps=cfg["rms_norm_eps"])
        self.attn = ExocoreAttention(cfg)
        self.mlp  = ExocoreMLP(cfg)

    def forward(self, x, rope, mask=None):
        x = x + self.attn(self.input_layernorm(x), rope, mask)
        x = x + self.mlp(self.post_attention_layernorm(x))
        return x


class ExocoreLM(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg          = cfg
        self.embed_tokens = nn.Embedding(cfg["vocab_size"], cfg["hidden_size"])
        self.layers       = nn.ModuleList([ExocoreBlock(cfg) for _ in range(cfg["num_hidden_layers"])])
        self.norm         = ExocoreRMSNorm(cfg["hidden_size"], eps=cfg["rms_norm_eps"])
        self.lm_head      = nn.Linear(cfg["hidden_size"], cfg["vocab_size"], bias=False)
        self.rope         = ExocoreRotaryEmbedding(
            cfg["head_dim"],
            max_seq=cfg.get("max_position_embeddings", 40960),
            base=cfg.get("rope_theta", 1000000),
        )

    def forward(self, input_ids):
        B, T = input_ids.shape
        x    = self.embed_tokens(input_ids)
        mask = torch.triu(torch.full((T, T), float("-inf"), device=x.device), diagonal=1)
        mask = mask.unsqueeze(0).unsqueeze(0)
        for block in self.layers:
            x = block(x, self.rope, mask)
        x = self.norm(x)
        return self.lm_head(x), None