#
#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
"""
Tests for ChunkFeedbackService - adjusting chunk weights based on user feedback.

Uses importlib to load chunk_feedback_service.py in isolation so that
test/testcases/test_web_api/common.py (a test-helper module) does not shadow
the project-level common/ package during collection.
"""
import importlib.util
import sys
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import MagicMock

import pytest

pytestmark = pytest.mark.p2

_REPO_ROOT = Path(__file__).resolve().parents[4]


def _load_feedback_module(monkeypatch):
    """Load chunk_feedback_service.py with lightweight stubs for its deps."""
    common_pkg = ModuleType("common")
    common_pkg.__path__ = [str(_REPO_ROOT / "common")]
    monkeypatch.setitem(sys.modules, "common", common_pkg)

    constants_mod = ModuleType("common.constants")
    constants_mod.PAGERANK_FLD = "pagerank_fea"
    monkeypatch.setitem(sys.modules, "common.constants", constants_mod)

    settings_mod = ModuleType("common.settings")
    settings_mod.docStoreConn = MagicMock()
    # Non-ES engines accept pagerank_fea=0; tests below override for elasticsearch/opensearch.
    settings_mod.DOC_ENGINE = "infinity"
    monkeypatch.setitem(sys.modules, "common.settings", settings_mod)
    common_pkg.settings = settings_mod

    rag_pkg = ModuleType("rag")
    rag_pkg.__path__ = []
    monkeypatch.setitem(sys.modules, "rag", rag_pkg)

    rag_nlp_pkg = ModuleType("rag.nlp")
    rag_nlp_pkg.__path__ = []
    rag_nlp_pkg.search = SimpleNamespace(index_name=lambda tid: f"idx-{tid}")
    monkeypatch.setitem(sys.modules, "rag.nlp", rag_nlp_pkg)

    rag_nlp_search_mod = ModuleType("rag.nlp.search")
    rag_nlp_search_mod.index_name = lambda tid: f"idx-{tid}"
    monkeypatch.setitem(sys.modules, "rag.nlp.search", rag_nlp_search_mod)

    services_pkg = ModuleType("api.db.services")
    services_pkg.__path__ = []
    monkeypatch.setitem(sys.modules, "api.db.services", services_pkg)

    spec = importlib.util.spec_from_file_location(
        "api.db.services.chunk_feedback_service",
        _REPO_ROOT / "api" / "db" / "services" / "chunk_feedback_service.py",
    )
    mod = importlib.util.module_from_spec(spec)
    monkeypatch.setitem(
        sys.modules, "api.db.services.chunk_feedback_service", mod
    )
    spec.loader.exec_module(mod)

    return mod, settings_mod


@pytest.fixture
def feedback_env(monkeypatch):
    """Provide (module, settings_stub) for chunk feedback tests."""
    return _load_feedback_module(monkeypatch)


class TestFeedbackRowsFromReference:
    """Chunk id + kb resolution via _feedback_rows_from_reference (single pass)."""

    def test_empty_reference(self, feedback_env):
        mod, _ = feedback_env
        assert mod.ChunkFeedbackService._feedback_rows_from_reference({}) == []
        assert mod.ChunkFeedbackService._feedback_rows_from_reference(None) == []

    def test_reference_with_id_and_dataset(self, feedback_env):
        mod, _ = feedback_env
        reference = {
            "chunks": [
                {"id": "chunk1", "content": "test", "dataset_id": "kb1"},
                {"id": "chunk2", "content": "test2", "dataset_id": "kb1"},
            ]
        }
        rows = mod.ChunkFeedbackService._feedback_rows_from_reference(reference)
        assert [r[0] for r in rows] == ["chunk1", "chunk2"]

    def test_reference_with_chunk_id_and_kb_id(self, feedback_env):
        mod, _ = feedback_env
        reference = {
            "chunks": [
                {"chunk_id": "chunk1", "content": "test", "kb_id": "kb1"},
                {"chunk_id": "chunk2", "content": "test2", "kb_id": "kb1"},
            ]
        }
        rows = mod.ChunkFeedbackService._feedback_rows_from_reference(reference)
        assert [r[0] for r in rows] == ["chunk1", "chunk2"]

    def test_reference_skips_chunks_without_kb(self, feedback_env):
        mod, _ = feedback_env
        reference = {
            "chunks": [
                {"id": "chunk1", "dataset_id": "kb1"},
                {"id": "chunk2", "content": "no kb"},
            ]
        }
        rows = mod.ChunkFeedbackService._feedback_rows_from_reference(reference)
        assert [r[0] for r in rows] == ["chunk1"]

    def test_reference_with_no_chunks(self, feedback_env):
        mod, _ = feedback_env
        reference = {"doc_aggs": [{"doc_id": "doc1"}]}
        assert mod.ChunkFeedbackService._feedback_rows_from_reference(reference) == []

    def test_chunk_id_to_kb_map_matches_row_pairs(self, feedback_env):
        mod, _ = feedback_env
        reference = {
            "chunks": [
                {"id": "a", "dataset_id": "kb1"},
                {"chunk_id": "b", "kb_id": "kb2"},
            ]
        }
        rows = mod.ChunkFeedbackService._feedback_rows_from_reference(reference)
        assert {r[0]: r[1] for r in rows} == {"a": "kb1", "b": "kb2"}


