"""
Qwen3-Reranker backend using packed inference with Transformers.

This backend implements the sequence stitching optimization described in
Qwen3-Reranker packed inference examples:
1. Share the query/instruction prefix across many documents.
2. Reset document ``position_ids`` relative to the shared prefix.
3. Use a custom causal attention mask so each document can attend to the
   prefix and itself, but never to other documents.

Compared with the standard per-pair batching path, this reduces repeated
prefix computation and removes inter-sample padding waste. For online search
requests like ``1 query + 400 docs``, the backend further packs documents into
multiple chunks under a configurable total token budget.
"""

from __future__ import annotations

import logging
import threading
import time
from typing import Any, Dict, List, Sequence, Tuple

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

logger = logging.getLogger("reranker.backends.qwen3_transformers_packed")

_DEFAULT_PREFIX = (
    "<|im_start|>system\n"
    "Judge whether the Document meets the requirements based on the Query and the Instruct "
    'provided. Note that the answer can only be "yes" or "no".'
    "<|im_end|>\n<|im_start|>user\n"
)
_DEFAULT_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
_DEFAULT_PAIR_PREFIX_TEMPLATE = "{prefix}<Instruct>: {instruction}\n<Query>: {query}\n<Document>: "


def _deduplicate_with_positions(texts: Sequence[str]) -> Tuple[List[str], List[int]]:
    unique_texts: List[str] = []
    position_to_unique: List[int] = []
    seen: Dict[str, int] = {}

    for text in texts:
        idx = seen.get(text)
        if idx is None:
            idx = len(unique_texts)
            seen[text] = idx
            unique_texts.append(text)
        position_to_unique.append(idx)

    return unique_texts, position_to_unique


