"""
Test custom LangChain LLM integration with LDR.

This tests the integration of custom LangChain LLMs with Local Deep Research,
ensuring that users can provide their own LLM implementations.
"""

import os
import pytest
from unittest.mock import Mock, patch
from langchain_core.language_models.llms import LLM
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from typing import Any, List, Optional

from local_deep_research.api.research_functions import (
    quick_summary,
    detailed_research,
)
from local_deep_research.database.session_context import get_user_db_session
from local_deep_research.settings import SettingsManager


class CustomTestLLM(LLM):
    """Custom LLM for testing."""

    @property
    def _llm_type(self) -> str:
        """Return identifier of llm."""
        return "custom_test_llm"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        """Run the LLM on the given prompt and input."""
        # Simple test response that varies based on prompt content
        if "quantum" in prompt.lower():
            return "Quantum computing uses quantum bits (qubits) that can exist in superposition states."
        elif "machine learning" in prompt.lower():
            return "Machine learning is a subset of AI that enables systems to learn from data."
        elif "climate" in prompt.lower():
            return "Climate change is primarily driven by greenhouse gas emissions."
        else:
            return f"This is a response from the custom LLM about: {prompt[:50]}..."

    @property
    def _identifying_params(self) -> dict:
        """Get the identifying parameters."""
        return {"model_name": "custom_test_llm", "version": "1.0"}


