"""Tests for GGUF Cross Encoder (Qwen3CrossEncoderGGUF)."""

from __future__ import annotations

import sys
from pathlib import Path
from typing import Any
from unittest import mock
from unittest.mock import MagicMock, patch

import numpy as np
import pytest

from qwen3_embed.rerank.cross_encoder.gguf_cross_encoder import (
    DEFAULT_INSTRUCTION,
    RERANK_TEMPLATE,
    SYSTEM_PROMPT,
    TOKEN_NO_ID,
    TOKEN_YES_ID,
    Qwen3CrossEncoderGGUF,
    _check_llama_cpp,
)

# ---------------------------------------------------------------------------
# Tests for _check_llama_cpp
# ---------------------------------------------------------------------------


def test_check_llama_cpp_missing():
    """Test that ImportError is raised when llama-cpp-python is missing."""
    with (
        mock.patch.dict(sys.modules, {"llama_cpp": None}),
        pytest.raises(ImportError, match="llama-cpp-python is required"),
    ):
        _check_llama_cpp()


def test_check_llama_cpp_present():
    """Test that no error is raised when llama-cpp-python is present."""
    # Mocking a successful import
    mock_module = mock.Mock()
    with mock.patch.dict(sys.modules, {"llama_cpp": mock_module}):
        try:
            _check_llama_cpp()
        except ImportError:
            pytest.fail("ImportError raised unexpectedly when llama_cpp is mocked as present")


def test_gguf_cross_encoder_init_missing_dependency():
    """Test that Qwen3CrossEncoderGGUF init fails if dependency is missing."""
    with (
        mock.patch.dict(sys.modules, {"llama_cpp": None}),
        pytest.raises(ImportError, match="llama-cpp-python is required"),
    ):
        # We don't need arguments because it should fail before using them
        Qwen3CrossEncoderGGUF()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_mock_llm(vocab_size: int = 10000) -> MagicMock:
    """Return a mock Llama instance for scoring tests."""
    mock_llm = MagicMock()
    # scores[n] is an array of vocab_size logits
    logits = np.zeros(vocab_size, dtype=np.float32)
    # Default: "yes" wins strongly
    logits[TOKEN_YES_ID] = 5.0
    logits[TOKEN_NO_ID] = 1.0
    mock_llm.scores.__getitem__ = MagicMock(return_value=logits)
    mock_llm.tokenize.return_value = [1, 2, 3]
    return mock_llm


def _make_model(**extra_attrs: Any) -> Qwen3CrossEncoderGGUF:
    """Create a Qwen3CrossEncoderGGUF instance with __init__ bypassed."""
    with patch(
        "qwen3_embed.rerank.cross_encoder.gguf_cross_encoder.Qwen3CrossEncoderGGUF.__init__",
        return_value=None,
    ):
        model = Qwen3CrossEncoderGGUF()
    model._llm = _make_mock_llm()
    for k, v in extra_attrs.items():
        setattr(model, k, v)
    return model


# ---------------------------------------------------------------------------
# Tests for __init__
# ---------------------------------------------------------------------------


