"""
Integration tests for EndeeVectorStore with real CrewAI workflows.
These tests use actual API credentials if available, or fall back to local mode.

Run with:
    pytest tests/test_integration.py
    # or
    python -m pytest tests/test_integration.py
"""

import os
import time

import pytest
from crewai import LLM, Agent, Crew, Process, Task
from crewai.memory import EntityMemory, ShortTermMemory

# Adjust import based on your actual package structure if needed
from crewai_endee import EndeeVectorStore


@pytest.mark.integration
class TestEndeeCrewAIIntegration:
    """Integration tests with real CrewAI workflows (Dense Only)."""

    @pytest.fixture(scope="class")
    def real_vector_store(self, precision: str | None = None):
        """Create a real EndeeVectorStore instance for integration testing."""
        # It will default to None (local mode) if not set.
        api_token = os.getenv("ENDEE_API_TOKEN")

        # We still need an embedder key (like Cohere or OpenAI) for vectors to work.
        if not os.getenv("COHERE_API_KEY"):
            pytest.skip("COHERE_API_KEY not set (required for embeddings in this test)")

        embedder_config = {
            "provider": "cohere",
            "config": {"model_name": "small", "api_key": os.getenv("COHERE_API_KEY")},
        }

        store = EndeeVectorStore(
            type="test_integration_index",
            api_token=api_token,  # Can be None for local/open source usage
            embedder_config=embedder_config,
            space_type="cosine",
            precision=precision,
        )

        # Reset to ensure clean state
        store.reset()
        time.sleep(2)

        yield store

        # Cleanup
        try:
            store.reset()
        except Exception:
            pass

    def test_end_to_end_crew_workflow(self, real_vector_store):
        """Test complete workflow: create crew, save memory, retrieve memory."""
        # Initialize memory systems
        short_memory = ShortTermMemory(storage=real_vector_store)
        entity_memory = EntityMemory(storage=real_vector_store)

        # Ensure we have an LLM key for the agent
        if not os.environ.get("GOOGLE_API_KEY"):
            pytest.skip("GOOGLE_API_KEY not set (required for LLM in this test)")

        # Create LLM
        llm = LLM(
            model="gemini-2.5-flash-lite",
            api_key=os.environ["GOOGLE_API_KEY"],
        )

        # Create agent
        agent = Agent(
            role="Research Assistant",
            goal="Store and recall programming language information",
            backstory="Expert in programming languages",
            llm=llm,
            verbose=False,
        )

        # Create task
        task = Task(
            description="Remember that Python was created by Guido van Rossum in 1991",
            agent=agent,
            expected_output="Confirmation of stored information",
        )

        # Create and run crew
        crew = Crew(
            agents=[agent],
            tasks=[task],
            process=Process.sequential,
            memory=True,
            short_term_memory=short_memory,
            entity_memory=entity_memory,
            verbose=False,
        )

        # Execute
        crew.kickoff()

        # Verify memory was stored
        time.sleep(3)  # Allow time for indexing

        search_results = real_vector_store.search("Python creator", limit=5)
        assert len(search_results) > 0, "Should find stored memories"

    def test_memory_persistence_across_sessions(self, real_vector_store):
        """Test that memory persists across different crew sessions."""
        # Session 1: Store information
        real_vector_store.save(
            "Go language was created by Google engineers",
            {"topic": "programming", "language": "Go"},
        )
        time.sleep(2)

        # Session 2: Retrieve information
        results = real_vector_store.search("Go programming language", limit=3)

        assert len(results) > 0
        assert "Go" in results[0]["content"]

    def test_multiple_document_retrieval(self, real_vector_store):
        """Test storing and retrieving multiple related documents."""
        documents = [
            ("Python is dynamically typed", {"lang": "Python", "feature": "typing"}),
            ("Python has simple syntax", {"lang": "Python", "feature": "syntax"}),
            ("Python is interpreted", {"lang": "Python", "feature": "execution"}),
        ]

        # Store all documents
        for text, metadata in documents:
            real_vector_store.save(text, metadata)
            time.sleep(0.5)

        time.sleep(3)  # Wait for indexing

        # Search for Python features
        results = real_vector_store.search("Python characteristics", limit=5)

        # We might not get exactly 3 depending on vector similarity, but should get some
        assert len(results) >= 1, "Should retrieve Python documents"

    def test_metadata_filtering(self, real_vector_store):
        """Test search with metadata filters."""
        # Store documents with different metadata
        real_vector_store.save(
            "JavaScript is weakly typed", {"lang": "JavaScript", "category": "web"}
        )
        real_vector_store.save(
            "Java is strongly typed", {"lang": "Java", "category": "enterprise"}
        )
        time.sleep(3)

        # Search with filter
        results = real_vector_store.search(
            query="typed language", limit=5, filter={"category": {"$eq": "web"}}
        )

        # Should only return JavaScript document
        if len(results) > 0:
            for res in results:
                assert "JavaScript" in res["content"] or "web" in str(res["metadata"])


