"""
Tests for advanced_search_system/candidate_exploration/diversity_explorer.py

Tests cover:
- DiversityExplorer initialization
- Diversity score calculation
- Category tracking
- Exploration strategy
"""

import pytest
from unittest.mock import Mock, patch
from collections import defaultdict


@pytest.fixture
def mock_model():
    """Create mock model."""
    return Mock()


@pytest.fixture
def mock_search_engine():
    """Create mock search engine."""
    mock = Mock()
    mock.search.return_value = []
    return mock


class TestDiversityExplorerInit:
    """Tests for DiversityExplorer initialization."""

    def test_init_default_params(self, mock_search_engine, mock_model):
        """Test initialization with default parameters."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
        )

        assert explorer.diversity_threshold == 0.7
        assert explorer.category_limit == 10
        assert explorer.similarity_threshold == 0.8
        assert isinstance(explorer.category_counts, defaultdict)
        assert explorer.diversity_categories == set()

    def test_init_custom_params(self, mock_search_engine, mock_model):
        """Test initialization with custom parameters."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            diversity_threshold=0.5,
            category_limit=5,
            similarity_threshold=0.9,
        )

        assert explorer.diversity_threshold == 0.5
        assert explorer.category_limit == 5
        assert explorer.similarity_threshold == 0.9


class TestDiversityExplorerCategorization:
    """Tests for candidate categorization."""

    def test_category_counts_initialized_empty(
        self, mock_search_engine, mock_model
    ):
        """Category counts should be empty initially."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
        )

        assert len(explorer.category_counts) == 0

    def test_diversity_categories_initialized_empty(
        self, mock_search_engine, mock_model
    ):
        """Diversity categories should be empty initially."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
        )

        assert len(explorer.diversity_categories) == 0


class TestDiversityExplorerScoring:
    """Tests for diversity score calculation."""

    def test_diversity_score_empty_candidates(
        self, mock_search_engine, mock_model
    ):
        """Diversity score for empty candidates should be defined."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
        )

        if hasattr(explorer, "_calculate_diversity_score"):
            score = explorer._calculate_diversity_score([])
            assert 0 <= score <= 1


@pytest.mark.timeout(120)
class TestDiversityExplorerExploration:
    """Tests for exploration methods."""

    def test_explore_returns_exploration_result(
        self, mock_search_engine, mock_model
    ):
        """Test that explore returns an ExplorationResult."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidate_exploration.base_explorer import (
            ExplorationResult,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            max_search_time=2.0,
        )

        # Mock the search to return empty results and ensure loop exits quickly
        with patch.object(explorer, "_execute_search", return_value=[]):
            with patch.object(
                explorer, "_extract_candidates_from_results", return_value=[]
            ):
                # Mock _should_continue_exploration to return False to exit loop
                with patch.object(
                    explorer, "_should_continue_exploration", return_value=False
                ):
                    result = explorer.explore("test query")

                    assert isinstance(result, ExplorationResult)

    def test_explore_tracks_exploration_paths(
        self, mock_search_engine, mock_model
    ):
        """Test that exploration paths are tracked."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            max_search_time=2.0,
        )

        with patch.object(explorer, "_execute_search", return_value=[]):
            with patch.object(
                explorer, "_extract_candidates_from_results", return_value=[]
            ):
                # Mock _should_continue_exploration to return False to exit loop
                with patch.object(
                    explorer, "_should_continue_exploration", return_value=False
                ):
                    result = explorer.explore("test query")

                    assert hasattr(result, "exploration_paths")
                    assert len(result.exploration_paths) >= 1

    def test_explore_completes_successfully(
        self, mock_search_engine, mock_model
    ):
        """Test that exploration completes."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            max_search_time=2.0,
        )

        with patch.object(explorer, "_execute_search", return_value=[]):
            with patch.object(
                explorer, "_extract_candidates_from_results", return_value=[]
            ):
                # Mock _should_continue_exploration to return False to exit loop
                with patch.object(
                    explorer, "_should_continue_exploration", return_value=False
                ):
                    result = explorer.explore("test query")

                    # Should complete successfully
                    assert result is not None


class TestDiversityExplorerShouldContinue:
    """Tests for _should_continue_exploration method."""

    def test_should_continue_true_initially(
        self, mock_search_engine, mock_model
    ):
        """Should continue when exploration just started."""
        import time

        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            max_search_time=2.0,
            max_candidates=100,
        )

        if hasattr(explorer, "_should_continue_exploration"):
            result = explorer._should_continue_exploration(time.time(), 0)
            assert result is True

    def test_should_continue_false_max_candidates(
        self, mock_search_engine, mock_model
    ):
        """Should stop when max candidates reached."""
        import time

        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            max_search_time=2.0,
            max_candidates=100,
        )

        if hasattr(explorer, "_should_continue_exploration"):
            result = explorer._should_continue_exploration(
                time.time(), explorer.max_candidates + 1
            )
            assert result is False

    def test_should_continue_respects_time(
        self, mock_search_engine, mock_model
    ):
        """Should respect time limit in exploration."""
        import time

        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            max_search_time=2.0,
            max_candidates=100,
        )

        if hasattr(explorer, "_should_continue_exploration"):
            # Recently started - should continue
            result = explorer._should_continue_exploration(time.time(), 0)
            assert result is True


