"""Encoder backbone (ModernBERT architecture) for Glide checkpoints whose decision_config.json has "backbone": "encoder".

The prompt TEXT is the same as for the decoder Glide models: it is rendered with the checkpoint's render tokenizer (render_tokenizer/: chat
template and single-token option codes), then tokenised by the encoder tokenizer as [CLS] + ids + [SEP]. Prompts longer than "max_len" tokens keep
the first "trunc_head" and the last (max_len - 2 - trunc_head) tokens. Every layer attends bidirectionally over the non-pad tokens; the readout reads
the last real token before [SEP] (the final "(" of "Answer: (", pooling "last") or [CLS] (pooling "cls").

Tensor names (HF convention, no backbone.* keys): model.embeddings.*, model.layers.<i>.*, model.final_norm.* (as in ModernBertForMaskedLM; the MLM
head is not included) in model-*.safetensors; the 255-way decision head in readout.safetensors {"weight": [255, hidden]}.
"""
from __future__ import annotations

import json
import os

import torch

ENC_PREFIX = "model."
RENDER_DIR = "render_tokenizer"
MAX_OPTIONS = 255


class EncoderTokenizer:
    """Prompt-side calls (chat template, vocab, encode of option codes) go to the render tokenizer; __call__ returns encoder ids."""

    def __init__(self, render, enc, max_len, trunc_head):
        self.render, self.enc = render, enc
        self.max_len, self.trunc_head = int(max_len), int(trunc_head)
        self.cls_id, self.sep_id, self.pad_token_id = enc.cls_token_id, enc.sep_token_id, enc.pad_token_id

    def apply_chat_template(self, *args, **kwargs):
        return self.render.apply_chat_template(*args, **kwargs)

    def get_vocab(self):
        return self.render.get_vocab()

    def encode(self, text, add_special_tokens=False, **kwargs):
        return self.render.encode(text, add_special_tokens=add_special_tokens, **kwargs)

    def _wrap(self, ids):
        n = self.max_len - 2
        if len(ids) > n:
            ids = ids[: self.trunc_head] + ids[len(ids) - (n - self.trunc_head):]
        return [self.cls_id] + list(ids) + [self.sep_id]

    def __call__(self, texts, add_special_tokens=False, **kwargs):
        ids = self.enc(list(texts), add_special_tokens=False)["input_ids"]
        return {"input_ids": [self._wrap(x) for x in ids]}


class EncoderNetwork(torch.nn.Module):
    def __init__(self, config, pooling="last"):
        from transformers.models.modernbert.modeling_modernbert import ModernBertModel
        super().__init__()
        self.model = ModernBertModel(config)
        self.readout = torch.nn.Linear(config.hidden_size, MAX_OPTIONS, bias=False)
        self.pooling = pooling
        self.visual = None

    def forward(self, input_ids, padding, pixel_values=None, image_grid_thw=None):
        if pixel_values is not None:
            raise ValueError("encoder Glide checkpoints are text-only")
        pos = (padding.long().cumsum(-1) - 1).clamp_(min=0)
        h = self.model(input_ids=input_ids, attention_mask=padding, position_ids=pos).last_hidden_state
        if self.pooling == "cls":
            idx = (padding.shape[1] - padding.long().sum(-1)).clamp_(max=padding.shape[1] - 1)
            v = h[torch.arange(h.shape[0], device=h.device), idx]
        else:
            v = h[:, -2]                    # left padding: [..., "(", [SEP]]
        return self.readout(v)


def load(path, config, device, dtype):
    """-> (EncoderNetwork in eval mode, EncoderTokenizer)."""
    from safetensors import safe_open
    from safetensors.torch import load_file
    from transformers import AutoConfig, AutoTokenizer
    render = AutoTokenizer.from_pretrained(os.path.join(path, config.get("render_tokenizer") or RENDER_DIR))
    enc = AutoTokenizer.from_pretrained(path)
    tok = EncoderTokenizer(render, enc, config.get("max_len", 7999), config.get("trunc_head", 1024))
    cfg = AutoConfig.from_pretrained(path)
    cfg._attn_implementation = "sdpa"
    net = EncoderNetwork(cfg, config.get("pooling", "last")).to(device=device, dtype=dtype)
    with open(os.path.join(path, "model.safetensors.index.json")) as f:
        files = sorted(set(json.load(f)["weight_map"].values()))
    where = {}
    handles = {f: safe_open(os.path.join(path, f), framework="pt", device="cpu") for f in files}
    for f, h in handles.items():
        for k in h.keys():
            where[k] = f
    missing, used = [], set()
    with torch.no_grad():
        for n, p in net.model.named_parameters():
            k = ENC_PREFIX + n
            if k not in where:
                missing.append(k)
                continue
            p.copy_(handles[where[k]].get_tensor(k).to(device=device, dtype=dtype))
            used.add(k)
        net.readout.weight.copy_(load_file(os.path.join(path, "readout.safetensors"))["weight"].to(device=device, dtype=dtype))
    if missing:
        raise RuntimeError(f"missing encoder weights: {missing[:5]} ({len(missing)} total)")
    extra = sorted(set(where) - used)
    if extra:
        raise RuntimeError(f"unused tensors in the checkpoint: {extra[:5]} ({len(extra)} total)")
    net.eval().requires_grad_(False)
    return net, tok