class TestUpdateChunkWeight:
    """Tests for update_chunk_weight method."""

    def test_update_weight_success(self, feedback_env):
        """Should update chunk weight successfully."""
        mod, settings_mod = feedback_env
        settings_mod.DOC_ENGINE = "mysql"
        mock_doc_store = MagicMock()
        mock_doc_store.adjust_chunk_pagerank_fea = None
        mock_doc_store.get.return_value = {"pagerank_fea": 10}
        mock_doc_store.update.return_value = True
        settings_mod.docStoreConn = mock_doc_store

        result = mod.ChunkFeedbackService.update_chunk_weight(
            tenant_id="tenant1",
            chunk_id="chunk1",
            kb_id="kb1",
            delta=1
        )

        assert result is True
        mock_doc_store.update.assert_called_once()

    def test_update_weight_chunk_not_found(self, feedback_env):
        """Should return False if chunk not found."""
        mod, settings_mod = feedback_env
        settings_mod.DOC_ENGINE = "mysql"
        mock_doc_store = MagicMock()
        mock_doc_store.adjust_chunk_pagerank_fea = None
        mock_doc_store.get.return_value = None
        settings_mod.docStoreConn = mock_doc_store

        result = mod.ChunkFeedbackService.update_chunk_weight(
            tenant_id="tenant1",
            chunk_id="chunk1",
            kb_id="kb1",
            delta=1
        )

        assert result is False

    def test_update_weight_clamp_max(self, feedback_env):
        """Should clamp weight to MAX_PAGERANK_WEIGHT."""
        mod, settings_mod = feedback_env
        settings_mod.DOC_ENGINE = "mysql"
        mock_doc_store = MagicMock()
        mock_doc_store.adjust_chunk_pagerank_fea = None
        mock_doc_store.get.return_value = {"pagerank_fea": mod.MAX_PAGERANK_WEIGHT}
        mock_doc_store.update.return_value = True
        settings_mod.docStoreConn = mock_doc_store

        mod.ChunkFeedbackService.update_chunk_weight(
            tenant_id="tenant1",
            chunk_id="chunk1",
            kb_id="kb1",
            delta=10  # Would exceed max
        )

        # Verify the new_value passed to update has clamped weight
        call_args = mock_doc_store.update.call_args
        new_value = call_args[0][1]
        assert new_value["pagerank_fea"] == mod.MAX_PAGERANK_WEIGHT

    def test_update_weight_clamp_min(self, feedback_env):
        """Should clamp weight to MIN_PAGERANK_WEIGHT."""
        mod, settings_mod = feedback_env
        settings_mod.DOC_ENGINE = "mysql"
        mock_doc_store = MagicMock()
        mock_doc_store.adjust_chunk_pagerank_fea = None
        mock_doc_store.get.return_value = {"pagerank_fea": 0}
        mock_doc_store.update.return_value = True
        settings_mod.docStoreConn = mock_doc_store

        mod.ChunkFeedbackService.update_chunk_weight(
            tenant_id="tenant1",
            chunk_id="chunk1",
            kb_id="kb1",
            delta=-10  # Would go below min
        )

        call_args = mock_doc_store.update.call_args
        new_value = call_args[0][1]
        assert new_value["pagerank_fea"] == mod.MIN_PAGERANK_WEIGHT

    def test_update_weight_elasticsearch_uses_atomic_adjust(self, feedback_env):
        """Elasticsearch uses script-based adjust (rank_feature zero handled in script)."""
        mod, settings_mod = feedback_env
        settings_mod.DOC_ENGINE = "elasticsearch"
        mock_doc_store = MagicMock()
        mock_adjust = MagicMock(return_value=True)
        mock_doc_store.adjust_chunk_pagerank_fea = mock_adjust
        settings_mod.docStoreConn = mock_doc_store

        assert mod.ChunkFeedbackService.update_chunk_weight(
            tenant_id="tenant1",
            chunk_id="chunk1",
            kb_id="kb1",
            delta=-1,
        )
        mock_adjust.assert_called_once_with(
            "chunk1",
            "idx-tenant1",
            "kb1",
            -1,
            mod.MIN_PAGERANK_WEIGHT,
            mod.MAX_PAGERANK_WEIGHT,
        )

    def test_update_weight_elasticsearch_forwards_row_id(self, feedback_env):
        """Elasticsearch adjust accepts and forwards row_id without TypeError."""
        mod, settings_mod = feedback_env
        settings_mod.DOC_ENGINE = "elasticsearch"
        mock_doc_store = MagicMock()
        mock_adjust = MagicMock(return_value=True)
        mock_doc_store.adjust_chunk_pagerank_fea = mock_adjust
        settings_mod.docStoreConn = mock_doc_store

        assert mod.ChunkFeedbackService.update_chunk_weight(
            tenant_id="tenant1",
            chunk_id="chunk1",
            kb_id="kb1",
            delta=-1,
            row_id=42,
        )
        mock_adjust.assert_called_once_with(
            "chunk1",
            "idx-tenant1",
            "kb1",
            -1,
            mod.MIN_PAGERANK_WEIGHT,
            mod.MAX_PAGERANK_WEIGHT,
            row_id=42,
        )

    def test_update_weight_opensearch_uses_atomic_adjust(self, feedback_env):
        mod, settings_mod = feedback_env
        settings_mod.DOC_ENGINE = "opensearch"
        mock_doc_store = MagicMock()
        mock_adjust = MagicMock(return_value=True)
        mock_doc_store.adjust_chunk_pagerank_fea = mock_adjust
        settings_mod.docStoreConn = mock_doc_store

        mod.ChunkFeedbackService.update_chunk_weight(
            tenant_id="tenant1",
            chunk_id="chunk1",
            kb_id="kb1",
            delta=-2,
        )
        mock_adjust.assert_called_once_with(
            "chunk1",
            "idx-tenant1",
            "kb1",
            -2,
            mod.MIN_PAGERANK_WEIGHT,
            mod.MAX_PAGERANK_WEIGHT,
        )

    def test_update_weight_opensearch_forwards_row_id(self, feedback_env):
        """OpenSearch adjust accepts and forwards row_id without TypeError."""
        mod, settings_mod = feedback_env
        settings_mod.DOC_ENGINE = "opensearch"
        mock_doc_store = MagicMock()
        mock_adjust = MagicMock(return_value=True)
        mock_doc_store.adjust_chunk_pagerank_fea = mock_adjust
        settings_mod.docStoreConn = mock_doc_store

        mod.ChunkFeedbackService.update_chunk_weight(
            tenant_id="tenant1",
            chunk_id="chunk1",
            kb_id="kb1",
            delta=-2,
            row_id=77,
        )
        mock_adjust.assert_called_once_with(
            "chunk1",
            "idx-tenant1",
            "kb1",
            -2,
            mod.MIN_PAGERANK_WEIGHT,
            mod.MAX_PAGERANK_WEIGHT,
            row_id=77,
        )

    def test_update_weight_infinity_uses_adjust_with_row_id(self, feedback_env):
        """Infinity path passes row_id to adjust_chunk_pagerank_fea."""
        mod, settings_mod = feedback_env
        settings_mod.DOC_ENGINE = "infinity"
        mock_doc_store = MagicMock()
        mock_adjust = MagicMock(return_value=True)
        mock_doc_store.adjust_chunk_pagerank_fea = mock_adjust
        settings_mod.docStoreConn = mock_doc_store

        ok = mod.ChunkFeedbackService.update_chunk_weight(
            tenant_id="tenant1",
            chunk_id="chunk1",
            kb_id="kb1",
            delta=1,
            row_id=42,
        )
        assert ok is True
        mock_adjust.assert_called_once_with(
            "chunk1",
            "idx-tenant1",
            "kb1",
            1,
            mod.MIN_PAGERANK_WEIGHT,
            mod.MAX_PAGERANK_WEIGHT,
            row_id=42,
        )