class TestDetermineCategoryMethod:
    """Tests for _determine_category method."""

    def test_categorizes_mountain_terms(self, mock_search_engine, mock_model):
        """Test categorizing mountain-related names."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        # Names must contain "mountain", "peak", "summit", or "hill" (case-insensitive)
        for name in ["Big Mountain View", "Rocky Peak", "Grand Summit"]:
            candidate = Candidate(name=name)
            result = explorer._determine_category(candidate)
            assert result == "mountain"

    def test_categorizes_water_terms(self, mock_search_engine, mock_model):
        """Test categorizing water-related names."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        for name in ["Lake Superior", "Colorado River", "Clear Creek"]:
            candidate = Candidate(name=name)
            result = explorer._determine_category(candidate)
            assert result == "water"

    def test_categorizes_park_terms(self, mock_search_engine, mock_model):
        """Test categorizing park-related names."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        for name in ["National Park", "Black Forest", "Nature Reserve"]:
            candidate = Candidate(name=name)
            result = explorer._determine_category(candidate)
            assert result == "park"

    def test_categorizes_trail_terms(self, mock_search_engine, mock_model):
        """Test categorizing trail-related names."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        candidate = Candidate(name="Appalachian Trail")
        result = explorer._determine_category(candidate)
        assert result == "trail"

    def test_categorizes_canyon_terms(self, mock_search_engine, mock_model):
        """Test categorizing canyon-related names."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        for name in ["Grand Canyon", "Box Gorge", "Hidden Valley"]:
            candidate = Candidate(name=name)
            result = explorer._determine_category(candidate)
            assert result == "canyon"

    def test_categorizes_viewpoint_terms(self, mock_search_engine, mock_model):
        """Test categorizing viewpoint-related names."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        for name in ["Eagle Cliff", "Sunset Bluff", "Scenic Overlook"]:
            candidate = Candidate(name=name)
            result = explorer._determine_category(candidate)
            assert result == "viewpoint"

    def test_categorizes_coastal_terms(self, mock_search_engine, mock_model):
        """Test categorizing coastal-related names."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        for name in ["Paradise Island", "Sunny Beach", "Rocky Coast"]:
            candidate = Candidate(name=name)
            result = explorer._determine_category(candidate)
            assert result == "coastal"

    def test_categorizes_place_terms(self, mock_search_engine, mock_model):
        """Test categorizing place-related names."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        for name in ["New York City", "Small Town", "Jefferson County"]:
            candidate = Candidate(name=name)
            result = explorer._determine_category(candidate)
            assert result == "place"

    def test_categorizes_other_terms(self, mock_search_engine, mock_model):
        """Test categorizing unknown names as other."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        candidate = Candidate(name="Some Random Name")
        result = explorer._determine_category(candidate)
        assert result == "other"


class TestCategorizeCandidatesMethod:
    """Tests for _categorize_candidates method."""

    def test_updates_category_counts(self, mock_search_engine, mock_model):
        """Test that category counts are updated."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        candidates = [
            Candidate(name="Hood Mountain"),
            Candidate(name="Lake Tahoe"),
            Candidate(name="Another Mountain Peak"),
        ]

        explorer._categorize_candidates(candidates)

        assert explorer.category_counts["mountain"] == 2
        assert explorer.category_counts["water"] == 1

    def test_adds_metadata_to_candidates(self, mock_search_engine, mock_model):
        """Test that metadata is added to candidates."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        candidate = Candidate(name="Grand Canyon")
        explorer._categorize_candidates([candidate])

        assert candidate.metadata is not None
        assert candidate.metadata["diversity_category"] == "canyon"

    def test_updates_diversity_categories_set(
        self, mock_search_engine, mock_model
    ):
        """Test that diversity_categories set is updated."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        candidates = [
            Candidate(name="Test Mountain"),
            Candidate(name="Lake Test"),
        ]

        explorer._categorize_candidates(candidates)

        assert "mountain" in explorer.diversity_categories
        assert "water" in explorer.diversity_categories


