"""

  Qwen3-Reranker via vLLM ``LLM.score()`` (pooling / cross-encoder score API).

  Matches vLLM ``examples/offline_inference/qwen3_reranker.py``: paired
  ``llm.score(query_texts, doc_texts)`` with the recommended prefix/suffix templates.
  Requires vLLM >= 0.17 (uses ``runner``/``convert`` auto, not legacy ``task="score"``).

  Dedicated venv: ``.venv-reranker-score`` + ``requirements_reranker_qwen3_vllm_score.txt``
  (see ``./scripts/setup_reranker_venv.sh qwen3_vllm_score``). Default ``model_name`` can match
  ``qwen3_vllm``; only the Python env differs for pinned high-performance vLLM.
  
  Reference: https://docs.vllm.ai/ — Qwen3 reranker example

  """
  
  from __future__ import annotations
  
  import logging
  import os
  import threading
  import time
  from typing import Any, Dict, List, Tuple
  
  logger = logging.getLogger("reranker.backends.qwen3_vllm_score")
  
  import torch
  from vllm import LLM
  
  from reranker.backends.qwen3_vllm import deduplicate_with_positions
  
  # Official vLLM Qwen3 reranker prompt layout (im_start blocks + assistant suffix).
  _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_QUERY_TEMPLATE = "{prefix}<Instruct>: {instruction}\n<Query>: {query}\n"
  _DEFAULT_DOCUMENT_TEMPLATE = "<Document>: {doc}{suffix}"

  # compact：与 qwen3_vllm._format_instruction 一致（instruction 作 system，user 内重复 Instruct）
  _IM_USER_START = "<|im_end|>\n<|im_start|>user\n"

  def _parse_env_bool(raw: str | None) -> bool | None:
      if raw is None:
          return None
      s = str(raw).strip().lower()
      if not s:
          return None
      if s in {"1", "true", "yes", "y", "on"}:
          return True
      if s in {"0", "false", "no", "n", "off"}:
          return False
      return None
  
  
  def _auto_triton_on_sm_lt_8_enabled(config: Dict[str, Any]) -> bool:
      """
      When True (default), sm < 8 injects TRITON_ATTN to avoid FA2-only paths that error on T4/V100.
  
      When False, vLLM may choose FLASHINFER on Turing; first ``score()`` can JIT-compile and needs
      ``ninja`` on PATH (``requirements_reranker_qwen3_vllm_score.txt``). Use
      ``./scripts/start_reranker.sh`` (prepends the backend venv's ``bin`` to ``PATH``) or
      ``source .../bin/activate``.
      """
      env = _parse_env_bool(os.getenv("RERANK_VLLM_AUTO_TRITON_ATTN"))
      if env is not None:
          return env
      raw = config.get("auto_triton_attn_on_sm_lt_8")
      if raw is None:
          return True
      if isinstance(raw, bool):
          return raw
      parsed = _parse_env_bool(str(raw))
      return True if parsed is None else parsed
  
  

  def _resolve_vllm_attention_config(config: Dict[str, Any]) -> Dict[str, Any] | None:
      """

      Optional explicit backend via vllm_attention_backend / RERANK_VLLM_ATTENTION_BACKEND.
  
      On compute capability < 8, vLLM may default to Flash-Attention 2, which is not supported on
      Turing/Volta; this module historically injected TRITON_ATTN. That can be slower than vLLM's
      other fallbacks — disable with auto_triton_attn_on_sm_lt_8: false or
      RERANK_VLLM_AUTO_TRITON_ATTN=0 if your stack runs without errors.

      """
      env = (os.getenv("RERANK_VLLM_ATTENTION_BACKEND") or "").strip()
      raw = config.get("vllm_attention_backend")
      if env:
          choice = env
      elif raw is not None and str(raw).strip() and str(raw).strip().lower() != "auto":
          choice = str(raw).strip()
      else:
          choice = ""
      if choice:
          backend = choice.strip().upper()
          if backend == "AUTO":
              choice = ""
          else:
              logger.info("[Qwen3_VLLM_SCORE] attention_config.backend=%s (from config/env)", backend)
              return {"backend": backend}
  
      major, minor = torch.cuda.get_device_capability()

      if major < 8 and _auto_triton_on_sm_lt_8_enabled(config):

          logger.info(
              "[Qwen3_VLLM_SCORE] GPU compute capability %d.%d < 8.0; using attention backend "
              "TRITON_ATTN (Flash-Attention 2 requires sm >= 80). "

              "To use vLLM default instead: auto_triton_attn_on_sm_lt_8: false or "
              "RERANK_VLLM_AUTO_TRITON_ATTN=0; or set vllm_attention_backend / "
              "RERANK_VLLM_ATTENTION_BACKEND.",

              major,
              minor,
          )
          return {"backend": "TRITON_ATTN"}

      if major < 8 and not _auto_triton_on_sm_lt_8_enabled(config):
          logger.info(
              "[Qwen3_VLLM_SCORE] GPU compute capability %d.%d < 8.0; auto TRITON_ATTN disabled — "
              "leaving attention backend to vLLM (no attention_config). "
              "If the first score() fails on 'ninja', install ninja in the score venv, ensure "
              "PATH includes that venv's bin (see start_reranker.sh), or use system ninja-build.",
              major,
              minor,
          )

      return None
  
  

  class Qwen3VLLMScoreRerankerBackend:
      """

      Qwen3 reranker using vLLM ``LLM.score()`` (pooling runner) for cross-encoder scores.

  
      Config from ``services.rerank.backends.qwen3_vllm_score``.
      """
  
      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")
          max_model_len = int(self._config.get("max_model_len", 2048))
          tensor_parallel_size = int(self._config.get("tensor_parallel_size", 1))
          gpu_memory_utilization = float(self._config.get("gpu_memory_utilization", 0.4))
          enable_prefix_caching = bool(self._config.get("enable_prefix_caching", False))
          enforce_eager = bool(self._config.get("enforce_eager", True))
          dtype = str(self._config.get("dtype", "float16")).strip().lower()
          use_hf_overrides = self._config.get("use_original_qwen3_hf_overrides")
          if use_hf_overrides is None:
              use_hf_overrides = True
          use_hf_overrides = bool(use_hf_overrides)
  
          self._instruction = str(
              self._config.get("instruction")
              or "Given a query, score the product for relevance"
          )

          _fmt = str(self._config.get("instruction_format") or "standard").strip().lower()
          if _fmt not in {"standard", "compact"}:
              raise ValueError(
                  f"instruction_format must be 'standard' or 'compact', got {_fmt!r}"
              )
          self._instruction_format = _fmt

          self._prefix = str(self._config.get("prompt_prefix") or _DEFAULT_PREFIX)
          self._suffix = str(self._config.get("prompt_suffix") or _DEFAULT_SUFFIX)
          self._query_template = str(self._config.get("query_template") or _DEFAULT_QUERY_TEMPLATE)
          self._document_template = str(
              self._config.get("document_template") or _DEFAULT_DOCUMENT_TEMPLATE
          )
  
          infer_batch_size = os.getenv("RERANK_VLLM_INFER_BATCH_SIZE") or self._config.get(
              "infer_batch_size", 64
          )
          sort_by_doc_length = os.getenv("RERANK_VLLM_SORT_BY_DOC_LENGTH")
          if sort_by_doc_length is None:
              sort_by_doc_length = self._config.get("sort_by_doc_length", True)
  
          self._infer_batch_size = int(infer_batch_size)
          self._sort_by_doc_length = str(sort_by_doc_length).strip().lower() in {
              "1",
              "true",
              "yes",
              "y",
              "on",
          }
  
          if not torch.cuda.is_available():
              raise RuntimeError(
                  "qwen3_vllm_score backend requires CUDA GPU, but torch.cuda.is_available() is False"
              )
          if dtype not in {"float16", "half", "auto"}:
              raise ValueError(
                  f"Unsupported dtype for qwen3_vllm_score: {dtype!r}. Use float16/half/auto."
              )
          if self._infer_batch_size <= 0:
              raise ValueError(f"infer_batch_size must be > 0, got {self._infer_batch_size}")
  
          runner = str(self._config.get("vllm_runner") or "auto").strip().lower()
          convert = str(self._config.get("vllm_convert") or "auto").strip().lower()
          if runner not in {"auto", "generate", "pooling", "draft"}:
              raise ValueError(f"Invalid vllm_runner: {runner!r}")
          if convert not in {"auto", "none", "embed", "classify"}:
              raise ValueError(f"Invalid vllm_convert: {convert!r}")
  
          logger.info(
              "[Qwen3_VLLM_SCORE] Loading model %s (LLM.score API, runner=%s, convert=%s, "

              "hf_overrides=%s, max_model_len=%s, tp=%s, gpu_mem=%.2f, dtype=%s, prefix_caching=%s, "
              "instruction_format=%s)",

              model_name,
              runner,
              convert,
              use_hf_overrides,
              max_model_len,
              tensor_parallel_size,
              gpu_memory_utilization,
              dtype,
              enable_prefix_caching,

              self._instruction_format,

          )
  
          # vLLM 0.17+ uses runner/convert instead of LLM(..., task="score"). With the official
          # Qwen3 reranker hf_overrides, architecture becomes *ForSequenceClassification -> pooling+classify.
          llm_kwargs: Dict[str, Any] = {
              "model": model_name,
              "runner": runner,
              "convert": convert,
              "tensor_parallel_size": tensor_parallel_size,
              "max_model_len": max_model_len,
              "gpu_memory_utilization": gpu_memory_utilization,
              "enable_prefix_caching": enable_prefix_caching,
              "enforce_eager": enforce_eager,
              "dtype": dtype,
          }
          hf_overrides: Dict[str, Any] = dict(self._config.get("hf_overrides") or {})
          if use_hf_overrides:
              hf_overrides = {
                  **hf_overrides,
                  "architectures": ["Qwen3ForSequenceClassification"],
                  "classifier_from_token": ["no", "yes"],
                  "is_original_qwen3_reranker": True,
              }
          if hf_overrides:
              llm_kwargs["hf_overrides"] = hf_overrides
  

          attn_cfg = _resolve_vllm_attention_config(self._config)
          if attn_cfg is not None:
              llm_kwargs["attention_config"] = attn_cfg
  

          self._llm = LLM(**llm_kwargs)
          # vLLM score path: single-process safety (mirrors generate backend until verified).
          self._infer_lock = threading.Lock()
  
          self._model_name = model_name
          logger.info("[Qwen3_VLLM_SCORE] Model ready | model=%s", model_name)
  
      def _format_pair(self, query: str, doc: str) -> Tuple[str, str]:

          if self._instruction_format == "compact":
              # Align with reranker.backends.qwen3_vllm._format_instruction query/doc split for LLM.score().
              compact_prefix = f"<|im_start|>system\n{self._instruction}{_IM_USER_START}"
              q_text = (
                  f"{compact_prefix}<Instruct>: {self._instruction}\n\n<Query>: {query}\n"
              )
              d_text = f"\n<Document>: {doc}{self._suffix}"
              return q_text, d_text

          q_text = self._query_template.format(
              prefix=self._prefix,
              instruction=self._instruction,
              query=query,
          )
          d_text = self._document_template.format(doc=doc, suffix=self._suffix)
          return q_text, d_text
  
      def _score_batch(self, pairs: List[Tuple[str, str]]) -> List[float]:
          if not pairs:
              return []
          queries: List[str] = []
          documents: List[str] = []
          for q, d in pairs:
              qt, dt = self._format_pair(q, d)
              queries.append(qt)
              documents.append(dt)
          with self._infer_lock:
              outputs = self._llm.score(queries, documents, use_tqdm=False)
          scores: List[float] = []
          for out in outputs:
              so = out.outputs
              scores.append(float(so.score))
          return scores
  
      @staticmethod
      def _estimate_doc_lengths(docs: List[str]) -> List[int]:
          if not docs:
              return []
          return [len(text) for text in docs]
  
      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_vllm_score",
                  "normalize": normalize,
                  "infer_batch_size": self._infer_batch_size,
                  "inference_batches": 0,
                  "sort_by_doc_length": self._sort_by_doc_length,

                  "instruction_format": self._instruction_format,

              }
  
          indexed_texts = [text for _, text in indexed]
          unique_texts, position_to_unique = deduplicate_with_positions(indexed_texts)
  
          lengths = self._estimate_doc_lengths(unique_texts)
          order = list(range(len(unique_texts)))
          if self._sort_by_doc_length and len(unique_texts) > 1:
              order = sorted(order, key=lambda i: lengths[i])
  
          unique_scores: List[float] = [0.0] * len(unique_texts)
          inference_batches = 0
          for start in range(0, len(order), self._infer_batch_size):
              batch_indices = order[start : start + self._infer_batch_size]
              inference_batches += 1
              pairs = [(query, unique_texts[i]) for i in batch_indices]
              batch_scores = self._score_batch(pairs)
              if len(batch_scores) != len(batch_indices):
                  raise RuntimeError(
                      f"Reranker score size mismatch: expected {len(batch_indices)}, got {len(batch_scores)}"
                  )
              for idx, score in zip(batch_indices, batch_scores):
                  unique_scores[idx] = float(score)
  
          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_vllm_score",
              "normalize": normalize,
              "infer_batch_size": self._infer_batch_size,
              "inference_batches": inference_batches,
              "sort_by_doc_length": self._sort_by_doc_length,

              "instruction_format": self._instruction_format,

          }
          return output_scores, meta