"""
Test custom LangChain retriever integration with LDR.

This tests the integration of custom LangChain retrievers with Local Deep Research,
ensuring that users can provide their own retriever implementations as search engines.
"""

import os
import pytest
from unittest.mock import Mock, patch
from typing import List
from langchain_core.retrievers import BaseRetriever, Document
from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
from pydantic import Field

from local_deep_research.api.research_functions import (
    quick_summary,
    detailed_research,
    generate_report,
)


class CustomTestRetriever(BaseRetriever):
    """Custom retriever for testing."""

    documents: List[Document] = Field(default_factory=list)

    def __init__(self, documents: List[Document] = None, **kwargs):
        """Initialize with optional documents."""
        super().__init__(**kwargs)
        self.documents = documents or self._get_default_documents()

    def _get_default_documents(self) -> List[Document]:
        """Get default test documents."""
        return [
            Document(
                page_content="Machine learning is a method of data analysis that automates analytical model building.",
                metadata={
                    "source": "internal_kb",
                    "title": "ML Basics",
                    "url": "internal://ml-basics",
                    "doc_id": "1",
                },
            ),
            Document(
                page_content="Deep learning is a subset of machine learning that uses neural networks with multiple layers.",
                metadata={
                    "source": "internal_kb",
                    "title": "Deep Learning Overview",
                    "url": "internal://dl-overview",
                    "doc_id": "2",
                },
            ),
            Document(
                page_content="Natural language processing enables computers to understand and process human language.",
                metadata={
                    "source": "internal_kb",
                    "title": "NLP Introduction",
                    "url": "internal://nlp-intro",
                    "doc_id": "3",
                },
            ),
        ]

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun = None,
    ) -> List[Document]:
        """Get documents relevant to the query."""
        # Simple keyword matching for testing
        relevant_docs = []
        query_lower = query.lower()

        for doc in self.documents:
            content_lower = doc.page_content.lower()
            if any(word in content_lower for word in query_lower.split()):
                relevant_docs.append(doc)

        # If no matches, return first document as fallback
        if not relevant_docs and self.documents:
            relevant_docs = [self.documents[0]]

        return relevant_docs

    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun = None,
    ) -> List[Document]:
        """Async version."""
        return self._get_relevant_documents(query, run_manager=run_manager)


