#  Copyright © 2025 Dr.-Ing. Paul Wilhelm <paul@wilhelm.dev>
#  This file is part of Archive Agent. See LICENSE for details.

import asyncio
from types import SimpleNamespace
from unittest.mock import Mock

from qdrant_client.models import ScoredPoint

from archive_agent.ai.query.AiQuery import AnswerItem, QuerySchema
from archive_agent.ai.AiManagerFactory import AiManagerFactory
from archive_agent.core.CliManager import CliManager
from archive_agent.db.QdrantManager import QdrantManager
import archive_agent.db.QdrantManager as qdrant_module
from archive_agent.util.knee_detection import find_score_cutoff_index


class FakeAi:
    def __init__(self, rerank_indices, query_schema=None):
        self._rerank_indices = rerank_indices
        self._query_schema = query_schema
        self.rerank_calls = []
        self.query_calls = []

    def embed(self, question):
        return [0.0]

    def rerank(self, question, indexed_chunks):
        self.rerank_calls.append(indexed_chunks)
        return SimpleNamespace(is_rejected=False, reranked_indices=self._rerank_indices)

    def query(self, question, points):
        self.query_calls.append((question, points))
        return self._query_schema


class FakeAiFactory(AiManagerFactory):
    def __init__(self, ai):
        self._ai = ai

    def get_ai(self):
        return self._ai


class FakeQdrantClient:
    def __init__(self, points):
        self._points = points

    async def collection_exists(self, collection_name):
        return True

    async def query_points(self, **kwargs):
        return SimpleNamespace(points=self._points)


def _make_payload(chunk_index, chunk_text):
    return {
        "file_path": "/tmp/test.txt",
        "file_mtime": 0.0,
        "chunk_index": chunk_index,
        "chunks_total": 4,
        "chunk_text": chunk_text,
        "page_range": [1],
        "line_range": None,
    }


def _make_point(score, chunk_index, chunk_text):
    point = Mock(spec=ScoredPoint)
    point.payload = _make_payload(chunk_index, chunk_text)
    point.score = score
    return point


def _make_manager(monkeypatch, points, ai):
    fake_qdrant = FakeQdrantClient(points)
    monkeypatch.setattr(qdrant_module, "AsyncQdrantClient", lambda *args, **kwargs: fake_qdrant)

    cli = CliManager(verbose=False)
    ai_factory = FakeAiFactory(ai)
    return QdrantManager(
        cli=cli,
        ai_factory=ai_factory,
        server_url="http://qdrant.test",
        collection="test",
        vector_size=1,
        retrieve_score_min=0.1,
        retrieve_chunks_max=10,
        retrieve_knee_enable=True,
        retrieve_knee_sensitivity=1.0,
        retrieve_knee_min_chunks=1,
        rerank_chunks_max=10,
        expand_chunks_radius=0,
    )


def _make_query_schema():
    return QuerySchema(
        question_rephrased="What is in scope?",
        answer_list=[
            AnswerItem(
                answer="Answer about retrieval.",
                chunk_ref_list=["<<< 0123456789ABCDEF >>>"]
            )
        ],
        answer_conclusion="Conclusion.",
        follow_up_questions_list=["What else?"],
        is_rejected=False,
        rejection_reason=""
    )


def test_search_applies_knee_cutoff_before_rerank(monkeypatch):
    points = [
        _make_point(0.95, 0, "chunk-0"),
        _make_point(0.94, 1, "chunk-1"),
        _make_point(0.93, 2, "chunk-2"),
        _make_point(0.10, 3, "chunk-3"),
        _make_point(0.09, 4, "chunk-4"),
        _make_point(0.08, 5, "chunk-5"),
    ]
    knee_cutoff = find_score_cutoff_index([p.score for p in points], min_chunks=1)
    assert knee_cutoff is not None
    assert knee_cutoff < len(points)

    rerank_indices = list(range(knee_cutoff - 1, -1, -1))
    ai = FakeAi(rerank_indices=rerank_indices)
    manager = _make_manager(monkeypatch, points, ai)

    result = asyncio.run(manager.search("what is tested"))

    assert result == [points[i] for i in rerank_indices]
    assert ai.rerank_calls == [
        {index: points[index].payload["chunk_text"] for index in range(knee_cutoff)}
    ]


