"""Sentence Transformers adapter for the complete, flat SlideOn checkpoint."""
from pathlib import Path

import torch
from sentence_transformers.base.modules import InputModule
from transformers import AutoConfig, AutoModel, AutoTokenizer


class SlideOn(InputModule):
    config_file_name = "config.json"
    forward_kwargs = {"task", "max_active_dims", "token_k"}
    save_in_root = True

    def __init__(self, checkpoint, model_kwargs=None, processor_kwargs=None,
                 config_kwargs=None, **hub_kwargs):
        super().__init__()
        config = AutoConfig.from_pretrained(
            checkpoint, trust_remote_code=True, **hub_kwargs, **(config_kwargs or {}))
        if config.model_type != "slideon":
            raise ValueError("Expected a flat SlideOn checkpoint; use slideon.export for legacy checkpoints")
        self.model = AutoModel.from_pretrained(
            checkpoint, config=config, trust_remote_code=True,
            **hub_kwargs, **(model_kwargs or {}))
        self.tokenizer = AutoTokenizer.from_pretrained(
            checkpoint, trust_remote_code=True, **hub_kwargs, **(processor_kwargs or {}))
        if self.tokenizer.mask_token_id is not None:
            self.tokenizer.pad_token_id = self.tokenizer.mask_token_id
        elif self.tokenizer.eos_token_id is not None:
            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
        self.settings = self.model.config.inference
        self.encoder_config = self.model.config.tokenizer_settings
        self.max_seq_length = max(self.settings["query_length"], self.settings["document_length"])

    def get_embedding_dimension(self):
        return self.model.get_embedding_dimension()

    @staticmethod
    def _task(task):
        task = task or "document"
        if task not in ("query", "document"):
            raise ValueError("task must be 'query' or 'document'")
        return task

    def preprocess(
        self, inputs, prompt=None, task=None, processing_kwargs=None, **kwargs
    ):
        task = self._task(task)
        if not all(isinstance(text, str) for text in inputs):
            raise TypeError("SlideOn accepts text strings only")
        if prompt is None:
            prompt = self.settings["prompts"][task]
        texts = [(prompt + text).strip() for text in inputs]
        prefix = self.encoder_config.get(task + "_prefix", "")
        length = min(self.max_seq_length, self.settings[task + "_length"])
        options = dict(
            padding=True,
            truncation=True,
            max_length=length - bool(prefix),
            return_tensors="pt",
        )
        options.update(processing_kwargs or {})
        features = dict(self.tokenizer(texts, **options))
        if prefix:
            prefix_id = self.tokenizer.convert_tokens_to_ids(prefix)
            for name, value in list(features.items()):
                fill = (
                    prefix_id
                    if name == "input_ids"
                    else (1 if name == "attention_mask" else 0)
                )
                marker = torch.full(
                    (value.shape[0], 1), fill, dtype=value.dtype, device=value.device
                )
                features[name] = torch.cat((value[:, :1], marker, value[:, 1:]), dim=1)
        features["slideon_task"] = task
        return features

    def forward(self, features, task=None, max_active_dims=None, token_k=None, **kwargs):
        if "top_k" in kwargs:
            raise TypeError("top_k has been removed; use max_active_dims (None for unpruned)")
        inputs = {key: features[key] for key in
                  ("input_ids", "attention_mask", "token_type_ids", "slideon_task") if key in features}
        features.update(self.model(**inputs, task=task, max_active_dims=max_active_dims, token_k=token_k))
        return features

    def save(self, output_path, safe_serialization=True, **kwargs):
        self.model.save_pretrained(output_path, safe_serialization=safe_serialization)
        self.tokenizer.save_pretrained(output_path)

    @classmethod
    def load(cls, model_name_or_path, subfolder="", token=None, cache_folder=None,
             revision=None, local_files_only=False, model_kwargs=None,
             processor_kwargs=None, config_kwargs=None, backend="torch", **kwargs):
        if backend != "torch":
            raise ValueError("SlideOn currently supports the torch backend only")
        return cls(model_name_or_path, model_kwargs, processor_kwargs, config_kwargs,
                   subfolder=subfolder, token=token, cache_dir=cache_folder,
                   revision=revision, local_files_only=local_files_only)