"""Unit tests for index_search_tools module."""

import pytest
from unittest.mock import MagicMock, AsyncMock, patch, call
from datetime import datetime, timezone
import zipfile
from io import BytesIO
import json
import tempfile
import shutil

from solace_agent_mesh.agent.tools.index_search_tools import (
    _validate_and_extract_zip,
    _get_next_index_search_turn,
    format_location_string,
    extract_locations,
    get_primary_location,
    format_location_range,
    _load_bm25_index,
    _perform_search,
    _format_results_for_llm,
    index_search,
    MAX_ZIP_SIZE,
    MAX_FILE_COUNT,
    MAX_COMPRESSION_RATIO,
    MAX_SINGLE_FILE_SIZE,
    MAX_UNCOMPRESSED_SIZE,
)


class TestValidateAndExtractZip:
    """Tests for ZIP validation and extraction security measures."""

    def test_zip_size_exceeds_limit(self):
        """Test that oversized ZIP files are rejected."""
        # Create a mock ZIP that reports size > MAX_ZIP_SIZE
        large_zip = b"x" * (MAX_ZIP_SIZE + 1)

        with pytest.raises(ValueError, match="ZIP file too large"):
            _validate_and_extract_zip(large_zip, "/tmp/test", "[test]")

    def test_too_many_files_in_zip(self):
        """Test that ZIP files with too many files are rejected."""
        # Create ZIP with too many files
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as zf:
            for i in range(MAX_FILE_COUNT + 1):
                zf.writestr(f"file{i}.txt", "content")

        zip_bytes = zip_buffer.getvalue()

        with pytest.raises(ValueError, match="Too many files in ZIP"):
            _validate_and_extract_zip(zip_bytes, "/tmp/test", "[test]")

    def test_excessive_compression_ratio(self):
        """Test that zip bombs (high compression ratio) are detected."""
        # Create a small compressed file that claims to expand to huge size
        # This simulates a zip bomb by manually creating a ZipInfo with fake sizes
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as zf:
            # Write small file
            zf.writestr("small.txt", "x" * 100)

        # Manually patch the ZipInfo to simulate zip bomb
        with tempfile.TemporaryDirectory() as temp_dir:
            zip_bytes = zip_buffer.getvalue()

            # Modify the zip to have suspicious compression ratio
            with zipfile.ZipFile(BytesIO(zip_bytes), 'r') as zf:
                info = zf.infolist()[0]
                # Can't modify ZipInfo directly in runtime, so this test
                # verifies the logic exists

                # The actual test would need to create a real zip bomb
                # For now, we verify the function accepts valid zips
                try:
                    _validate_and_extract_zip(zip_bytes, temp_dir, "[test]")
                    # Should succeed for normal files
                    assert True
                except ValueError:
                    pytest.fail("Normal ZIP should not raise ValueError")

    def test_path_traversal_in_zip(self):
        """Test that path traversal attempts in ZIP files are blocked."""
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as zf:
            # Try to write file with path traversal
            zf.writestr("../../../etc/passwd", "malicious")

        zip_bytes = zip_buffer.getvalue()

        with pytest.raises(ValueError, match="Path traversal detected"):
            with tempfile.TemporaryDirectory() as temp_dir:
                _validate_and_extract_zip(zip_bytes, temp_dir, "[test]")

    def test_valid_zip_extraction(self):
        """Test that valid ZIP files are extracted successfully."""
        # Create a valid ZIP
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as zf:
            zf.writestr("index/test.txt", "test content")
            zf.writestr("manifest.json", '{"test": "data"}')

        zip_bytes = zip_buffer.getvalue()

        with tempfile.TemporaryDirectory() as temp_dir:
            # Should not raise
            _validate_and_extract_zip(zip_bytes, temp_dir, "[test]")

            # Verify files were extracted
            assert (temp_dir + "/index/test.txt")
            assert (temp_dir + "/manifest.json")