class Qwen3TransformersPackedRerankerBackend:
    """
    Qwen3-Reranker packed inference backend using Transformers.

    Config from ``services.rerank.backends.qwen3_transformers_packed``.
    """

    def __init__(self, config: Dict[str, Any]) -> None:
        self._config = config or {}
        model_name = str(self._config.get("model_name") or "Qwen/Qwen3-Reranker-0.6B")
        self._instruction = str(
            self._config.get("instruction")
            or "Rank products by query with category & style match prioritized"
        )
        self._prefix = str(self._config.get("prompt_prefix") or _DEFAULT_PREFIX)
        self._suffix = str(self._config.get("prompt_suffix") or _DEFAULT_SUFFIX)
        self._pair_prefix_template = str(
            self._config.get("pair_prefix_template") or _DEFAULT_PAIR_PREFIX_TEMPLATE
        )

        max_model_len = int(self._config.get("max_model_len", 4096))
        max_doc_len = int(self._config.get("max_doc_len", 160))
        max_docs_per_pack = int(self._config.get("max_docs_per_pack", 0))
        use_fp16 = bool(self._config.get("use_fp16", True))
        device = self._config.get("device")
        attn_impl = str(self._config.get("attn_implementation") or "eager").strip()
        sort_by_doc_length = self._config.get("sort_by_doc_length", True)

        self._model_name = model_name
        self._max_model_len = max_model_len
        self._max_doc_len = max_doc_len
        self._max_docs_per_pack = max_docs_per_pack
        self._sort_by_doc_length = str(sort_by_doc_length).strip().lower() in {
            "1",
            "true",
            "yes",
            "y",
            "on",
        }
        self._attn_impl = attn_impl

        logger.info(
            "[Qwen3_Transformers_Packed] Loading model %s (max_model_len=%s, max_doc_len=%s, "
            "max_docs_per_pack=%s, fp16=%s, attn_impl=%s)",
            model_name,
            max_model_len,
            max_doc_len,
            max_docs_per_pack,
            use_fp16,
            attn_impl,
        )

        self._tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
        self._tokenizer.pad_token = self._tokenizer.eos_token

        self._prefix_tokens = self._tokenizer.encode(self._prefix, add_special_tokens=False)
        self._suffix_tokens = self._tokenizer.encode(self._suffix, add_special_tokens=False)
        self._suffix_len = len(self._suffix_tokens)

        if not torch.cuda.is_available():
            raise RuntimeError(
                "qwen3_transformers_packed backend requires CUDA GPU, "
                "but torch.cuda.is_available() is False"
            )

        kwargs: Dict[str, Any] = {}
        if use_fp16:
            kwargs["torch_dtype"] = torch.float16
        if attn_impl:
            kwargs["attn_implementation"] = attn_impl

        self._model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs).eval()
        target_device = str(device).strip() if device is not None else "cuda"
        if not target_device.startswith("cuda"):
            raise ValueError(
                "qwen3_transformers_packed backend is GPU-only. "
                f"Unsupported device setting: {target_device!r}"
            )
        self._model = self._model.to(target_device)
        self._device = next(self._model.parameters()).device
        if self._device.type != "cuda":
            raise RuntimeError(
                "qwen3_transformers_packed backend failed to place model on CUDA. "
                f"Current device: {self._device}"
            )

        self._token_true_id = self._tokenizer.convert_tokens_to_ids("yes")
        self._token_false_id = self._tokenizer.convert_tokens_to_ids("no")
        if self._token_true_id is None or self._token_false_id is None:
            raise RuntimeError("Failed to resolve Qwen3 reranker classifier token ids for yes/no")

        prefix_budget = len(self._prefix_tokens) + self._suffix_len + 1
        if self._max_model_len <= prefix_budget:
            raise ValueError(
                "max_model_len is too small for packed reranking. "
                f"Need > {prefix_budget}, got {self._max_model_len}."
            )
        if self._max_doc_len <= 0:
            raise ValueError(f"max_doc_len must be > 0, got {self._max_doc_len}")
        if self._max_docs_per_pack < 0:
            raise ValueError(
                f"max_docs_per_pack must be >= 0, got {self._max_docs_per_pack}"
            )

        self._infer_lock = threading.Lock()

        logger.info(
            "[Qwen3_Transformers_Packed] Model ready | model=%s device=%s",
            model_name,
            self._device,
        )

    def _build_pair_prefix_tokens(self, query: str) -> List[int]:
        pair_prefix = self._pair_prefix_template.format(
            prefix=self._prefix,
            instruction=self._instruction,
            query=query,
        )
        return self._tokenizer.encode(pair_prefix, add_special_tokens=False)

    def _tokenize_documents(self, docs: Sequence[str], query_prefix_len: int) -> List[List[int]]:
        max_doc_tokens = min(
            self._max_doc_len,
            max(1, self._max_model_len - query_prefix_len - self._suffix_len),
        )
        tokenized = self._tokenizer(
            list(docs),
            padding=False,
            truncation=True,
            max_length=max_doc_tokens,
            add_special_tokens=False,
            return_attention_mask=False,
        )
        return [list(ids) for ids in tokenized["input_ids"]]

    def _build_pack_plan(
        self,
        query_prefix_len: int,
        doc_tokens: Sequence[Sequence[int]],
    ) -> List[List[int]]:
        order = list(range(len(doc_tokens)))
        if self._sort_by_doc_length and len(order) > 1:
            order.sort(key=lambda idx: len(doc_tokens[idx]))

        packs: List[List[int]] = []
        current_pack: List[int] = []
        current_len = query_prefix_len
        for idx in order:
            packed_doc_len = len(doc_tokens[idx]) + self._suffix_len
            if packed_doc_len <= 0:
                continue

            over_docs_cap = self._max_docs_per_pack > 0 and len(current_pack) >= self._max_docs_per_pack
            over_token_cap = current_pack and (current_len + packed_doc_len > self._max_model_len)
            if over_docs_cap or over_token_cap:
                packs.append(current_pack)
                current_pack = []
                current_len = query_prefix_len

            if query_prefix_len + packed_doc_len > self._max_model_len:
                raise ValueError(
                    "Packed doc still exceeds max_model_len after truncation. "
                    f"query_prefix_len={query_prefix_len}, doc_len={packed_doc_len}, "
                    f"max_model_len={self._max_model_len}"
                )

            current_pack.append(idx)
            current_len += packed_doc_len

        if current_pack:
            packs.append(current_pack)
        return packs

    def _build_pack_inputs(
        self,
        query_prefix_tokens: Sequence[int],
        doc_tokens: Sequence[Sequence[int]],
        doc_indices: Sequence[int],
    ) -> Tuple[Dict[str, torch.Tensor], torch.Tensor]:
        prefix_len = len(query_prefix_tokens)
        input_ids_list = list(query_prefix_tokens)
        position_ids_list = list(range(prefix_len))
        spans: List[Tuple[int, int]] = []
        current_len = prefix_len

        for idx in doc_indices:
            doc_with_suffix = list(doc_tokens[idx]) + self._suffix_tokens
            start = current_len
            end = start + len(doc_with_suffix)
            spans.append((start, end))
            input_ids_list.extend(doc_with_suffix)
            position_ids_list.extend(range(prefix_len, prefix_len + len(doc_with_suffix)))
            current_len = end

        total_len = len(input_ids_list)
        device = self._device
        neg_inf = torch.finfo(torch.float32).min

        allowed = torch.zeros((total_len, total_len), dtype=torch.bool, device=device)
        prefix_causal = torch.tril(
            torch.ones((prefix_len, prefix_len), dtype=torch.bool, device=device)
        )
        allowed[:prefix_len, :prefix_len] = prefix_causal
        for start, end in spans:
            allowed[start:end, :prefix_len] = True
            doc_len = end - start
            allowed[start:end, start:end] = torch.tril(
                torch.ones((doc_len, doc_len), dtype=torch.bool, device=device)
            )

        attention_mask = torch.full(
            (total_len, total_len),
            neg_inf,
            dtype=torch.float32,
            device=device,
        )
        attention_mask.masked_fill_(allowed, 0.0)

        inputs = {
            "input_ids": torch.tensor([input_ids_list], dtype=torch.long, device=device),
            "position_ids": torch.tensor([position_ids_list], dtype=torch.long, device=device),
            "attention_mask": attention_mask.view(1, 1, total_len, total_len),
        }
        logits_ids = torch.tensor(
            [end - 1 for _, end in spans],
            dtype=torch.long,
            device=device,
        )
        return inputs, logits_ids

    @torch.no_grad()
    def _score_pack(
        self,
        query_prefix_tokens: Sequence[int],
        doc_tokens: Sequence[Sequence[int]],
        doc_indices: Sequence[int],
    ) -> Tuple[List[float], int]:
        inputs, logits_ids = self._build_pack_inputs(
            query_prefix_tokens=query_prefix_tokens,
            doc_tokens=doc_tokens,
            doc_indices=doc_indices,
        )
        outputs = self._model(**inputs)
        scores = outputs.logits[0, logits_ids, :]
        true_vector = scores[:, self._token_true_id]
        false_vector = scores[:, self._token_false_id]
        pair_scores = torch.stack([false_vector, true_vector], dim=1)
        pair_scores = torch.nn.functional.log_softmax(pair_scores, dim=1)
        return pair_scores[:, 1].exp().tolist(), int(inputs["input_ids"].shape[1])

    def score_with_meta(
        self,
        query: str,
        docs: List[str],
        normalize: bool = True,
    ) -> Tuple[List[float], Dict[str, Any]]:
        start_ts = time.time()
        total_docs = len(docs) if docs else 0
        output_scores: List[float] = [0.0] * total_docs

        query = "" if query is None else str(query).strip()
        indexed: List[Tuple[int, str]] = []
        for i, doc in enumerate(docs or []):
            if doc is None:
                continue
            text = str(doc).strip()
            if not text:
                continue
            indexed.append((i, text))

        if not query or not indexed:
            elapsed_ms = (time.time() - start_ts) * 1000.0
            return output_scores, {
                "input_docs": total_docs,
                "usable_docs": len(indexed),
                "unique_docs": 0,
                "dedup_ratio": 0.0,
                "elapsed_ms": round(elapsed_ms, 3),
                "model": self._model_name,
                "backend": "qwen3_transformers_packed",
                "normalize": normalize,
                "packed_batches": 0,
                "max_model_len": self._max_model_len,
                "max_doc_len": self._max_doc_len,
                "sort_by_doc_length": self._sort_by_doc_length,
            }

        indexed_texts = [text for _, text in indexed]
        unique_texts, position_to_unique = _deduplicate_with_positions(indexed_texts)

        query_prefix_tokens = self._build_pair_prefix_tokens(query)
        doc_tokens = self._tokenize_documents(unique_texts, query_prefix_len=len(query_prefix_tokens))
        pack_plan = self._build_pack_plan(
            query_prefix_len=len(query_prefix_tokens),
            doc_tokens=doc_tokens,
        )

        unique_scores: List[float] = [0.0] * len(unique_texts)
        pack_lengths: List[int] = []
        with self._infer_lock:
            for pack_doc_indices in pack_plan:
                batch_scores, pack_seq_len = self._score_pack(
                    query_prefix_tokens=query_prefix_tokens,
                    doc_tokens=doc_tokens,
                    doc_indices=pack_doc_indices,
                )
                if len(batch_scores) != len(pack_doc_indices):
                    raise RuntimeError(
                        "Packed reranker score size mismatch: "
                        f"expected {len(pack_doc_indices)}, got {len(batch_scores)}"
                    )
                for idx, score in zip(pack_doc_indices, batch_scores):
                    unique_scores[idx] = float(score)
                pack_lengths.append(pack_seq_len)

        for (orig_idx, _), unique_idx in zip(indexed, position_to_unique):
            output_scores[orig_idx] = float(unique_scores[unique_idx])

        elapsed_ms = (time.time() - start_ts) * 1000.0
        dedup_ratio = 0.0
        if indexed:
            dedup_ratio = 1.0 - (len(unique_texts) / float(len(indexed)))

        meta = {
            "input_docs": total_docs,
            "usable_docs": len(indexed),
            "unique_docs": len(unique_texts),
            "dedup_ratio": round(dedup_ratio, 4),
            "elapsed_ms": round(elapsed_ms, 3),
            "model": self._model_name,
            "backend": "qwen3_transformers_packed",
            "normalize": normalize,
            "packed_batches": len(pack_plan),
            "packed_max_seq_len": max(pack_lengths) if pack_lengths else 0,
            "packed_avg_seq_len": round(sum(pack_lengths) / len(pack_lengths), 3)
            if pack_lengths
            else 0.0,
            "max_model_len": self._max_model_len,
            "max_doc_len": self._max_doc_len,
            "max_docs_per_pack": self._max_docs_per_pack,
            "sort_by_doc_length": self._sort_by_doc_length,
            "attn_implementation": self._attn_impl,
        }
        return output_scores, meta