class TestGGUFCrossEncoderInit:
    def test_init_creates_llm_cpu(self, tmp_path: Path):
        """Test __init__ creates Llama with n_gpu_layers=0 for CPU device."""
        from qwen3_embed.common.types import Device

        model_file = tmp_path / "qwen3-reranker-0.6b-q4-k-m.gguf"
        model_file.touch()

        mock_llama_cls = MagicMock()
        mock_llama_module = MagicMock()
        mock_llama_module.Llama = mock_llama_cls

        with (
            patch.dict(sys.modules, {"llama_cpp": mock_llama_module}),
            patch.object(Qwen3CrossEncoderGGUF, "download_model", return_value=str(tmp_path)),
            patch(
                "qwen3_embed.rerank.cross_encoder.gguf_cross_encoder.define_cache_dir",
                return_value=Path("/tmp"),
            ),
        ):
            model = Qwen3CrossEncoderGGUF(
                model_name="n24q02m/Qwen3-Reranker-0.6B-GGUF",
                cuda=Device.CPU,
            )

        mock_llama_cls.assert_called_once()
        call_kwargs = mock_llama_cls.call_args[1]
        assert call_kwargs["n_gpu_layers"] == 0
        assert call_kwargs["logits_all"] is False
        assert call_kwargs["verbose"] is False
        assert model._llm is mock_llama_cls.return_value

    def test_init_creates_llm_gpu_auto(self, tmp_path: Path):
        """Test __init__ creates Llama with n_gpu_layers=-1 for Device.AUTO."""
        from qwen3_embed.common.types import Device

        model_file = tmp_path / "qwen3-reranker-0.6b-q4-k-m.gguf"
        model_file.touch()

        mock_llama_cls = MagicMock()
        mock_llama_module = MagicMock()
        mock_llama_module.Llama = mock_llama_cls

        with (
            patch.dict(sys.modules, {"llama_cpp": mock_llama_module}),
            patch.object(Qwen3CrossEncoderGGUF, "download_model", return_value=str(tmp_path)),
            patch(
                "qwen3_embed.rerank.cross_encoder.gguf_cross_encoder.define_cache_dir",
                return_value=Path("/tmp"),
            ),
        ):
            Qwen3CrossEncoderGGUF(
                model_name="n24q02m/Qwen3-Reranker-0.6B-GGUF",
                cuda=Device.AUTO,
            )

        call_kwargs = mock_llama_cls.call_args[1]
        assert call_kwargs["n_gpu_layers"] == -1

    def test_init_cuda_false_forces_cpu(self, tmp_path: Path):
        """Test __init__ uses n_gpu_layers=0 when cuda=False."""
        model_file = tmp_path / "qwen3-reranker-0.6b-q4-k-m.gguf"
        model_file.touch()

        mock_llama_cls = MagicMock()
        mock_llama_module = MagicMock()
        mock_llama_module.Llama = mock_llama_cls

        with (
            patch.dict(sys.modules, {"llama_cpp": mock_llama_module}),
            patch.object(Qwen3CrossEncoderGGUF, "download_model", return_value=str(tmp_path)),
            patch(
                "qwen3_embed.rerank.cross_encoder.gguf_cross_encoder.define_cache_dir",
                return_value=Path("/tmp"),
            ),
        ):
            Qwen3CrossEncoderGGUF(
                model_name="n24q02m/Qwen3-Reranker-0.6B-GGUF",
                cuda=False,
            )

        call_kwargs = mock_llama_cls.call_args[1]
        assert call_kwargs["n_gpu_layers"] == 0

    def test_init_raises_file_not_found(self, tmp_path: Path):
        """Test __init__ raises FileNotFoundError if GGUF file is absent."""
        mock_llama_module = MagicMock()

        with (
            patch.dict(sys.modules, {"llama_cpp": mock_llama_module}),
            patch.object(Qwen3CrossEncoderGGUF, "download_model", return_value=str(tmp_path)),
            patch(
                "qwen3_embed.rerank.cross_encoder.gguf_cross_encoder.define_cache_dir",
                return_value=Path("/tmp"),
            ),
            pytest.raises(FileNotFoundError, match="GGUF model file not found"),
        ):
            Qwen3CrossEncoderGGUF(model_name="n24q02m/Qwen3-Reranker-0.6B-GGUF")

    def test_init_threads_none_uses_zero(self, tmp_path: Path):
        """Test __init__ passes n_threads=0 when threads=None."""
        model_file = tmp_path / "qwen3-reranker-0.6b-q4-k-m.gguf"
        model_file.touch()

        mock_llama_cls = MagicMock()
        mock_llama_module = MagicMock()
        mock_llama_module.Llama = mock_llama_cls

        with (
            patch.dict(sys.modules, {"llama_cpp": mock_llama_module}),
            patch.object(Qwen3CrossEncoderGGUF, "download_model", return_value=str(tmp_path)),
            patch(
                "qwen3_embed.rerank.cross_encoder.gguf_cross_encoder.define_cache_dir",
                return_value=Path("/tmp"),
            ),
        ):
            Qwen3CrossEncoderGGUF(
                model_name="n24q02m/Qwen3-Reranker-0.6B-GGUF",
                threads=None,
            )

        call_kwargs = mock_llama_cls.call_args[1]
        assert call_kwargs["n_threads"] == 0

    def test_init_threads_custom(self, tmp_path: Path):
        """Test __init__ passes custom threads count."""
        model_file = tmp_path / "qwen3-reranker-0.6b-q4-k-m.gguf"
        model_file.touch()

        mock_llama_cls = MagicMock()
        mock_llama_module = MagicMock()
        mock_llama_module.Llama = mock_llama_cls

        with (
            patch.dict(sys.modules, {"llama_cpp": mock_llama_module}),
            patch.object(Qwen3CrossEncoderGGUF, "download_model", return_value=str(tmp_path)),
            patch(
                "qwen3_embed.rerank.cross_encoder.gguf_cross_encoder.define_cache_dir",
                return_value=Path("/tmp"),
            ),
        ):
            Qwen3CrossEncoderGGUF(
                model_name="n24q02m/Qwen3-Reranker-0.6B-GGUF",
                threads=8,
            )

        call_kwargs = mock_llama_cls.call_args[1]
        assert call_kwargs["n_threads"] == 8