class TestApplyFeedback:
    """Tests for apply_feedback method."""

    def test_apply_feedback_disabled(self, feedback_env, monkeypatch):
        """Should return early when feature is disabled."""
        mod, _ = feedback_env
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_ENABLED", False)

        result = mod.ChunkFeedbackService.apply_feedback(
            tenant_id="tenant1",
            reference={"chunks": [{"id": "chunk1", "dataset_id": "kb1"}]},
            is_positive=True
        )

        assert result["success_count"] == 0
        assert result["fail_count"] == 0
        assert result.get("disabled") is True

    def test_apply_positive_feedback(self, feedback_env, monkeypatch):
        """Relevance mode splits the per-event budget across chunks (equal when no scores)."""
        mod, _ = feedback_env
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_ENABLED", True)
        mock_update = MagicMock(return_value=True)
        monkeypatch.setattr(
            mod.ChunkFeedbackService, "update_chunk_weight", mock_update
        )

        reference = {
            "chunks": [
                {"id": "chunk1", "dataset_id": "kb1"},
                {"id": "chunk2", "dataset_id": "kb1"},
            ]
        }
        result = mod.ChunkFeedbackService.apply_feedback(
            tenant_id="tenant1",
            reference=reference,
            is_positive=True
        )

        assert result["success_count"] == 1
        assert result["fail_count"] == 0
        assert mock_update.call_count == 1
        mock_update.assert_called_once_with("tenant1", "chunk1", "kb1", 1, row_id=None)

    def test_apply_negative_feedback(self, feedback_env, monkeypatch):
        """Should apply negative feedback with full budget when only one chunk."""
        mod, _ = feedback_env
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_ENABLED", True)
        mock_update = MagicMock(return_value=True)
        monkeypatch.setattr(
            mod.ChunkFeedbackService, "update_chunk_weight", mock_update
        )

        reference = {"chunks": [{"id": "chunk1", "dataset_id": "kb1"}]}
        result = mod.ChunkFeedbackService.apply_feedback(
            tenant_id="tenant1",
            reference=reference,
            is_positive=False
        )

        assert result["success_count"] == 1
        mock_update.assert_called_with("tenant1", "chunk1", "kb1", -1, row_id=None)

    def test_apply_feedback_no_chunks(self, feedback_env, monkeypatch):
        """Should handle empty chunk list gracefully."""
        mod, _ = feedback_env
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_ENABLED", True)

        result = mod.ChunkFeedbackService.apply_feedback(
            tenant_id="tenant1",
            reference={},
            is_positive=True
        )

        assert result["success_count"] == 0
        assert result["fail_count"] == 0
        assert result["chunk_ids"] == []

    def test_apply_feedback_partial_failure(self, feedback_env, monkeypatch):
        """Should count failures correctly (uniform gives each chunk a unit)."""
        mod, _ = feedback_env
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_ENABLED", True)
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_WEIGHTING", "uniform")
        mock_update = MagicMock(side_effect=[True, False])
        monkeypatch.setattr(
            mod.ChunkFeedbackService, "update_chunk_weight", mock_update
        )

        reference = {
            "chunks": [
                {"id": "chunk1", "dataset_id": "kb1"},
                {"id": "chunk2", "dataset_id": "kb1"},
            ]
        }
        result = mod.ChunkFeedbackService.apply_feedback(
            tenant_id="tenant1",
            reference=reference,
            is_positive=True
        )

        assert result["success_count"] == 1
        assert result["fail_count"] == 1

    def test_apply_positive_feedback_uniform_mode(self, feedback_env, monkeypatch):
        """uniform: each cited chunk gets the full increment (legacy)."""
        mod, _ = feedback_env
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_ENABLED", True)
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_WEIGHTING", "uniform")
        mock_update = MagicMock(return_value=True)
        monkeypatch.setattr(
            mod.ChunkFeedbackService, "update_chunk_weight", mock_update
        )
        reference = {
            "chunks": [
                {"id": "chunk1", "dataset_id": "kb1"},
                {"id": "chunk2", "dataset_id": "kb1"},
            ]
        }
        mod.ChunkFeedbackService.apply_feedback(
            tenant_id="tenant1", reference=reference, is_positive=True
        )
        mock_update.assert_any_call("tenant1", "chunk1", "kb1", mod.UPVOTE_WEIGHT_INCREMENT, row_id=None)
        mock_update.assert_any_call("tenant1", "chunk2", "kb1", mod.UPVOTE_WEIGHT_INCREMENT, row_id=None)

    def test_apply_positive_feedback_relevance_weighted(self, feedback_env, monkeypatch):
        """Higher retrieval similarity receives a larger share of the budget."""
        mod, _ = feedback_env
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_ENABLED", True)
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_WEIGHTING", "relevance")
        mock_update = MagicMock(return_value=True)
        monkeypatch.setattr(
            mod.ChunkFeedbackService, "update_chunk_weight", mock_update
        )
        reference = {
            "chunks": [
                {"id": "a", "dataset_id": "kb1", "similarity": 0.9},
                {"id": "b", "dataset_id": "kb1", "similarity": 0.1},
            ]
        }
        mod.ChunkFeedbackService.apply_feedback(
            tenant_id="tenant1", reference=reference, is_positive=True
        )
        mock_update.assert_called_once_with("tenant1", "a", "kb1", 1, row_id=None)

    def test_apply_feedback_passes_row_id_from_reference(self, feedback_env, monkeypatch):
        """row_id from retrieval results flows through to update_chunk_weight."""
        mod, _ = feedback_env
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_ENABLED", True)
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_WEIGHTING", "relevance")
        mock_update = MagicMock(return_value=True)
        monkeypatch.setattr(
            mod.ChunkFeedbackService, "update_chunk_weight", mock_update
        )
        reference = {
            "chunks": [
                {"id": "c1", "dataset_id": "kb1", "similarity": 0.8, "row_id": 99},
            ]
        }
        mod.ChunkFeedbackService.apply_feedback(
            tenant_id="tenant1", reference=reference, is_positive=True
        )
        mock_update.assert_called_once_with("tenant1", "c1", "kb1", 1, row_id=99)


