from __future__ import annotations

import json
from collections.abc import Mapping

from torch import Tensor, nn


def serialize_state(state: Mapping[str, object] | str) -> str:
    """Serialize structured state in stable key order for the text encoder."""
    if isinstance(state, str):
        return state
    parts = []
    for key in sorted(state):
        value = state[key]
        if value is None:
            continue
        if isinstance(value, (Mapping, list, tuple)):
            value = json.dumps(value, ensure_ascii=False, sort_keys=True)
        parts.append(f"[{key}] {value}")
    return "\n".join(parts)


class TransformerTextEncoder(nn.Module):
    """Shared transformer encoder with attention-mask-aware mean pooling."""

    def __init__(
        self, model: nn.Module, tokenizer: object, hidden_size: int, max_length: int = 512
    ):
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
        self.hidden_size = hidden_size
        self.max_length = max_length

    @classmethod
    def from_pretrained(cls, name: str, max_length: int = 512) -> "TransformerTextEncoder":
        from transformers import AutoModel, AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(name)
        model = AutoModel.from_pretrained(name)
        return cls(model, tokenizer, model.config.hidden_size, max_length)

    def forward(self, texts: list[str]) -> Tensor:
        batch = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt",
        )
        device = next(self.model.parameters()).device
        batch = {key: value.to(device) for key, value in batch.items()}
        outputs = self.model(**batch)
        token_embeddings = outputs.last_hidden_state
        mask = batch["attention_mask"].unsqueeze(-1).to(token_embeddings.dtype)
        pooled = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1)
        return nn.functional.layer_norm(pooled, (self.hidden_size,))