"""
Qwen3-Reranker-0.6B backend using vLLM.

Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B
Requires: vllm>=0.8.5, transformers; GPU recommended.
"""

from __future__ import annotations

import logging
import math
import threading
import time
from typing import Any, Dict, List, Optional, Tuple

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

try:
    import torch
    from transformers import AutoTokenizer
    from vllm import LLM, SamplingParams
    from vllm.inputs.data import TokensPrompt
except ImportError as e:
    raise ImportError(
        "Qwen3-vLLM reranker backend requires vllm>=0.8.5 and transformers. "
        "Install with: pip install vllm transformers"
    ) from e


def _format_instruction(instruction: str, query: str, doc: str) -> List[Dict[str, str]]:
    """Build chat messages for one (query, doc) pair."""
    return [
        {
            "role": "system",
            "content": "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\".",
        },
        {
            "role": "user",
            "content": f"<Instruct>: {instruction}\n\n<Query>: {query}\n\n<Document>: {doc}",
        },
    ]


class Qwen3VLLMRerankerBackend:
    """
    Qwen3-Reranker-0.6B with vLLM inference.
    Config from services.rerank.backends.qwen3_vllm.
    """

    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()
        self._instruction = str(
            self._config.get("instruction")
            or "Given a web search query, retrieve relevant passages that answer the query"
        )
        if not torch.cuda.is_available():
            raise RuntimeError("qwen3_vllm 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: {dtype!r}. Use float16/half/auto.")

        logger.info(
            "[Qwen3_VLLM] Loading model %s (max_model_len=%s, tp=%s, gpu_mem=%.2f, dtype=%s, prefix_caching=%s)",
            model_name,
            max_model_len,
            tensor_parallel_size,
            gpu_memory_utilization,
            dtype,
            enable_prefix_caching,
        )

        self._llm = LLM(
            model=model_name,
            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,
        )
        self._tokenizer = AutoTokenizer.from_pretrained(model_name)
        self._tokenizer.padding_side = "left"
        self._tokenizer.pad_token = self._tokenizer.eos_token

        # Suffix for generation prompt (assistant answer)
        self._suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
        self._suffix_tokens = self._tokenizer.encode(
            self._suffix, add_special_tokens=False
        )
        self._max_prompt_len = max_model_len - len(self._suffix_tokens)

        self._true_token = self._tokenizer("yes", add_special_tokens=False).input_ids[0]
        self._false_token = self._tokenizer("no", add_special_tokens=False).input_ids[0]
        self._sampling_params = SamplingParams(
            temperature=0,
            max_tokens=1,
            logprobs=20,
            allowed_token_ids=[self._true_token, self._false_token],
        )
        # vLLM generate path is unstable under concurrent calls in this process model.
        # Serialize infer calls to avoid engine-core protocol corruption.
        self._infer_lock = threading.Lock()

        self._model_name = model_name
        logger.info("[Qwen3_VLLM] Model ready | model=%s", model_name)

    def _process_inputs(
        self,
        pairs: List[Tuple[str, str]],
    ) -> List[TokensPrompt]:
        """Build tokenized prompts for vLLM from (query, doc) pairs. Batch apply_chat_template."""
        messages_batch = [
            _format_instruction(self._instruction, q, d) for q, d in pairs
        ]
        tokenized = self._tokenizer.apply_chat_template(
            messages_batch,
            tokenize=True,
            add_generation_prompt=False,
            enable_thinking=False,
        )
        # Single conv returns flat list; batch returns list of lists
        if tokenized and not isinstance(tokenized[0], list):
            tokenized = [tokenized]
        prompts = [
            TokensPrompt(
                prompt_token_ids=ids[: self._max_prompt_len] + self._suffix_tokens
            )
            for ids in tokenized
        ]
        return prompts

    def _compute_scores(
        self,
        prompts: List[TokensPrompt],
    ) -> List[float]:
        """Run vLLM generate and compute yes/no probability per prompt."""
        if not prompts:
            return []
        outputs = self._llm.generate(prompts, self._sampling_params, use_tqdm=False)
        scores = []
        for i in range(len(outputs)):
            out = outputs[i]
            if not out.outputs:
                scores.append(0.0)
                continue
            final_logits = out.outputs[0].logprobs
            if not final_logits:
                scores.append(0.0)
                continue
            last = final_logits[-1]
            # Match official: missing token -> logprob = -10
            if self._true_token not in last:
                true_logit = -10
            else:
                true_logit = last[self._true_token].logprob
            if self._false_token not in last:
                false_logit = -10
            else:
                false_logit = last[self._false_token].logprob
            true_score = math.exp(true_logit)
            false_score = math.exp(false_logit)
            score = true_score / (true_score + false_score)
            scores.append(float(score))
        return scores

    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",
                "normalize": normalize,
            }

        # Deduplicate by text, keep mapping to original indices
        unique_texts: List[str] = []
        position_to_unique: List[int] = []
        prev: Optional[str] = None
        for _idx, text in indexed:
            if text != prev:
                unique_texts.append(text)
                prev = text
            position_to_unique.append(len(unique_texts) - 1)

        pairs = [(query, t) for t in unique_texts]
        with self._infer_lock:
            prompts = self._process_inputs(pairs)
            unique_scores = self._compute_scores(prompts)

        for (orig_idx, _), unique_idx in zip(indexed, position_to_unique):
            # Score is already P(yes) in [0,1] from yes/(yes+no)
            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",
            "normalize": normalize,
        }
        return output_scores, meta