class TestFindUnderrepresentedCategories:
    """Tests for _find_underrepresented_categories method."""

    def test_empty_counts_returns_empty(self, mock_search_engine, mock_model):
        """Test that empty counts returns empty list."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        result = explorer._find_underrepresented_categories()

        assert result == []

    def test_finds_underrepresented(self, mock_search_engine, mock_model):
        """Test finding underrepresented categories."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        # Set up uneven distribution
        explorer.category_counts["mountain"] = 10
        explorer.category_counts["water"] = 10
        explorer.category_counts["trail"] = 1  # Underrepresented

        result = explorer._find_underrepresented_categories()

        assert "trail" in result
        assert "mountain" not in result
        assert "water" not in result


class TestGenerateDiversityQueries:
    """Tests for _generate_diversity_queries method."""

    def test_generates_queries_for_missing_categories(
        self, mock_search_engine, mock_model
    ):
        """Test generating queries for missing categories."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        # Create candidates with only one category
        candidate = Candidate(name="Mount Test")
        candidate.metadata = {"diversity_category": "mountain"}
        candidates = [candidate]

        queries = explorer._generate_diversity_queries("hiking", candidates)

        assert len(queries) > 0

    def test_uses_entity_type_as_base(self, mock_search_engine, mock_model):
        """Test using entity type as base for queries."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        queries = explorer._generate_diversity_queries(
            "base query", [], entity_type="locations"
        )

        # Entity type should be used in queries
        for query in queries:
            assert "locations" in query.lower() or "location" in query.lower()


class TestFilterForDiversity:
    """Tests for _filter_for_diversity method."""

    def test_filters_over_limit_category(self, mock_search_engine, mock_model):
        """Test filtering candidates from over-limit categories."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model, category_limit=2
        )
        explorer.category_counts["mountain"] = 2  # At limit

        new_candidates = [Candidate(name="Another Mountain Peak")]
        existing = []

        filtered = explorer._filter_for_diversity(new_candidates, existing)

        assert len(filtered) == 0

    def test_keeps_candidates_under_limit(self, mock_search_engine, mock_model):
        """Test keeping candidates from under-limit categories."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            category_limit=10,
        )
        explorer.category_counts["water"] = 1  # Under limit

        new_candidates = [Candidate(name="New Lake")]
        existing = []

        filtered = explorer._filter_for_diversity(new_candidates, existing)

        assert len(filtered) == 1


class TestIsSufficientlyDifferent:
    """Tests for _is_sufficiently_different method."""

    def test_similar_candidates_returns_false(
        self, mock_search_engine, mock_model
    ):
        """Test that very similar candidates return False."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            similarity_threshold=0.3,
        )

        candidate = Candidate(name="Mount Rainier Washington")
        existing = [Candidate(name="Mount Rainier Peak")]

        result = explorer._is_sufficiently_different(candidate, existing)

        assert result is False

    def test_different_candidates_returns_true(
        self, mock_search_engine, mock_model
    ):
        """Test that different candidates return True."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine,
            model=mock_model,
            similarity_threshold=0.8,
        )

        candidate = Candidate(name="Grand Canyon Arizona")
        existing = [Candidate(name="Lake Superior Michigan")]

        result = explorer._is_sufficiently_different(candidate, existing)

        assert result is True

    def test_empty_existing_returns_true(self, mock_search_engine, mock_model):
        """Test that empty existing list returns True."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        candidate = Candidate(name="Test Place")

        result = explorer._is_sufficiently_different(candidate, [])

        assert result is True


class TestFinalDiversitySelection:
    """Tests for _final_diversity_selection method."""

    def test_empty_candidates_returns_empty(
        self, mock_search_engine, mock_model
    ):
        """Test that empty candidates returns empty list."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        result = explorer._final_diversity_selection([])

        assert result == []

    def test_selects_from_multiple_categories(
        self, mock_search_engine, mock_model
    ):
        """Test selecting from multiple categories."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model, max_candidates=4
        )

        candidates = []
        for i in range(3):
            candidate = Candidate(name=f"Mountain {i}")
            candidate.metadata = {"diversity_category": "mountain"}
            candidates.append(candidate)

        for i in range(3):
            candidate = Candidate(name=f"Lake {i}")
            candidate.metadata = {"diversity_category": "water"}
            candidates.append(candidate)

        result = explorer._final_diversity_selection(candidates)

        # Should have representation from both categories
        categories = [c.metadata.get("diversity_category") for c in result]
        assert "mountain" in categories
        assert "water" in categories


class TestGenerateExplorationQueries:
    """Tests for generate_exploration_queries method."""

    def test_returns_list(self, mock_search_engine, mock_model):
        """Test that it returns a list."""
        from local_deep_research.advanced_search_system.candidate_exploration.diversity_explorer import (
            DiversityExplorer,
        )
        from local_deep_research.advanced_search_system.candidates.base_candidate import (
            Candidate,
        )

        explorer = DiversityExplorer(
            search_engine=mock_search_engine, model=mock_model
        )

        candidates = [Candidate(name="Test")]

        result = explorer.generate_exploration_queries("base query", candidates)

        assert isinstance(result, list)