@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 TestCustomLangChainLLM:
    """Test suite for custom LangChain LLM integration."""

    @pytest.fixture
    def mock_session(self):
        """Create a mock database session."""
        session = Mock()
        return session

    @pytest.fixture
    def settings_snapshot(self):
        """Create a settings snapshot for testing."""
        return {
            "llm.provider": {"value": "none", "type": "str"},
            "llm.model": {"value": "custom_test_llm", "type": "str"},
            "llm.temperature": {"value": 0.7, "type": "float"},
            "llm.custom.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": ["wikipedia"], "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": "wikipedia", "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"},
        }

    def test_custom_llm_basic_usage(self, settings_snapshot):
        """Test basic usage of custom LLM with quick_summary."""
        # Create custom LLM instance
        custom_llm = CustomTestLLM()

        # Mock the search results
        with patch(
            "local_deep_research.search_system.get_search"
        ) as _mock_get_search:
            mock_search_engine = Mock()
            mock_search_engine.run.return_value = [
                {
                    "url": "https://example.com/quantum",
                    "title": "Quantum Computing Basics",
                    "content": "Quantum computing is a revolutionary technology...",
                    "source": "wikipedia",
                }
            ]

            _mock_get_search.return_value = mock_search_engine

            # Run quick summary with custom LLM
            result = quick_summary(
                query="What is quantum computing?",
                research_id=12345,
                llms={"custom": custom_llm},
                settings_snapshot=settings_snapshot,
                search_tool="wikipedia",
                iterations=1,
                questions_per_iteration=2,
            )

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

    def test_custom_llm_with_detailed_research(self, settings_snapshot):
        """Test custom LLM with detailed_research function."""
        custom_llm = CustomTestLLM()

        with patch(
            "local_deep_research.config.search_config.get_search"
        ) as mock_search:
            # Mock multiple search results
            mock_search.return_value = [
                {
                    "url": "https://example.com/ml1",
                    "title": "Machine Learning Introduction",
                    "content": "Machine learning is transforming industries...",
                    "source": "wikipedia",
                },
                {
                    "url": "https://example.com/ml2",
                    "title": "ML Applications",
                    "content": "Applications of machine learning include...",
                    "source": "wikipedia",
                },
            ]

            with patch(
                "random.randint",
                return_value=67890,
            ):
                result = detailed_research(
                    query="Explain machine learning applications",
                    llms={"custom": custom_llm},
                    settings_snapshot=settings_snapshot,
                    iterations=2,
                    questions_per_iteration=3,
                )

        assert result is not None
        assert result["research_id"] == 67890
        assert "machine learning" in result["summary"].lower()
        assert len(result["sources"]) >= 2
        assert "findings" in result

    def test_custom_llm_with_custom_factory(self, settings_snapshot):
        """Test using a custom LLM factory function."""

        def create_custom_llm(
            model_name=None, temperature=None, settings_snapshot=None
        ):
            """Factory function for creating custom LLM."""
            # Access settings from snapshot
            api_key = settings_snapshot.get("llm.custom.api_key", {}).get(
                "value"
            )
            assert api_key == "test-key"  # Verify settings access

            # Create and configure custom LLM
            llm = CustomTestLLM()
            # In real implementation, would use api_key and other settings
            return llm

        # Use factory to create LLM
        custom_llm = create_custom_llm(
            model_name="custom_test_llm",
            temperature=0.7,
            settings_snapshot=settings_snapshot,
        )

        with patch(
            "local_deep_research.search_system.get_search"
        ) as _mock_get_search:
            mock_search_engine = Mock()
            mock_search_engine.run.return_value = [
                {
                    "url": "https://example.com/climate",
                    "title": "Climate Change Overview",
                    "content": "Climate change affects global weather patterns...",
                    "source": "wikipedia",
                }
            ]

            with patch(
                "random.randint",
                return_value=11111,
            ):
                result = quick_summary(
                    query="Impact of climate change",
                    llms={"custom": custom_llm},
                    settings_snapshot=settings_snapshot,
                    iterations=1,
                )

        assert result["research_id"] == 11111
        assert "climate" in result["summary"].lower()

    def test_custom_llm_error_handling(self, settings_snapshot):
        """Test error handling with custom LLM."""

        class FailingLLM(LLM):
            """LLM that raises errors for testing."""

            @property
            def _llm_type(self) -> str:
                return "failing_llm"

            def _call(
                self, prompt: str, stop=None, run_manager=None, **kwargs
            ) -> str:
                raise RuntimeError("LLM call failed")

            @property
            def _identifying_params(self) -> dict:
                return {"model_name": "failing_llm"}

        failing_llm = FailingLLM()

        with patch(
            "local_deep_research.search_system.get_search"
        ) as _mock_get_search:
            mock_search_engine = Mock()
            mock_search_engine.run.return_value = [
                {
                    "url": "https://example.com/test",
                    "title": "Test Article",
                    "content": "Test content...",
                    "source": "wikipedia",
                }
            ]

            # Should handle the error gracefully
            with pytest.raises(RuntimeError, match="LLM call failed"):
                quick_summary(
                    query="Test query",
                    llms={"failing": failing_llm},
                    settings_snapshot=settings_snapshot,
                    iterations=1,
                )

    def test_custom_llm_with_real_session(self):
        """Test custom LLM with real session context (integration test)."""
        # Skip if no test database is available
        pytest.skip("Requires test database setup")

        # This would be an integration test with real database
        with get_user_db_session(
            username="testuser", password="testpass"
        ) as session:
            settings_manager = SettingsManager(session, "testuser")
            settings_snapshot = settings_manager.get_all_settings()

            custom_llm = CustomTestLLM()

            result = quick_summary(
                query="Test query",
                llms={"custom": custom_llm},
                settings_snapshot=settings_snapshot,
                iterations=1,
            )

            assert result is not None

    def test_custom_llm_streaming(self, settings_snapshot):
        """Test custom LLM with streaming support."""

        class StreamingTestLLM(LLM):
            """LLM with streaming support for testing."""

            @property
            def _llm_type(self) -> str:
                return "streaming_test_llm"

            def _call(
                self, prompt: str, stop=None, run_manager=None, **kwargs
            ) -> str:
                # Simulate streaming by calling callbacks
                response = (
                    "This is a streaming response about quantum computing."
                )
                if run_manager:
                    for token in response.split():
                        run_manager.on_llm_new_token(token + " ")
                return response

            @property
            def _identifying_params(self) -> dict:
                return {"model_name": "streaming_test_llm", "streaming": True}

        streaming_llm = StreamingTestLLM()

        with patch(
            "local_deep_research.search_system.get_search"
        ) as _mock_get_search:
            mock_search_engine = Mock()
            mock_search_engine.run.return_value = [
                {
                    "url": "https://example.com/quantum",
                    "title": "Quantum Info",
                    "content": "Quantum information...",
                    "source": "wikipedia",
                }
            ]

            with patch(
                "random.randint",
                return_value=99999,
            ):
                result = quick_summary(
                    query="Quantum computing basics",
                    llms={"streaming": streaming_llm},
                    settings_snapshot=settings_snapshot,
                    iterations=1,
                )

        assert result["research_id"] == 99999
        assert "streaming response" in result["summary"]
