"""Cross-encoder specialist: state + question + one option in a single encoder pass."""

from __future__ import annotations

import json
import os

import torch
import torch.nn as nn
from safetensors.torch import load_file, save_file
from transformers import AutoConfig, AutoModel


class PlutoCrossEncoder(nn.Module):
    def __init__(self, backbone_name: str | None = None, backbone_config=None, vocab_size: int | None = None):
        super().__init__()
        if backbone_config is None:
            self.backbone = AutoModel.from_pretrained(backbone_name, attn_implementation="sdpa")
        else:
            self.backbone = AutoModel.from_config(backbone_config, attn_implementation="sdpa")
        if vocab_size is not None:
            self.backbone.resize_token_embeddings(vocab_size)
        h = self.backbone.config.hidden_size
        self.head = nn.Sequential(
            nn.Linear(h, h),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.LayerNorm(h),
            nn.Linear(h, 1),
        )

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0]
        return self.head(h).squeeze(-1).float()


def save_cross_encoder(model: PlutoCrossEncoder, tokenizer, out_dir: str, extra: dict) -> None:
    os.makedirs(out_dir, exist_ok=True)
    save_file({k: v.detach().cpu() for k, v in model.state_dict().items()}, os.path.join(out_dir, "model.safetensors"))
    tokenizer.save_pretrained(out_dir)
    model.backbone.config.to_json_file(os.path.join(out_dir, "backbone_config.json"))
    cfg = {
        "format": "pluto-ce-1",
        "vocab_size": model.backbone.config.vocab_size,
        **extra,
    }
    with open(os.path.join(out_dir, "pluto_ce_config.json"), "w", encoding="utf-8") as f:
        json.dump(cfg, f, indent=2)


def load_cross_encoder(path: str, device: str | torch.device = "cpu") -> tuple[PlutoCrossEncoder, dict]:
    cfg = json.load(open(os.path.join(path, "pluto_ce_config.json"), encoding="utf-8"))
    backbone_cfg = AutoConfig.for_model(
        **json.load(open(os.path.join(path, "backbone_config.json"), encoding="utf-8"))
    )
    model = PlutoCrossEncoder(backbone_config=backbone_cfg, vocab_size=cfg["vocab_size"])
    model.load_state_dict(load_file(os.path.join(path, "model.safetensors")), strict=True)
    model.to(device).eval()
    return model, cfg