@pytest.mark.integration
class TestHybridIntegration:
    """Integration tests specifically for Hybrid Search (Dense + Sparse)."""

    @pytest.fixture(scope="class")
    def hybrid_store(self):
        """Create a Hybrid EndeeVectorStore instance."""
        api_token = os.getenv("ENDEE_API_TOKEN")

        if not os.getenv("COHERE_API_KEY"):
            pytest.skip("COHERE_API_KEY not set (required for embeddings)")

        embedder_config = {
            "provider": "cohere",
            "config": {"model_name": "small", "api_key": os.getenv("COHERE_API_KEY")},
        }

        # Initialize with hybrid parameters
        store = EndeeVectorStore(
            type="test_integration_hybrid",
            api_token=api_token,
            embedder_config=embedder_config,
            space_type="cosine",
            sparse_model_name="splade_pp",  # Enable Hybrid
            # sparse_dim=30522 # Optional, will auto-detect
        )

        store.reset()
        time.sleep(2)

        yield store

        try:
            store.reset()
        except Exception:
            pass

    def test_hybrid_storage_and_retrieval(self, hybrid_store):
        """Test saving and searching in a hybrid index."""
        # Data with specific keywords that sparse search handles well
        hybrid_store.save("The error code is XJ-99-ZQ.", {"type": "error_log"})
        hybrid_store.save("The weather is sunny in California.", {"type": "weather"})

        time.sleep(3)  # Wait for indexing (sparse generation takes a moment)

        # Search for the specific code (Sparse vectors excel at exact keyword matches)
        results = hybrid_store.search("XJ-99-ZQ", limit=1)

        assert len(results) > 0
        assert "XJ-99-ZQ" in results[0]["content"]

    def test_hybrid_vs_dense_relevance(self, hybrid_store):
        """Test that hybrid search finds relevant documents using keyword signals."""
        # 'FlibberFlabber' is a made-up word. Dense vectors might struggle or hallucinate
        # similarity, but sparse vectors will see it as a unique token.
        hybrid_store.save("This document contains FlibberFlabber.", {"id": 1})
        hybrid_store.save("This document contains nothing interesting.", {"id": 2})

        time.sleep(3)

        results = hybrid_store.search("FlibberFlabber", limit=1)

        assert len(results) > 0
        assert "FlibberFlabber" in results[0]["content"]