# ---------------------------------------------------------------------------
# Tests for _format_rerank_input (line 150)
# ---------------------------------------------------------------------------


class TestFormatRerankInput:
    def test_format_default_instruction(self):
        """Test _format_rerank_input uses default instruction when none provided."""
        result = Qwen3CrossEncoderGGUF._format_rerank_input("what is AI?", "AI is...")

        assert SYSTEM_PROMPT in result
        assert DEFAULT_INSTRUCTION in result
        assert "what is AI?" in result
        assert "AI is..." in result

    def test_format_custom_instruction(self):
        """Test _format_rerank_input uses custom instruction."""
        custom_instruction = "Judge relevance strictly."
        result = Qwen3CrossEncoderGGUF._format_rerank_input(
            "query", "document", instruction=custom_instruction
        )

        assert custom_instruction in result
        assert DEFAULT_INSTRUCTION not in result

    def test_format_template_structure(self):
        """Test _format_rerank_input produces correct chat-template structure."""
        result = Qwen3CrossEncoderGGUF._format_rerank_input("my query", "my document")

        expected = RERANK_TEMPLATE.format(
            system=SYSTEM_PROMPT,
            instruction=DEFAULT_INSTRUCTION,
            query="my query",
            document="my document",
        )
        assert result == expected

    def test_format_returns_string(self):
        """Test _format_rerank_input returns a string."""
        result = Qwen3CrossEncoderGGUF._format_rerank_input("q", "d")
        assert isinstance(result, str)


# ---------------------------------------------------------------------------
# Tests for _score_text (lines 162-181)
# ---------------------------------------------------------------------------


class TestScoreText:
    def test_score_text_yes_wins(self):
        """Test _score_text returns P(yes)>0.5 when yes logit > no logit."""
        model = _make_model()
        # yes_logit=5.0, no_logit=1.0 → P(yes) >> 0.5
        score = model._score_text("some text")

        assert isinstance(score, float)
        assert score > 0.5

    def test_score_text_no_wins(self):
        """Test _score_text returns P(yes)<0.5 when no logit > yes logit."""
        model = _make_model()
        logits = np.zeros(10000, dtype=np.float32)
        logits[TOKEN_YES_ID] = 1.0
        logits[TOKEN_NO_ID] = 5.0
        model._llm.scores.__getitem__ = MagicMock(return_value=logits)

        score = model._score_text("some text")

        assert score < 0.5

    def test_score_text_equal_logits_returns_half(self):
        """Test _score_text returns ~0.5 when yes and no logits are equal."""
        model = _make_model()
        logits = np.zeros(10000, dtype=np.float32)
        logits[TOKEN_YES_ID] = 2.0
        logits[TOKEN_NO_ID] = 2.0
        model._llm.scores.__getitem__ = MagicMock(return_value=logits)

        score = model._score_text("some text")

        assert abs(score - 0.5) < 1e-5

    def test_score_text_calls_tokenize_and_eval(self):
        """Test _score_text calls tokenize, reset, and eval on the LLM."""
        model = _make_model()
        model._score_text("evaluate this")

        model._llm.tokenize.assert_called_once_with(b"evaluate this", add_bos=False)
        model._llm.reset.assert_called_once()
        model._llm.eval.assert_called_once()

    def test_score_text_uses_last_token_logits(self):
        """Test _score_text indexes scores with len(tokens)-1 (last token)."""
        model = _make_model()
        tokens = [1, 2, 3, 4, 5]
        model._llm.tokenize.return_value = tokens

        model._score_text("some input")

        # Should use index len(tokens)-1 = 4
        model._llm.scores.__getitem__.assert_called_with(len(tokens) - 1)

    def test_score_text_returns_float_in_range(self):
        """Test _score_text output is in [0, 1]."""
        model = _make_model()
        score = model._score_text("test")

        assert 0.0 <= score <= 1.0