@pytest.mark.skipif(
    os.environ.get("CI") == "true"
    or os.environ.get("GITHUB_ACTIONS") == "true",
    reason="Langchain integration tests skipped in CI - testing advanced features",
)
class TestCustomLangChainRetriever:
    """Test suite for custom LangChain retriever integration."""

    @pytest.fixture
    def settings_snapshot(self):
        """Create a settings snapshot for testing."""
        return {
            "llm.provider": {"value": "openai", "type": "str"},
            "llm.model": {"value": "gpt-3.5-turbo", "type": "str"},
            "llm.temperature": {"value": 0.7, "type": "float"},
            "llm.openai.api_key": {"value": "test-key", "type": "str"},
            "research.iterations": {"value": 2, "type": "int"},
            "research.questions_per_iteration": {"value": 3, "type": "int"},
            "research.search_engines": {"value": ["custom_kb"], "type": "list"},
            "research.local_context": {"value": 2000, "type": "int"},
            "research.web_context": {"value": 2000, "type": "int"},
            "llm.context_window_unrestricted": {"value": False, "type": "bool"},
            "llm.context_window_size": {"value": 8192, "type": "int"},
            "llm.local_context_window_size": {"value": 4096, "type": "int"},
            "llm.supports_max_tokens": {"value": True, "type": "bool"},
            "llm.max_tokens": {"value": 4096, "type": "int"},
            "rate_limiting.llm_enabled": {"value": False, "type": "bool"},
            "search.tool": {"value": "auto", "type": "str"},
            "search.max_results": {"value": 10, "type": "int"},
            "search.cross_engine_max_results": {"value": 100, "type": "int"},
            "search.cross_engine_use_reddit": {"value": False, "type": "bool"},
            "search.cross_engine_min_date": {"value": None, "type": "str"},
            "search.region": {"value": "us", "type": "str"},
            "search.time_period": {"value": "y", "type": "str"},
            "search.safe_search": {"value": True, "type": "bool"},
            "search.snippets_only": {"value": True, "type": "bool"},
            "search.search_language": {"value": "English", "type": "str"},
            "search.max_filtered_results": {"value": 20, "type": "int"},
        }

    @pytest.fixture
    def mock_llm(self):
        """Create a mock LLM for testing."""
        llm = Mock()
        llm.invoke.return_value = Mock(
            content="This is a summary of machine learning concepts."
        )
        llm.batch.return_value = [
            Mock(content="What are the key algorithms in machine learning?"),
            Mock(content="How does deep learning differ from traditional ML?"),
            Mock(content="What are the applications of NLP?"),
        ]
        return llm

    def test_custom_retriever_basic_usage(self, settings_snapshot, mock_llm):
        """Test basic usage of custom retriever with quick_summary."""
        # Create custom retriever
        custom_retriever = CustomTestRetriever()

        with patch(
            "local_deep_research.config.llm_config.get_llm",
            return_value=mock_llm,
        ):
            with patch(
                "random.randint",
                return_value=12345,
            ):
                # Run quick summary with custom retriever
                result = quick_summary(
                    query="What is machine learning?",
                    retrievers={"custom_kb": custom_retriever},
                    search_tool="custom_kb",
                    settings_snapshot=settings_snapshot,
                    iterations=1,
                    questions_per_iteration=2,
                )

        # Verify results
        assert result is not None
        assert result["research_id"] == 12345
        assert "summary" in result
        assert "machine learning" in result["summary"].lower()
        assert "sources" in result
        assert len(result["sources"]) > 0

        # Check that sources come from our retriever
        for source in result["sources"]:
            if hasattr(source, "metadata"):
                assert source.metadata.get("source") == "internal_kb"

    def test_custom_retriever_with_multiple_retrievers(
        self, settings_snapshot, mock_llm
    ):
        """Test using multiple custom retrievers."""
        # Create different retrievers with different document sets
        tech_docs = [
            Document(
                page_content="Python is a popular programming language for machine learning.",
                metadata={"source": "tech_kb", "title": "Python for ML"},
            ),
            Document(
                page_content="TensorFlow and PyTorch are leading deep learning frameworks.",
                metadata={"source": "tech_kb", "title": "DL Frameworks"},
            ),
        ]

        research_docs = [
            Document(
                page_content="Recent research shows promising results in transformer models.",
                metadata={
                    "source": "research_kb",
                    "title": "Transformer Research",
                },
            ),
            Document(
                page_content="Attention mechanisms have revolutionized NLP.",
                metadata={"source": "research_kb", "title": "Attention in NLP"},
            ),
        ]

        tech_retriever = CustomTestRetriever(documents=tech_docs)
        research_retriever = CustomTestRetriever(documents=research_docs)

        with patch(
            "local_deep_research.config.llm_config.get_llm",
            return_value=mock_llm,
        ):
            with patch(
                "random.randint",
                return_value=67890,
            ):
                result = detailed_research(
                    query="What are the latest developments in deep learning?",
                    retrievers={
                        "tech_kb": tech_retriever,
                        "research_kb": research_retriever,
                    },
                    search_tool="auto",  # Use all retrievers
                    settings_snapshot=settings_snapshot,
                    iterations=2,
                )

        assert result["research_id"] == 67890
        assert (
            len(result["sources"]) >= 2
        )  # Should have sources from both retrievers

        # Check sources come from both retrievers
        source_types = set()
        for source in result["sources"]:
            if hasattr(source, "metadata"):
                source_types.add(source.metadata.get("source"))

        assert "tech_kb" in source_types or "research_kb" in source_types

    def test_custom_retriever_with_web_search_hybrid(
        self, settings_snapshot, mock_llm
    ):
        """Test hybrid search combining custom retriever with web search."""
        custom_retriever = CustomTestRetriever()

        # Mock web search results
        web_results = [
            {
                "url": "https://example.com/ml-trends",
                "title": "ML Trends 2024",
                "content": "Latest trends in machine learning include...",
                "source": "wikipedia",
            }
        ]

        with patch(
            "local_deep_research.config.llm_config.get_llm",
            return_value=mock_llm,
        ):
            with patch(
                "local_deep_research.config.search_config.get_search",
                return_value=web_results,
            ):
                with patch(
                    "random.randint",
                    return_value=11111,
                ):
                    result = quick_summary(
                        query="Machine learning trends and internal best practices",
                        retrievers={"internal_kb": custom_retriever},
                        search_tool="meta",
                        meta_search_config={
                            "retrievers": ["internal_kb"],
                            "engines": ["wikipedia"],
                        },
                        settings_snapshot=settings_snapshot,
                        iterations=1,
                    )

        assert result["research_id"] == 11111
        assert (
            len(result["sources"]) >= 2
        )  # Should have both retriever and web sources

    def test_custom_retriever_empty_results(self, settings_snapshot, mock_llm):
        """Test handling when custom retriever returns no results."""
        # Create retriever with no documents
        empty_retriever = CustomTestRetriever(documents=[])

        with patch(
            "local_deep_research.config.llm_config.get_llm",
            return_value=mock_llm,
        ):
            with patch(
                "random.randint",
                return_value=22222,
            ):
                # Should handle gracefully even with no results
                result = quick_summary(
                    query="Non-existent topic",
                    retrievers={"empty_kb": empty_retriever},
                    search_tool="empty_kb",
                    settings_snapshot=settings_snapshot,
                    iterations=1,
                )

        assert result["research_id"] == 22222
        assert "summary" in result  # Should still generate a summary

    def test_custom_retriever_with_llm_integration(self, settings_snapshot):
        """Test custom retriever with both custom LLM and retriever."""
        from test_custom_langchain_llm import CustomTestLLM

        # Create custom components
        custom_llm = CustomTestLLM()
        custom_retriever = CustomTestRetriever()

        with patch(
            "random.randint",
            return_value=33333,
        ):
            # Use both custom LLM and retriever
            result = quick_summary(
                query="Explain machine learning concepts",
                llms={"custom": custom_llm},
                retrievers={"custom_kb": custom_retriever},
                search_tool="custom_kb",
                settings_snapshot=settings_snapshot,
                iterations=1,
            )

        assert result["research_id"] == 33333
        assert "machine learning" in result["summary"].lower()
        assert len(result["sources"]) > 0

    def test_custom_retriever_error_handling(self, settings_snapshot, mock_llm):
        """Test error handling with failing retriever."""

        class FailingRetriever(BaseRetriever):
            """Retriever that raises errors for testing."""

            def _get_relevant_documents(
                self, query: str, **kwargs
            ) -> List[Document]:
                raise RuntimeError("Retriever failed")

            async def _aget_relevant_documents(
                self, query: str, **kwargs
            ) -> List[Document]:
                raise RuntimeError("Retriever failed")

        failing_retriever = FailingRetriever()

        with patch(
            "local_deep_research.config.llm_config.get_llm",
            return_value=mock_llm,
        ):
            # Should handle the error gracefully
            with pytest.raises(RuntimeError, match="Retriever failed"):
                quick_summary(
                    query="Test query",
                    retrievers={"failing_kb": failing_retriever},
                    search_tool="failing_kb",
                    settings_snapshot=settings_snapshot,
                    iterations=1,
                )

    def test_custom_retriever_with_report_generation(
        self, settings_snapshot, mock_llm
    ):
        """Test custom retriever with report generation."""
        custom_retriever = CustomTestRetriever()

        # Mock more complex LLM responses for report generation
        mock_llm.invoke.side_effect = [
            Mock(
                content="# Executive Summary\n\nMachine learning is transforming industries..."
            ),
            Mock(
                content="## Key Findings\n\n1. ML adoption is accelerating\n2. Deep learning dominates"
            ),
            Mock(
                content="## Recommendations\n\n1. Invest in ML infrastructure\n2. Train staff"
            ),
        ]

        with patch(
            "local_deep_research.config.llm_config.get_llm",
            return_value=mock_llm,
        ):
            with patch(
                "random.randint",
                return_value=44444,
            ):
                result = generate_report(
                    query="Machine learning adoption strategy",
                    retrievers={"internal_docs": custom_retriever},
                    search_tool="internal_docs",
                    settings_snapshot=settings_snapshot,
                    report_type="research_report",
                    iterations=2,
                )

        assert result["research_id"] == 44444
        assert "report" in result
        assert "summary" in result

    def test_custom_retriever_metadata_handling(
        self, settings_snapshot, mock_llm
    ):
        """Test that retriever metadata is properly preserved."""
        # Create retriever with rich metadata
        docs_with_metadata = [
            Document(
                page_content="Advanced ML techniques for production systems.",
                metadata={
                    "source": "internal_kb",
                    "title": "ML in Production",
                    "author": "ML Team",
                    "date": "2024-01-15",
                    "tags": ["ml", "production", "best-practices"],
                    "department": "Engineering",
                },
            )
        ]

        metadata_retriever = CustomTestRetriever(documents=docs_with_metadata)

        with patch(
            "local_deep_research.config.llm_config.get_llm",
            return_value=mock_llm,
        ):
            with patch(
                "random.randint",
                return_value=55555,
            ):
                result = quick_summary(
                    query="ML production best practices",
                    retrievers={"internal": metadata_retriever},
                    search_tool="internal",
                    settings_snapshot=settings_snapshot,
                    iterations=1,
                )

        assert result["research_id"] == 55555

        # Check that metadata is preserved in sources
        for source in result["sources"]:
            if (
                hasattr(source, "metadata")
                and source.metadata.get("source") == "internal_kb"
            ):
                # Verify rich metadata is preserved
                assert "author" in source.metadata
                assert "date" in source.metadata
                assert "tags" in source.metadata