class TestIndexSearchTurnTracking:
    """Tests for search turn counter (citation collision prevention)."""

    def test_first_search_returns_zero(self):
        """Test that first search in session returns turn 0."""
        mock_context = MagicMock()
        mock_context.state = {}

        turn = _get_next_index_search_turn(mock_context)

        assert turn == 0
        assert mock_context.state["index_search_turn_counter"] == 1

    def test_subsequent_searches_increment(self):
        """Test that subsequent searches increment turn number."""
        mock_context = MagicMock()
        mock_context.state = {"index_search_turn_counter": 3}

        turn = _get_next_index_search_turn(mock_context)

        assert turn == 3
        assert mock_context.state["index_search_turn_counter"] == 4

    def test_no_context_returns_zero(self):
        """Test fallback when no context provided."""
        turn = _get_next_index_search_turn(None)
        assert turn == 0


class TestLocationFormatting:
    """Tests for location string formatting."""

    def test_format_page_location(self):
        """Test formatting of page locations."""
        assert format_location_string("physical_page_5", "page") == "Page 5"
        assert format_location_string("physical_page_100", "page") == "Page 100"

    def test_format_paragraph_location(self):
        """Test formatting of paragraph locations."""
        assert format_location_string("physical_paragraph_3", "paragraph") == "Paragraph 3"

    def test_format_slide_location(self):
        """Test formatting of slide locations."""
        assert format_location_string("physical_slide_10", "slide") == "Slide 10"

    def test_format_line_range_location(self):
        """Test formatting of line range locations."""
        assert format_location_string("lines_1-50", "line_range") == "Lines 1-50"

    def test_fallback_for_unknown_format(self):
        """Test fallback for unrecognized location format."""
        assert format_location_string("unknown_format", "page") == "unknown_format"

    def test_extract_multiple_locations(self):
        """Test extracting locations from citation_map."""
        chunk = {
            "citation_map": [
                {"location": "physical_page_5"},
                {"location": "physical_page_6"},
                {"location": "physical_page_7"}
            ],
            "citation_type": "page"
        }

        locations = extract_locations(chunk)
        assert locations == ["Page 5", "Page 6", "Page 7"]

    def test_get_primary_location(self):
        """Test getting first location from chunk."""
        chunk = {
            "citation_map": [
                {"location": "physical_page_5"},
                {"location": "physical_page_6"}
            ],
            "citation_type": "page"
        }

        primary = get_primary_location(chunk)
        assert primary == "Page 5"

    def test_format_single_page_range(self):
        """Test formatting single page as range."""
        chunk = {
            "citation_map": [{"location": "physical_page_5"}],
            "citation_type": "page"
        }

        range_str = format_location_range(chunk)
        assert range_str == "Page 5"

    def test_format_multiple_page_range(self):
        """Test formatting multiple pages as range."""
        chunk = {
            "citation_map": [
                {"location": "physical_page_5"},
                {"location": "physical_page_6"},
                {"location": "physical_page_7"}
            ],
            "citation_type": "page"
        }

        range_str = format_location_range(chunk)
        assert range_str == "Pages 5-7"