@pytest.mark.integration
class TestEndeeVectorStorePerformance:
    """Performance and stress tests."""

    @pytest.fixture
    def performance_store(self):
        """Create store for performance testing."""
        # Allow optional token
        api_token = os.getenv("ENDEE_API_TOKEN")

        if not os.getenv("COHERE_API_KEY"):
            pytest.skip("COHERE_API_KEY not set (required for embeddings)")

        embedder_config = {
            "provider": "cohere",
            "config": {"model_name": "small", "api_key": os.getenv("COHERE_API_KEY")},
        }

        store = EndeeVectorStore(
            type="test_performance_index",
            api_token=api_token,
            embedder_config=embedder_config,
        )

        store.reset()
        time.sleep(2)

        yield store

        try:
            store.reset()
        except Exception:
            pass

    @pytest.mark.slow
    def test_bulk_document_storage(self, performance_store):
        """Test storing many documents."""
        num_docs = 20  # Reduced for faster CI, increase for stress testing

        start_time = time.time()

        for i in range(num_docs):
            performance_store.save(
                f"Document {i}: This is test content about topic {i}",
                {"doc_id": i, "category": f"cat_{i % 5}"},
            )

        elapsed = time.time() - start_time

        print(f"\nStored {num_docs} documents in {elapsed:.2f}s")
        print(f"Average: {elapsed / num_docs:.3f}s per document")

        # Allow time for indexing
        time.sleep(5)

        # Verify all documents are searchable
        results = performance_store.search("test content", limit=num_docs)
        assert len(results) > 0

    @pytest.mark.slow
    def test_search_performance(self, performance_store):
        """Test search performance with various query types."""
        # Prepare test data
        for i in range(10):
            performance_store.save(
                f"Programming language {i} with features {i}", {"lang_id": i}
            )

        time.sleep(3)

        # Test different search scenarios
        queries = [
            "programming language features",
            "language 5",
            "features 8",
        ]

        total_time = 0
        for query in queries:
            start = time.time()
            results = performance_store.search(query, limit=5)
            elapsed = time.time() - start
            total_time += elapsed

            print(
                f"\nQuery: '{query}' took {elapsed:.3f}s, found {len(results)} results"
            )

        avg_time = total_time / len(queries)
        print(f"\nAverage search time: {avg_time:.3f}s")

        # Assert reasonable performance (adjust threshold as needed)
        assert avg_time < 5.0, f"Average search time {avg_time}s exceeds 5s threshold"


@pytest.mark.integration
class TestEndeeVectorStoreErrorRecovery:
    """Test error recovery and edge cases in real scenarios."""

    @pytest.fixture
    def recovery_store(self):
        """Create store for error recovery testing."""
        # Allow optional token
        api_token = os.getenv("ENDEE_API_TOKEN")

        if not os.getenv("COHERE_API_KEY"):
            pytest.skip("COHERE_API_KEY not set (required for embeddings)")

        embedder_config = {
            "provider": "cohere",
            "config": {"model_name": "small", "api_key": os.getenv("COHERE_API_KEY")},
        }

        store = EndeeVectorStore(
            type="test_recovery_index",
            api_token=api_token,
            embedder_config=embedder_config,
        )

        store.reset()
        time.sleep(2)

        yield store

        try:
            store.reset()
        except Exception:
            pass

    def test_empty_search_results(self, recovery_store):
        """Test handling of search with no results."""
        results = recovery_store.search("nonexistent query xyz", limit=5)

        assert isinstance(results, list)
        assert len(results) == 0

    def test_very_long_text(self, recovery_store):
        """Test storing and retrieving very long text."""
        long_text = "A" * 10000  # 10KB of text

        recovery_store.save(long_text, {"type": "long_text"})
        time.sleep(2)

        results = recovery_store.search("AAAA", limit=1)

        assert len(results) > 0
        # Text should be truncated to MAX_LENGTH_BYTES
        assert (
            len(results[0]["content"].encode("utf-8"))
            <= recovery_store.MAX_LENGTH_BYTES
        )


# Helper function to run integration tests
def run_integration_tests():
    """
    Helper function to run integration tests.
    Usage: python -c "from tests.test_integration import run_integration_tests; run_integration_tests()"
    """
    pytest.main([__file__, "--run-integration", "-v", "-s"])


if __name__ == "__main__":
    run_integration_tests()
