"""
Qwen3-Reranker GGUF backend using llama-cpp-python.

Reference:
- https://huggingface.co/DevQuasar/Qwen.Qwen3-Reranker-4B-GGUF
- https://huggingface.co/Qwen/Qwen3-Reranker-4B
- https://huggingface.co/ggml-org/Qwen3-Reranker-0.6B-Q8_0-GGUF
- https://huggingface.co/Qwen/Qwen3-Reranker-0.6B
"""

from __future__ import annotations

import logging
import math
import os
import threading
import time
from pathlib import Path
from typing import Any, Dict, List, Tuple


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


_BACKEND_DEFAULTS: Dict[str, Dict[str, str]] = {
    "qwen3_gguf": {
        "repo_id": "DevQuasar/Qwen.Qwen3-Reranker-4B-GGUF",
        "filename": "*Q8_0.gguf",
        "local_dir": "./models/reranker/qwen3-reranker-4b-gguf",
    },
    "qwen3_gguf_06b": {
        "repo_id": "ggml-org/Qwen3-Reranker-0.6B-Q8_0-GGUF",
        "filename": "qwen3-reranker-0.6b-q8_0.gguf",
        "local_dir": "./models/reranker/qwen3-reranker-0.6b-q8_0-gguf",
    },
}


def deduplicate_with_positions(texts: List[str]) -> Tuple[List[str], List[int]]:
    """Deduplicate texts globally while preserving first-seen order."""
    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


def _format_instruction(instruction: str, query: str, doc: str) -> str:
    return "<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {doc}".format(
        instruction=instruction,
        query=query,
        doc=doc,
    )