class TestLoadBM25Index:
    """Tests for BM25 index loading."""

    @pytest.mark.asyncio
    async def test_index_not_found(self):
        """Test handling when index artifact doesn't exist."""
        mock_artifact_service = AsyncMock()
        mock_artifact_service.load_artifact = AsyncMock(
            return_value={"status": "error", "message": "Not found"}
        )

        # Mock load_artifact_content_or_metadata
        with patch('solace_agent_mesh.agent.tools.index_search_tools.load_artifact_content_or_metadata') as mock_load:
            mock_load.return_value = {"status": "error", "message": "Not found"}

            retriever, manifest = await _load_bm25_index(
                mock_artifact_service,
                "test-app",
                "user-123",
                "session-456",
                None
            )

            assert retriever is None
            assert manifest is None

    @pytest.mark.asyncio
    async def test_corrupted_zip_file(self):
        """Test handling of corrupted ZIP files."""
        mock_artifact_service = AsyncMock()

        # Return corrupted ZIP bytes
        with patch('solace_agent_mesh.agent.tools.index_search_tools.load_artifact_content_or_metadata') as mock_load:
            mock_load.return_value = {
                "status": "success",
                "raw_bytes": b"not a valid zip file"
            }

            retriever, manifest = await _load_bm25_index(
                mock_artifact_service,
                "test-app",
                "user-123",
                "session-456",
                None
            )

            assert retriever is None
            assert manifest is None

    @pytest.mark.asyncio
    async def test_missing_manifest_in_zip(self):
        """Test handling when manifest.json is missing from ZIP."""
        # Create ZIP without manifest
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as zf:
            zf.writestr("index/data.txt", "test")

        with patch('solace_agent_mesh.agent.tools.index_search_tools.load_artifact_content_or_metadata') as mock_load:
            mock_load.return_value = {
                "status": "success",
                "raw_bytes": zip_buffer.getvalue()
            }

            retriever, manifest = await _load_bm25_index(
                AsyncMock(),
                "test-app",
                "user-123",
                "session-456",
                None
            )

            assert retriever is None
            assert manifest is None


