# coding=utf-8
"""
A sentence-transformers input module that reproduces the LLM2Vec encoding
pipeline exactly: instruction handling, left-padded tokenization, and a
mean pooling that covers only the content tokens (BOS and instruction
tokens are excluded from the pooled vector).

This is a faithful port of `llm2vec.LLM2Vec.encode` -> the model produced
here yields the same embeddings as the original llm2vec-based usage, but
loads with a plain `SentenceTransformer(...)` call and needs no extra
package beyond `sentence-transformers` + `transformers`.
"""
from __future__ import annotations

from typing import Any

import torch
from transformers import AutoModel, AutoTokenizer

from sentence_transformers.base.modules.input_module import InputModule


class LLM2VecTransformer(InputModule):
    """LLM2Vec encoder packaged as a sentence-transformers module."""

    # Internal separator used by llm2vec to split instruction from content.
    SEP = "!@#$%^&*()"

    config_file_name: str = "sentence_bert_config.json"
    config_keys: list[str] = ["max_seq_length", "pooling_mode", "model_kwargs"]
    save_in_root: bool = True

    def __init__(
        self,
        model_name_or_path: str,
        max_seq_length: int = 8124,
        pooling_mode: str = "mean",
        model_kwargs: dict[str, Any] | None = None,
        tokenizer_kwargs: dict[str, Any] | None = None,
        hub_kwargs: dict[str, Any] | None = None,
        **kwargs,
    ) -> None:
        super().__init__()
        if pooling_mode != "mean":
            raise ValueError("Only 'mean' pooling is supported by this model.")
        self.max_seq_length = max_seq_length
        self.pooling_mode = pooling_mode
        # `model_kwargs` is persisted in the module config; `hub_kwargs`
        # (token, revision, trust_remote_code, ...) are loading-only.
        self.model_kwargs = model_kwargs or {}
        hub_kwargs = hub_kwargs or {}

        self.auto_model = AutoModel.from_pretrained(
            model_name_or_path, **{**self.model_kwargs, **hub_kwargs}
        )
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name_or_path, **{**(tokenizer_kwargs or {}), **hub_kwargs}
        )
        # llm2vec configures the tokenizer this way before encoding.
        self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.padding_side = "left"

    # -- encoding -----------------------------------------------------------
    def preprocess(
        self,
        inputs: list[str],
        prompt: str | None = None,
        **kwargs,
    ) -> dict[str, torch.Tensor | Any]:
        """Tokenize, reproducing llm2vec `_convert_to_str` + `tokenize`."""
        sep = self.SEP
        # llm2vec format: "<instruction> !@#$%^&*()<text>" for queries,
        # "!@#$%^&*()<text>" for documents (no instruction).
        if prompt:
            combined = [f"{prompt.strip()} {sep}{text}" for text in inputs]
        else:
            combined = [f"{sep}{text}" for text in inputs]

        # Split off the content (everything after the separator).
        content = [t.split(sep)[1] if len(t.split(sep)) > 1 else "" for t in combined]
        originals = ["".join(t.split(sep)) for t in combined]

        features = self.tokenizer(
            originals,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.max_seq_length,
        )

        # `embed_mask` marks the content tokens (the trailing tokens of each
        # sequence), excluding BOS, instruction and padding.
        embed_mask = torch.zeros_like(features["attention_mask"])
        for i, text in enumerate(content):
            ids = self.tokenizer(
                [text],
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=self.max_seq_length,
                add_special_tokens=False,
            )["input_ids"][0]
            if len(ids) > 0:
                embed_mask[i, -len(ids):] = 1
        features["embed_mask"] = embed_mask
        return features

    def forward(self, features: dict[str, Any], **kwargs) -> dict[str, Any]:
        """Run the encoder and mean-pool over content tokens only."""
        outputs = self.auto_model(
            input_ids=features["input_ids"],
            attention_mask=features["attention_mask"],
        )
        last_hidden = outputs.last_hidden_state
        embed_mask = features["embed_mask"]

        # Identical to llm2vec's mean pooling (left-padded, content at the end).
        seq_lengths = embed_mask.sum(dim=-1)
        embeddings = torch.stack(
            [
                last_hidden[i, -int(length):, :].mean(dim=0)
                for i, length in enumerate(seq_lengths)
            ],
            dim=0,
        )
        features["sentence_embedding"] = embeddings
        return features

    # -- introspection ------------------------------------------------------
    def get_sentence_embedding_dimension(self) -> int:
        return self.auto_model.config.hidden_size

    @property
    def max_seq_length_(self) -> int:  # pragma: no cover - convenience
        return self.max_seq_length

    # -- persistence --------------------------------------------------------
    def save(self, output_path: str, *args, safe_serialization: bool = True, **kwargs) -> None:
        self.auto_model.save_pretrained(output_path, safe_serialization=safe_serialization)
        self.tokenizer.save_pretrained(output_path)
        self.save_config(output_path)

    @classmethod
    def load(
        cls,
        model_name_or_path: str,
        subfolder: str = "",
        token: bool | str | None = None,
        cache_folder: str | None = None,
        revision: str | None = None,
        local_files_only: bool = False,
        trust_remote_code: bool = False,
        model_kwargs: dict[str, Any] | None = None,
        processor_kwargs: dict[str, Any] | None = None,
        config_kwargs: dict[str, Any] | None = None,
        backend: str = "torch",
        **kwargs,
    ) -> "LLM2VecTransformer":
        config = cls.load_config(
            model_name_or_path=model_name_or_path,
            subfolder=subfolder,
            token=token,
            cache_folder=cache_folder,
            revision=revision,
            local_files_only=local_files_only,
        )
        hub_kwargs = {
            "token": token,
            "cache_dir": cache_folder,
            "revision": revision,
            "local_files_only": local_files_only,
            "trust_remote_code": trust_remote_code,
        }
        persistent_model_kwargs = {
            **config.get("model_kwargs", {}),
            **(model_kwargs or {}),
        }
        return cls(
            model_name_or_path,
            max_seq_length=config.get("max_seq_length", 8124),
            pooling_mode=config.get("pooling_mode", "mean"),
            model_kwargs=persistent_model_kwargs,
            tokenizer_kwargs=processor_kwargs or {},
            hub_kwargs=hub_kwargs,
        )