class Qwen3GGUFRerankerBackend:
    """
    Qwen3-Reranker GGUF backend using llama.cpp through llama-cpp-python.

    Tuned for short-query / short-doc reranking on a single GPU.
    Config from services.rerank.backends.<backend_name>.
    """

    def __init__(self, config: Dict[str, Any]) -> None:
        self._config = config or {}
        self._backend_name = str(self._config.get("_backend_name") or "qwen3_gguf").strip()
        defaults = _BACKEND_DEFAULTS.get(self._backend_name, _BACKEND_DEFAULTS["qwen3_gguf"])
        self._repo_id = str(self._config.get("repo_id") or defaults["repo_id"]).strip()
        self._filename = str(self._config.get("filename") or defaults["filename"]).strip()
        self._model_path = str(self._config.get("model_path") or "").strip()
        self._cache_dir = str(self._config.get("cache_dir") or "").strip() or None
        self._local_dir = str(self._config.get("local_dir") or defaults["local_dir"]).strip() or None
        self._instruction = str(
            self._config.get("instruction")
            or "Rank products by query with category & style match prioritized"
        )
        self._infer_batch_size = int(
            os.getenv("RERANK_GGUF_INFER_BATCH_SIZE") or self._config.get("infer_batch_size", 8)
        )
        sort_by_doc_length = os.getenv("RERANK_GGUF_SORT_BY_DOC_LENGTH")
        if sort_by_doc_length is None:
            sort_by_doc_length = self._config.get("sort_by_doc_length", True)
        self._sort_by_doc_length = str(sort_by_doc_length).strip().lower() in {
            "1",
            "true",
            "yes",
            "y",
            "on",
        }
        self._length_sort_mode = str(self._config.get("length_sort_mode") or "char").strip().lower()
        self._reuse_query_state = bool(self._config.get("reuse_query_state", False))

        n_ctx = int(self._config.get("n_ctx", self._config.get("max_model_len", 384)))
        n_batch = int(self._config.get("n_batch", min(n_ctx, 384)))
        n_ubatch = int(self._config.get("n_ubatch", min(n_batch, 128)))
        n_gpu_layers = int(self._config.get("n_gpu_layers", 24))
        main_gpu = int(self._config.get("main_gpu", 0))
        n_threads = int(self._config.get("n_threads", 2))
        n_threads_batch = int(self._config.get("n_threads_batch", 4))
        flash_attn = bool(self._config.get("flash_attn", True))
        offload_kqv = bool(self._config.get("offload_kqv", True))
        use_mmap = bool(self._config.get("use_mmap", True))
        use_mlock = bool(self._config.get("use_mlock", False))
        verbose = bool(self._config.get("verbose", False))
        enable_warmup = bool(self._config.get("enable_warmup", True))

        if self._infer_batch_size <= 0:
            raise ValueError(f"infer_batch_size must be > 0, got {self._infer_batch_size}")
        if n_ctx <= 0:
            raise ValueError(f"n_ctx must be > 0, got {n_ctx}")
        if n_batch <= 0 or n_ubatch <= 0:
            raise ValueError(f"n_batch/n_ubatch must be > 0, got {n_batch}/{n_ubatch}")

        try:
            from llama_cpp import Llama
        except Exception as exc:  # pragma: no cover - depends on optional dependency
            raise RuntimeError(
                f"{self._backend_name} backend requires llama-cpp-python. "
                f"Install the {self._backend_name} backend venv first via "
                f"scripts/setup_reranker_venv.sh {self._backend_name}."
            ) from exc

        self._llama_class = Llama
        self._n_ctx = n_ctx
        self._n_batch = n_batch
        self._n_ubatch = n_ubatch
        self._n_gpu_layers = n_gpu_layers
        self._enable_warmup = enable_warmup
        self._infer_lock = threading.Lock()

        logger.info(
            "[Qwen3_GGUF] Loading backend=%s repo=%s filename=%s model_path=%s n_ctx=%s n_batch=%s n_ubatch=%s n_gpu_layers=%s flash_attn=%s offload_kqv=%s reuse_query_state=%s",
            self._backend_name,
            self._repo_id,
            self._filename,
            self._model_path or None,
            n_ctx,
            n_batch,
            n_ubatch,
            n_gpu_layers,
            flash_attn,
            offload_kqv,
            self._reuse_query_state,
        )

        llm_kwargs = {
            "n_ctx": n_ctx,
            "n_batch": n_batch,
            "n_ubatch": n_ubatch,
            "n_gpu_layers": n_gpu_layers,
            "main_gpu": main_gpu,
            "n_threads": n_threads,
            "n_threads_batch": n_threads_batch,
            "logits_all": True,
            "offload_kqv": offload_kqv,
            "flash_attn": flash_attn,
            "use_mmap": use_mmap,
            "use_mlock": use_mlock,
            "verbose": verbose,
        }
        llm_kwargs = {key: value for key, value in llm_kwargs.items() if value is not None}
        self._llm = self._load_model(llm_kwargs)
        self._model_name = self._model_path or f"{self._repo_id}:{self._filename}"

        self._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"
        )
        self._suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
        self._prefix_tokens = self._tokenize(self._prefix, special=True)
        self._suffix_tokens = self._tokenize(self._suffix, special=True)
        self._request_prefix_template = "<Instruct>: {instruction}\n<Query>: {query}\n<Document>: "
        self._effective_max_len = self._n_ctx - len(self._prefix_tokens) - len(self._suffix_tokens)
        if self._effective_max_len <= 16:
            raise RuntimeError(
                f"n_ctx={self._n_ctx} is too small after prompt overhead; effective={self._effective_max_len}"
            )

        self._true_token = self._single_token_id("yes")
        self._false_token = self._single_token_id("no")

        if self._enable_warmup:
            self._warmup()

        logger.info(
            "[Qwen3_GGUF] Model ready | backend=%s model=%s effective_max_len=%s infer_batch_size=%s sort_by_doc_length=%s",
            self._backend_name,
            self._model_name,
            self._effective_max_len,
            self._infer_batch_size,
            self._sort_by_doc_length,
        )

    def _load_model(self, llm_kwargs: Dict[str, Any]):
        if self._model_path:
            return self._llama_class(model_path=self._model_path, **llm_kwargs)
        if self._local_dir:
            matches = sorted(
                path for path in Path(self._local_dir).glob(self._filename) if path.is_file()
            )
            if matches:
                local_model_path = str(matches[0].resolve())
                logger.info("[Qwen3_GGUF] Using local GGUF file: %s", local_model_path)
                return self._llama_class(model_path=local_model_path, **llm_kwargs)
        return self._llama_class.from_pretrained(
            repo_id=self._repo_id,
            filename=self._filename,
            local_dir=self._local_dir,
            cache_dir=self._cache_dir,
            **llm_kwargs,
        )

    def _tokenize(self, text: str, *, special: bool) -> List[int]:
        return list(
            self._llm.tokenize(
                text.encode("utf-8"),
                add_bos=False,
                special=special,
            )
        )

    def _single_token_id(self, text: str) -> int:
        token_ids = self._tokenize(text, special=False)
        if len(token_ids) != 1:
            raise RuntimeError(f"Expected {text!r} to be one token, got {token_ids}")
        return int(token_ids[0])

    def _warmup(self) -> None:
        try:
            prompt = self._build_prompt_tokens("warmup query", "warmup document")
            with self._infer_lock:
                self._eval_logits(prompt)
        except Exception as exc:  # pragma: no cover - defensive
            logger.warning("[Qwen3_GGUF] Warmup failed: %s", exc)

    def _build_request_prefix_tokens(self, query: str) -> List[int]:
        request_prefix = self._request_prefix_template.format(
            instruction=self._instruction,
            query=query,
        )
        return self._tokenize(request_prefix, special=False)

    def _build_prompt_tokens(self, query: str, doc: str) -> List[int]:
        pair = _format_instruction(self._instruction, query, doc)
        pair_tokens = self._tokenize(pair, special=False)
        pair_tokens = pair_tokens[: self._effective_max_len]
        return self._prefix_tokens + pair_tokens + self._suffix_tokens

    def _eval_logits(self, prompt_tokens: List[int]) -> List[float]:
        self._llm.reset()
        self._llm.eval(prompt_tokens)
        logits = self._llm.eval_logits
        if not logits:
            raise RuntimeError("llama.cpp returned empty logits")
        return list(logits[-1])

    def _score_prompt(self, prompt_tokens: List[int]) -> float:
        logits = self._eval_logits(prompt_tokens)
        true_logit = float(logits[self._true_token])
        false_logit = float(logits[self._false_token])
        max_logit = max(true_logit, false_logit)
        true_exp = math.exp(true_logit - max_logit)
        false_exp = math.exp(false_logit - max_logit)
        return float(true_exp / (true_exp + false_exp))

    def _supports_query_state_reuse(self) -> bool:
        return (
            self._reuse_query_state
            and hasattr(self._llm, "save_state")
            and hasattr(self._llm, "load_state")
        )

    def _build_query_state_locked(self, query: str):
        request_prefix_tokens = self._build_request_prefix_tokens(query)
        max_doc_tokens = self._effective_max_len - len(request_prefix_tokens)
        if max_doc_tokens <= 0:
            return None, 0
        self._llm.reset()
        self._llm.eval(self._prefix_tokens + request_prefix_tokens)
        return self._llm.save_state(), max_doc_tokens

    def _score_doc_with_state_locked(self, state, doc_tokens: List[int], max_doc_tokens: int) -> float:
        self._llm.load_state(state)
        self._llm.eval(doc_tokens[:max_doc_tokens] + self._suffix_tokens)
        logits = self._llm.eval_logits
        if not logits:
            raise RuntimeError("llama.cpp returned empty logits")
        final_logits = list(logits[-1])
        true_logit = float(final_logits[self._true_token])
        false_logit = float(final_logits[self._false_token])
        max_logit = max(true_logit, false_logit)
        true_exp = math.exp(true_logit - max_logit)
        false_exp = math.exp(false_logit - max_logit)
        return float(true_exp / (true_exp + false_exp))

    def _estimate_doc_lengths(self, docs: List[str]) -> List[int]:
        if self._length_sort_mode == "token":
            return [len(self._tokenize(text, special=False)) for text in docs]
        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": self._backend_name,
                "normalize": normalize,
                "infer_batch_size": self._infer_batch_size,
                "inference_batches": 0,
                "sort_by_doc_length": self._sort_by_doc_length,
                "n_ctx": self._n_ctx,
                "n_batch": self._n_batch,
                "n_ubatch": self._n_ubatch,
                "n_gpu_layers": self._n_gpu_layers,
            }

        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)
        unique_doc_tokens = [self._tokenize(text, special=False) for text in unique_texts]
        inference_batches = 0
        with self._infer_lock:
            query_state = None
            max_doc_tokens = self._effective_max_len
            if self._supports_query_state_reuse():
                query_state, max_doc_tokens = self._build_query_state_locked(query)
            for start in range(0, len(order), self._infer_batch_size):
                batch_indices = order[start : start + self._infer_batch_size]
                inference_batches += 1
                for idx in batch_indices:
                    if query_state is not None:
                        unique_scores[idx] = self._score_doc_with_state_locked(
                            query_state,
                            unique_doc_tokens[idx],
                            max_doc_tokens,
                        )
                    else:
                        prompt = self._build_prompt_tokens(query, unique_texts[idx])
                        unique_scores[idx] = self._score_prompt(prompt)

        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": self._backend_name,
            "normalize": normalize,
            "infer_batch_size": self._infer_batch_size,
            "inference_batches": inference_batches,
            "sort_by_doc_length": self._sort_by_doc_length,
            "length_sort_mode": self._length_sort_mode,
            "n_ctx": self._n_ctx,
            "n_batch": self._n_batch,
            "n_ubatch": self._n_ubatch,
            "n_gpu_layers": self._n_gpu_layers,
            "reuse_query_state": query_state is not None,
        }
        return output_scores, meta