class TestPerformSearch:
    """Tests for BM25 search execution."""

    @pytest.mark.asyncio
    async def test_search_returns_results(self):
        """Test that search returns properly formatted results."""
        import numpy as np

        # Mock BM25 retriever
        mock_retriever = MagicMock()
        mock_retriever.scores = {"num_docs": 3}
        mock_retriever.retrieve = MagicMock(
            return_value=(
                np.array([[0, 1, 2]]),  # corpus indices (numpy array)
                np.array([[10.5, 8.3, 5.2]])  # scores (numpy array)
            )
        )

        # Mock manifest
        manifest = {
            "chunks": [
                {
                    "chunk_text": "First result text",
                    "filename": "doc1.txt",
                    "source_file": "doc1.txt",
                    "chunk_id": 0,
                    "doc_id": 0,
                    "chunk_start": 0,
                    "chunk_end": 100,
                    "citation_type": "page",
                    "citation_map": [{"location": "physical_page_1"}],
                    "version": 1
                },
                {
                    "chunk_text": "Second result text",
                    "filename": "doc1.txt",
                    "source_file": "doc1.txt",
                    "chunk_id": 1,
                    "doc_id": 0,
                    "chunk_start": 100,
                    "chunk_end": 200,
                    "citation_type": "page",
                    "citation_map": [{"location": "physical_page_2"}],
                    "version": 1
                },
                {
                    "chunk_text": "Third result text",
                    "filename": "doc2.txt",
                    "source_file": "doc2.txt",
                    "chunk_id": 0,
                    "doc_id": 1,
                    "chunk_start": 0,
                    "chunk_end": 100,
                    "citation_type": "paragraph",
                    "citation_map": [{"location": "physical_paragraph_1"}],
                    "version": 1
                }
            ]
        }

        with patch('solace_agent_mesh.agent.tools.index_search_tools.bm25s') as mock_bm25s:
            mock_bm25s.tokenize = MagicMock(return_value=["query", "tokens"])

            results = await _perform_search(
                mock_retriever,
                manifest,
                "test query",
                top_k=3,
                min_score=0.0,
                search_turn=0
            )

            assert len(results) == 3
            assert results[0]["citation_id"] == "idx0r0"
            assert results[1]["citation_id"] == "idx0r1"
            assert results[2]["citation_id"] == "idx0r2"
            assert results[0]["chunk_text"] == "First result text"
            assert results[0]["relevance_score"] == 1.0  # Highest score normalized to 1.0

    @pytest.mark.asyncio
    async def test_search_filters_by_min_score(self):
        """Test that results below min_score are filtered out."""
        import numpy as np

        mock_retriever = MagicMock()
        mock_retriever.scores = {"num_docs": 3}
        mock_retriever.retrieve = MagicMock(
            return_value=(
                np.array([[0, 1, 2]]),
                np.array([[10.0, 5.0, 2.0]])  # Third result below threshold
            )
        )

        manifest = {
            "chunks": [
                {"chunk_text": "High score", "filename": "doc.txt", "source_file": "doc.txt",
                 "chunk_id": 0, "doc_id": 0, "chunk_start": 0, "chunk_end": 100,
                 "citation_type": "page", "citation_map": [], "version": 1},
                {"chunk_text": "Medium score", "filename": "doc.txt", "source_file": "doc.txt",
                 "chunk_id": 1, "doc_id": 0, "chunk_start": 100, "chunk_end": 200,
                 "citation_type": "page", "citation_map": [], "version": 1},
                {"chunk_text": "Low score", "filename": "doc.txt", "source_file": "doc.txt",
                 "chunk_id": 2, "doc_id": 0, "chunk_start": 200, "chunk_end": 300,
                 "citation_type": "page", "citation_map": [], "version": 1}
            ]
        }

        with patch('solace_agent_mesh.agent.tools.index_search_tools.bm25s') as mock_bm25s:
            mock_bm25s.tokenize = MagicMock(return_value=["query"])

            results = await _perform_search(
                mock_retriever,
                manifest,
                "test query",
                top_k=3,
                min_score=4.0,  # Filter out score of 2.0
                search_turn=0
            )

            # Only 2 results should pass the threshold
            assert len(results) == 2
            assert results[0]["citation_id"] == "idx0r0"
            assert results[1]["citation_id"] == "idx0r1"


    @pytest.mark.asyncio
    async def test_search_clamps_top_k_to_corpus_size(self):
        """Test that top_k is clamped when corpus has fewer chunks than requested."""
        import numpy as np

        mock_retriever = MagicMock()
        mock_retriever.scores = {"num_docs": 2}
        mock_retriever.retrieve = MagicMock(
            return_value=(
                np.array([[0, 1]]),
                np.array([[10.0, 5.0]])
            )
        )

        manifest = {
            "chunks": [
                {"chunk_text": "First", "filename": "doc.txt", "source_file": "doc.txt",
                 "chunk_id": 0, "doc_id": 0, "chunk_start": 0, "chunk_end": 100,
                 "citation_type": "page", "citation_map": [], "version": 1},
                {"chunk_text": "Second", "filename": "doc.txt", "source_file": "doc.txt",
                 "chunk_id": 1, "doc_id": 0, "chunk_start": 100, "chunk_end": 200,
                 "citation_type": "page", "citation_map": [], "version": 1},
            ]
        }

        with patch('solace_agent_mesh.agent.tools.index_search_tools.bm25s') as mock_bm25s:
            mock_bm25s.tokenize = MagicMock(return_value=["query"])

            results = await _perform_search(
                mock_retriever,
                manifest,
                "test query",
                top_k=5,  # Requesting 5 but only 2 in corpus
                min_score=0.0,
                search_turn=0
            )

            # Should clamp to k=2 and return both results without error
            mock_retriever.retrieve.assert_called_once()
            call_args = mock_retriever.retrieve.call_args
            assert call_args[1]["k"] == 2 or call_args[0][1] == 2
            assert len(results) == 2

    @pytest.mark.asyncio
    async def test_search_returns_empty_for_zero_corpus(self):
        """Test that an empty corpus returns no results."""
        mock_retriever = MagicMock()
        mock_retriever.scores = {"num_docs": 0}

        manifest = {"chunks": []}

        with patch('solace_agent_mesh.agent.tools.index_search_tools.bm25s') as mock_bm25s:
            mock_bm25s.tokenize = MagicMock(return_value=["query"])

            results = await _perform_search(
                mock_retriever,
                manifest,
                "test query",
                top_k=5,
                min_score=0.0,
                search_turn=0
            )

            assert results == []
            mock_retriever.retrieve.assert_not_called()