# ---------------------------------------------------------------------------
# Tests for rerank (lines 204-207)
# ---------------------------------------------------------------------------


class TestRerank:
    def test_rerank_yields_scores_per_document(self):
        """Test rerank yields one score per document."""
        model = _make_model()
        documents = ["doc1", "doc2", "doc3"]
        scores = list(model.rerank("my query", documents))

        assert len(scores) == 3
        assert all(isinstance(s, float) for s in scores)

    def test_rerank_uses_default_instruction(self):
        """Test rerank uses DEFAULT_INSTRUCTION when no instruction kwarg provided."""
        model = _make_model()
        model._score_text = MagicMock(return_value=0.9)  # type: ignore[method-assign]

        list(model.rerank("q", ["doc"]))

        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
        assert DEFAULT_INSTRUCTION in call_arg

    def test_rerank_custom_instruction(self):
        """Test rerank passes custom instruction to _format_rerank_input."""
        model = _make_model()
        model._score_text = MagicMock(return_value=0.9)  # type: ignore[method-assign]
        custom_instruction = "Custom scoring instruction."

        list(model.rerank("q", ["doc"], instruction=custom_instruction))

        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
        assert custom_instruction in call_arg

    def test_rerank_score_range(self):
        """Test rerank scores are in [0, 1]."""
        model = _make_model()
        scores = list(model.rerank("query", ["doc1", "doc2"]))

        for score in scores:
            assert 0.0 <= score <= 1.0

    def test_rerank_empty_documents(self):
        """Test rerank with empty documents yields nothing."""
        model = _make_model()
        scores = list(model.rerank("query", []))

        assert scores == []


# ---------------------------------------------------------------------------
# Tests for rerank_pairs (lines 227-230)
# ---------------------------------------------------------------------------


class TestRerankPairs:
    def test_rerank_pairs_yields_scores(self):
        """Test rerank_pairs yields one score per pair."""
        model = _make_model()
        pairs = [("query1", "doc1"), ("query2", "doc2")]
        scores = list(model.rerank_pairs(pairs))

        assert len(scores) == 2
        assert all(isinstance(s, float) for s in scores)

    def test_rerank_pairs_uses_default_instruction(self):
        """Test rerank_pairs uses DEFAULT_INSTRUCTION by default."""
        model = _make_model()
        model._score_text = MagicMock(return_value=0.7)  # type: ignore[method-assign]

        list(model.rerank_pairs([("q", "d")]))

        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
        assert DEFAULT_INSTRUCTION in call_arg

    def test_rerank_pairs_custom_instruction(self):
        """Test rerank_pairs passes custom instruction."""
        model = _make_model()
        model._score_text = MagicMock(return_value=0.7)  # type: ignore[method-assign]
        custom_instruction = "Rate this pair."

        list(model.rerank_pairs([("q", "d")], instruction=custom_instruction))

        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
        assert custom_instruction in call_arg

    def test_rerank_pairs_score_range(self):
        """Test rerank_pairs scores are in [0, 1]."""
        model = _make_model()
        scores = list(model.rerank_pairs([("q1", "d1"), ("q2", "d2")]))

        for score in scores:
            assert 0.0 <= score <= 1.0

    def test_rerank_pairs_empty(self):
        """Test rerank_pairs with empty iterable yields nothing."""
        model = _make_model()
        scores = list(model.rerank_pairs([]))

        assert scores == []

    def test_rerank_pairs_formats_input_correctly(self):
        """Test rerank_pairs formats each pair as a proper chat-template string."""
        model = _make_model()
        model._score_text = MagicMock(return_value=0.5)  # type: ignore[method-assign]

        list(model.rerank_pairs([("test query", "test document")]))

        call_arg = model._score_text.call_args[0][0]  # type: ignore[unresolved-attribute]
        assert "test query" in call_arg
        assert "test document" in call_arg
        assert SYSTEM_PROMPT in call_arg