class TestThumbFlipFeedback:
    """Verify that toggling thumbup↔thumbdown applies undo + new (two calls)."""

    @staticmethod
    def _simulate_feedback(mod, monkeypatch, reference, prior_thumb, new_thumb):
        """Reproduce the chat_api thumb-flip logic in isolation."""
        monkeypatch.setattr(mod, "CHUNK_FEEDBACK_ENABLED", True)
        mock_update = MagicMock(return_value=True)
        monkeypatch.setattr(mod.ChunkFeedbackService, "update_chunk_weight", mock_update)

        calls = []

        apply_chunk_feedback = False
        if new_thumb is True:
            apply_chunk_feedback = prior_thumb is not True
        else:
            apply_chunk_feedback = prior_thumb is not False

        if apply_chunk_feedback and reference:
            if isinstance(prior_thumb, bool) and prior_thumb != new_thumb:
                r = mod.ChunkFeedbackService.apply_feedback(
                    tenant_id="t1", reference=reference, is_positive=not prior_thumb,
                )
                calls.append(("undo", r))
            r = mod.ChunkFeedbackService.apply_feedback(
                tenant_id="t1", reference=reference, is_positive=new_thumb is True,
            )
            calls.append(("new", r))

        return calls, mock_update

    def test_toggle_thumbup_to_thumbdown(self, feedback_env, monkeypatch):
        """thumbup→thumbdown: undo (+1→-1) then apply new (-1). Two calls."""
        mod, _ = feedback_env
        ref = {"chunks": [{"id": "c1", "dataset_id": "kb1"}]}
        calls, mock = self._simulate_feedback(mod, monkeypatch, ref, True, False)
        assert len(calls) == 2
        assert calls[0][0] == "undo"
        assert calls[1][0] == "new"

    def test_toggle_thumbdown_to_thumbup(self, feedback_env, monkeypatch):
        """thumbdown→thumbup: undo (-1→+1) then apply new (+1). Two calls."""
        mod, _ = feedback_env
        ref = {"chunks": [{"id": "c1", "dataset_id": "kb1"}]}
        calls, mock = self._simulate_feedback(mod, monkeypatch, ref, False, True)
        assert len(calls) == 2
        assert calls[0][0] == "undo"
        assert calls[1][0] == "new"

    def test_no_prior_to_thumbup(self, feedback_env, monkeypatch):
        """None→thumbup: single apply, no undo."""
        mod, _ = feedback_env
        ref = {"chunks": [{"id": "c1", "dataset_id": "kb1"}]}
        calls, mock = self._simulate_feedback(mod, monkeypatch, ref, None, True)
        assert len(calls) == 1
        assert calls[0][0] == "new"

    def test_same_thumb_no_op(self, feedback_env, monkeypatch):
        """thumbup→thumbup: no feedback at all (apply_chunk_feedback is False)."""
        mod, _ = feedback_env
        ref = {"chunks": [{"id": "c1", "dataset_id": "kb1"}]}
        calls, mock = self._simulate_feedback(mod, monkeypatch, ref, True, True)
        assert len(calls) == 0