class TestFormatResultsForLLM:
    """Tests for LLM result formatting."""

    def test_format_includes_all_results(self):
        """Test that all results are included in formatted output."""
        results = [
            {
                "citation_id": "idx0r0",
                "chunk_text": "First result",
                "source_file": "doc1.pdf",
                "filename": "doc1.pdf.converted.txt",
                "location_range": "Page 1",
                "score": 10.5
            },
            {
                "citation_id": "idx0r1",
                "chunk_text": "Second result",
                "source_file": "doc2.pdf",
                "filename": "doc2.pdf.converted.txt",
                "location_range": "Page 2",
                "score": 8.3
            }
        ]

        formatted = _format_results_for_llm(
            "test query",
            0,
            results,
            ["idx0r0", "idx0r1"]
        )

        assert "RESULT 1" in formatted
        assert "RESULT 2" in formatted
        assert "[[cite:idx0r0]]" in formatted
        assert "[[cite:idx0r1]]" in formatted
        assert "First result" in formatted
        assert "Second result" in formatted
        assert "doc1.pdf" in formatted
        assert "doc2.pdf" in formatted

    def test_format_includes_citation_instructions(self):
        """Test that citation instructions are included."""
        results = [
            {
                "citation_id": "idx0r0",
                "chunk_text": "Test",
                "source_file": "doc.pdf",
                "filename": "doc.pdf",
                "location_range": "Page 1",
                "score": 10.0
            }
        ]

        formatted = _format_results_for_llm("query", 0, results, ["idx0r0"])

        assert "IMPORTANT CITATION RULES" in formatted
        assert "USE [[cite:idx0r0]] to cite facts" in formatted


class TestIndexSearchEndToEnd:
    """End-to-end tests for index_search function."""

    @pytest.mark.asyncio
    async def test_search_with_no_context_returns_error(self):
        """Test that search without tool context returns error."""
        result = await index_search(
            query="test query",
            tool_context=None
        )

        assert result["status"] == "error"
        assert result["error_code"] == "NO_CONTEXT"

    @pytest.mark.asyncio
    async def test_search_with_no_index_returns_error(self):
        """Test that search with no index returns appropriate error."""
        mock_context = MagicMock()
        mock_context._invocation_context = MagicMock()
        mock_context._invocation_context.artifact_service = AsyncMock()
        mock_context._invocation_context.app_name = "test-app"
        mock_context._invocation_context.user_id = "user-123"
        mock_context._invocation_context.session_id = "session-456"
        mock_context.state = {}

        with patch('solace_agent_mesh.agent.tools.index_search_tools.get_original_session_id') as mock_get_session:
            mock_get_session.return_value = "session-456"

            with patch('solace_agent_mesh.agent.tools.index_search_tools._load_bm25_index') as mock_load:
                mock_load.return_value = (None, None)

                result = await index_search(
                    query="test query",
                    tool_context=mock_context
                )

                assert result["status"] == "error"
                assert result["error_code"] == "INDEX_NOT_FOUND"
                assert "No document index found" in result["message"]

    @pytest.mark.asyncio
    async def test_search_with_no_results(self):
        """Test search that returns no results."""
        mock_context = MagicMock()
        mock_context._invocation_context = MagicMock()
        mock_context._invocation_context.artifact_service = AsyncMock()
        mock_context._invocation_context.app_name = "test-app"
        mock_context._invocation_context.user_id = "user-123"
        mock_context._invocation_context.session_id = "session-456"
        mock_context.state = {}

        mock_retriever = MagicMock()
        mock_manifest = {"chunks": []}

        with patch('solace_agent_mesh.agent.tools.index_search_tools.get_original_session_id') as mock_get_session:
            mock_get_session.return_value = "session-456"

            with patch('solace_agent_mesh.agent.tools.index_search_tools._load_bm25_index') as mock_load:
                mock_load.return_value = (mock_retriever, mock_manifest)

                with patch('solace_agent_mesh.agent.tools.index_search_tools._perform_search') as mock_search:
                    mock_search.return_value = []

                    result = await index_search(
                        query="test query",
                        tool_context=mock_context
                    )

                    assert result["status"] == "success"
                    assert result["num_results"] == 0
                    assert len(result["results"]) == 0
                    assert "No relevant results found" in result["message"]