def test_search_without_knee_cutoff_keeps_points(monkeypatch):
    points = [
        _make_point(0.9, 0, "chunk-0"),
        _make_point(0.85, 1, "chunk-1"),
        _make_point(0.8, 2, "chunk-2"),
        _make_point(0.75, 3, "chunk-3"),
        _make_point(0.7, 4, "chunk-4"),
        _make_point(0.65, 5, "chunk-5"),
    ]
    knee_cutoff = find_score_cutoff_index([p.score for p in points], min_chunks=1)
    assert knee_cutoff is None

    rerank_indices = list(range(len(points)))
    ai = FakeAi(rerank_indices=rerank_indices)
    manager = _make_manager(monkeypatch, points, ai)

    result = asyncio.run(manager.search("what is tested"))

    assert result == [points[i] for i in rerank_indices]
    assert ai.rerank_calls == [
        {i: points[i].payload["chunk_text"] for i in range(len(points))}
    ]


def test_expand_points_fetches_neighbors(monkeypatch):
    center_point = _make_point(0.9, 1, "center")
    prev_point = _make_point(0.8, 0, "prev")
    next_point = _make_point(0.7, 2, "next")
    ai = FakeAi(rerank_indices=[])
    manager = _make_manager(monkeypatch, [center_point], ai)
    manager.expand_chunks_radius = 1

    async def fake_get_points(file_path, chunk_indices):
        if chunk_indices == [0]:
            return [prev_point]
        if chunk_indices == [2]:
            return [next_point]
        return []

    monkeypatch.setattr(manager, "_get_points", fake_get_points)

    result = asyncio.run(manager._expand_points([center_point]))

    assert result == [prev_point, center_point, next_point]


def test_dedup_points_keeps_first_occurrence(monkeypatch):
    point_a = _make_point(0.9, 1, "dup-a")
    point_b = _make_point(0.8, 1, "dup-b")
    ai = FakeAi(rerank_indices=[])
    manager = _make_manager(monkeypatch, [], ai)

    result = manager._dedup_points([point_a, point_b])

    assert result == [point_a]


def test_query_uses_search_results_and_formats_answer(monkeypatch):
    points = [
        _make_point(0.9, 0, "chunk-0"),
        _make_point(0.8, 1, "chunk-1"),
    ]
    query_schema = _make_query_schema()
    ai = FakeAi(rerank_indices=[], query_schema=query_schema)
    manager = _make_manager(monkeypatch, points, ai)
    manager.expand_chunks_radius = 0

    async def fake_search(question):
        return points

    monkeypatch.setattr(manager, "search", fake_search)

    result_schema, answer_text = asyncio.run(manager.query("question"))

    assert result_schema == query_schema
    assert "Answer about retrieval." in answer_text
    assert ai.query_calls == [("question", points)]


def test_get_chunks_by_file_returns_sorted_points(monkeypatch):
    """Test that get_chunks_by_file returns points sorted by chunk_index."""
    from archive_agent.db.QdrantSchema import parse_payload

    # Create points in non-sorted order
    points = [
        _make_point(0.9, 2, "chunk-2"),
        _make_point(0.8, 0, "chunk-0"),
        _make_point(0.7, 1, "chunk-1"),
    ]
    ai = FakeAi(rerank_indices=[])
    manager = _make_manager(monkeypatch, points, ai)

    result = asyncio.run(manager.get_chunks_by_file("/tmp/test.txt"))

    # Should be sorted by chunk_index
    assert len(result) == 3
    assert parse_payload(result[0].payload).chunk_index == 0
    assert parse_payload(result[1].payload).chunk_index == 1
    assert parse_payload(result[2].payload).chunk_index == 2


def test_get_chunks_by_file_returns_empty_for_nonexistent_file(monkeypatch):
    """Test that get_chunks_by_file returns empty list for non-existent file."""
    ai = FakeAi(rerank_indices=[])
    manager = _make_manager(monkeypatch, [], ai)

    result = asyncio.run(manager.get_chunks_by_file("/nonexistent/file.txt"))

    assert result == []
