"""
Extended tests for AdaptiveQueryGenerator - Adaptive query generation system.

Tests cover:
- QueryPattern dataclass
- AdaptiveQueryGenerator initialization
- Default pattern initialization
- Pattern-based query generation
- Semantic expansion
- LLM-based query generation
- Pattern updates
- Fallback query generation
- Constraint combination optimization

These tests import and test the ACTUAL AdaptiveQueryGenerator class with mocked LLM.
"""

from collections import defaultdict
from unittest.mock import MagicMock


from local_deep_research.advanced_search_system.query_generation.adaptive_query_generator import (
    AdaptiveQueryGenerator,
    QueryPattern,
)
from local_deep_research.advanced_search_system.constraints.base_constraint import (
    Constraint,
    ConstraintType,
)


class TestQueryPatternDataclass:
    """Tests for QueryPattern dataclass."""

    def test_query_pattern_has_template(self):
        """QueryPattern should have template field."""
        pattern = QueryPattern(
            template='"{entity}" {property}',
            constraint_types=[
                ConstraintType.NAME_PATTERN,
                ConstraintType.PROPERTY,
            ],
            success_rate=0.7,
            example_queries=["test query"],
            discovered_entities=set(),
        )
        assert pattern.template == '"{entity}" {property}'

    def test_query_pattern_has_constraint_types(self):
        """QueryPattern should have constraint_types field."""
        pattern = QueryPattern(
            template="template",
            constraint_types=[
                ConstraintType.NAME_PATTERN,
                ConstraintType.PROPERTY,
            ],
            success_rate=0.6,
            example_queries=[],
            discovered_entities=set(),
        )
        assert len(pattern.constraint_types) == 2
        assert ConstraintType.NAME_PATTERN in pattern.constraint_types

    def test_query_pattern_has_success_rate(self):
        """QueryPattern should have success_rate field."""
        pattern = QueryPattern(
            template="template",
            constraint_types=[],
            success_rate=0.75,
            example_queries=[],
            discovered_entities=set(),
        )
        assert pattern.success_rate == 0.75

    def test_query_pattern_has_example_queries(self):
        """QueryPattern should have example_queries field."""
        pattern = QueryPattern(
            template="template",
            constraint_types=[],
            success_rate=0.5,
            example_queries=["query1", "query2"],
            discovered_entities=set(),
        )
        assert len(pattern.example_queries) == 2

    def test_query_pattern_has_discovered_entities(self):
        """QueryPattern should have discovered_entities field."""
        pattern = QueryPattern(
            template="template",
            constraint_types=[],
            success_rate=0.5,
            example_queries=[],
            discovered_entities={"entity1", "entity2"},
        )
        assert len(pattern.discovered_entities) == 2


class TestAdaptiveQueryGeneratorInitialization:
    """Tests for AdaptiveQueryGenerator initialization."""

    def test_model_assignment(self):
        """Should assign model on initialization."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)
        assert generator.model == mock_model

    def test_successful_patterns_initialized(self):
        """successful_patterns should be initialized as list with defaults."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)
        assert isinstance(generator.successful_patterns, list)
        assert len(generator.successful_patterns) > 0  # Has default patterns

    def test_failed_queries_initialized(self):
        """failed_queries should be initialized as set."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)
        assert isinstance(generator.failed_queries, set)
        assert len(generator.failed_queries) == 0

    def test_semantic_expansions_initialized(self):
        """semantic_expansions should be initialized as dict."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)
        assert isinstance(generator.semantic_expansions, dict)
        assert len(generator.semantic_expansions) == 0

    def test_constraint_combinations_initialized(self):
        """constraint_combinations should be initialized as defaultdict."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)
        assert isinstance(generator.constraint_combinations, defaultdict)


class TestDefaultPatternInitialization:
    """Tests for _initialize_default_patterns method."""

    def test_creates_default_patterns(self):
        """Should create default patterns."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        assert len(generator.successful_patterns) >= 3

    def test_entity_location_pattern(self):
        """Should have entity-location pattern."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        templates = [p.template for p in generator.successful_patterns]
        has_entity_pattern = any("entity" in t for t in templates)
        assert has_entity_pattern

    def test_event_temporal_pattern(self):
        """Should have event-temporal pattern."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        templates = [p.template for p in generator.successful_patterns]
        has_event_pattern = any("event" in t for t in templates)
        assert has_event_pattern

    def test_name_comparison_pattern(self):
        """Should have name-comparison pattern."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        templates = [p.template for p in generator.successful_patterns]
        has_and_pattern = any("AND" in t for t in templates)
        assert has_and_pattern


class TestGenerateQuery:
    """Tests for generate_query method."""

    def test_tries_pattern_first(self):
        """Should try pattern-based generation first."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "fallback query"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="mountain",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="height",
            ),
            Constraint(
                id="c3",
                type=ConstraintType.LOCATION,
                description="Location",
                value="Colorado",
            ),
        ]

        query = generator.generate_query(constraints)

        # Should generate from pattern, not invoke LLM
        assert isinstance(query, str)
        assert len(query) > 0

    def test_falls_back_to_expansion(self):
        """Should try semantic expansion when pattern fails."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "synonym1\nsynonym2\nsynonym3"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        # Use constraint types that don't match default patterns
        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.EXISTENCE,
                description="Exists",
                value="feature",
            ),
        ]

        query = generator.generate_query(constraints)

        assert isinstance(query, str)

    def test_falls_back_to_llm(self):
        """Should fall back to LLM generation."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "LLM generated query"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)
        generator.failed_queries.add(
            '"feature" places locations'
        )  # Block expansion result

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.EXISTENCE,
                description="Exists",
                value="feature",
            ),
        ]

        _query = generator.generate_query(constraints)

        # Should invoke LLM
        mock_model.invoke.assert_called()

    def test_skips_failed_queries(self):
        """Should skip queries in failed_queries set."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "new query"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        # Add pattern result to failed queries
        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="mountain",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="formed ice age",
            ),
            Constraint(
                id="c3",
                type=ConstraintType.LOCATION,
                description="Location",
                value="Colorado",
            ),
        ]

        # First get what the pattern would generate
        first_query = generator._generate_from_patterns(constraints)
        if first_query:
            generator.failed_queries.add(first_query)

        # Now generate should try next option
        query = generator.generate_query(constraints)

        assert query != first_query or query is None


class TestGenerateFromPatterns:
    """Tests for _generate_from_patterns method."""

    def test_returns_none_for_no_match(self):
        """Should return None when no patterns match."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        # Clear patterns
        generator.successful_patterns = []

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.PROPERTY,
                description="Test",
                value="value",
            ),
        ]

        result = generator._generate_from_patterns(constraints)

        assert result is None

    def test_selects_highest_success_rate(self):
        """Should select pattern with highest success rate."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        # Add custom patterns with same constraint types but different rates
        generator.successful_patterns = [
            QueryPattern(
                template='low: "{entity}"',
                constraint_types=[ConstraintType.NAME_PATTERN],
                success_rate=0.5,
                example_queries=[],
                discovered_entities=set(),
            ),
            QueryPattern(
                template='high: "{entity}"',
                constraint_types=[ConstraintType.NAME_PATTERN],
                success_rate=0.9,
                example_queries=[],
                discovered_entities=set(),
            ),
        ]

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="test",
            ),
        ]

        result = generator._generate_from_patterns(constraints)

        assert result is not None
        assert "high:" in result

    def test_fills_template_vars(self):
        """Should fill template variables."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        generator.successful_patterns = [
            QueryPattern(
                template='"{entity}" {property}',
                constraint_types=[
                    ConstraintType.NAME_PATTERN,
                    ConstraintType.PROPERTY,
                ],
                success_rate=0.8,
                example_queries=[],
                discovered_entities=set(),
            ),
        ]

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="mountain",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="height",
            ),
        ]

        result = generator._generate_from_patterns(constraints)

        assert result == '"mountain" height'

    def test_handles_missing_vars(self):
        """Should handle missing template variables gracefully."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        generator.successful_patterns = [
            QueryPattern(
                template="{entity} {missing_var}",
                constraint_types=[ConstraintType.NAME_PATTERN],
                success_rate=0.8,
                example_queries=[],
                discovered_entities=set(),
            ),
        ]

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="test",
            ),
        ]

        result = generator._generate_from_patterns(constraints)

        # Should return None on KeyError
        assert result is None

    def test_maps_constraint_types_to_vars(self):
        """Should map constraint types to template variables."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        # Test each constraint type mapping
        generator.successful_patterns = [
            QueryPattern(
                template='"{name_pattern}" {property} {location}',
                constraint_types=[
                    ConstraintType.NAME_PATTERN,
                    ConstraintType.PROPERTY,
                    ConstraintType.LOCATION,
                ],
                success_rate=0.8,
                example_queries=[],
                discovered_entities=set(),
            ),
        ]

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="mountain",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="tall",
            ),
            Constraint(
                id="c3",
                type=ConstraintType.LOCATION,
                description="Location",
                value="Colorado",
            ),
        ]

        result = generator._generate_from_patterns(constraints)

        assert result == '"mountain" tall Colorado'


class TestGenerateWithExpansion:
    """Tests for _generate_with_expansion method."""

    def test_creates_or_clause(self):
        """Should create OR clause with expansions."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "peak\nsummit\nhill"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="mountain",
            ),
        ]

        result = generator._generate_with_expansion(constraints)

        assert "OR" in result
        assert "mountain" in result

    def test_uses_quotes_without_expansion(self):
        """Should use quotes when no expansion available."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = ""  # No expansions
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="specific_term",
            ),
        ]

        result = generator._generate_with_expansion(constraints)

        # Should still produce a query
        assert isinstance(result, str)

    def test_joins_terms_with_and(self):
        """Should join terms with AND."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = ""
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="term1",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="term2",
            ),
        ]

        result = generator._generate_with_expansion(constraints)

        assert "AND" in result

    def test_caches_semantic_expansions(self):
        """Should cache semantic expansions."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "synonym\nrelated"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="test",
            ),
        ]

        # First call
        generator._generate_with_expansion(constraints)

        # Should be cached
        assert "test" in generator.semantic_expansions

        # Second call should use cache
        call_count_before = mock_model.invoke.call_count
        generator._generate_with_expansion(constraints)
        call_count_after = mock_model.invoke.call_count

        # Should not call LLM again for same term
        assert call_count_after == call_count_before


class TestGetSemanticExpansions:
    """Tests for _get_semantic_expansions method."""

    def test_returns_quoted_expansions(self):
        """Should return quoted expansions."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "synonym\nrelated\nalternative"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        expansions = generator._get_semantic_expansions(
            "test", ConstraintType.NAME_PATTERN
        )

        assert all(e.startswith('"') and e.endswith('"') for e in expansions)

    def test_limits_to_three_expansions(self):
        """Should limit to 3 expansions."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "a\nb\nc\nd\ne"  # 5 expansions
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        expansions = generator._get_semantic_expansions(
            "test", ConstraintType.NAME_PATTERN
        )

        assert len(expansions) <= 3


class TestGenerateWithLLM:
    """Tests for _generate_with_llm method."""

    def test_formats_constraints(self):
        """Should format constraints for prompt."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "generated query"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern description",
                value="test",
            ),
        ]

        generator._generate_with_llm(constraints)

        # Check prompt contains constraint info
        call_args = mock_model.invoke.call_args[0][0]
        assert "name_pattern" in call_args
        assert "test" in call_args

    def test_includes_failed_queries_context(self):
        """Should include failed queries in context."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "generated query"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="value",
            ),
        ]
        context = {"failed_queries": ["query1", "query2"]}

        generator._generate_with_llm(constraints, context)

        call_args = mock_model.invoke.call_args[0][0]
        assert "query1" in call_args

    def test_includes_successful_queries_context(self):
        """Should include successful queries in context."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "generated query"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="value",
            ),
        ]
        context = {"successful_queries": ["good1", "good2"]}

        generator._generate_with_llm(constraints, context)

        call_args = mock_model.invoke.call_args[0][0]
        assert "good1" in call_args


class TestUpdatePatterns:
    """Tests for update_patterns method."""

    def test_adds_to_failed_on_failure(self):
        """Should add to failed_queries on failure."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        generator.update_patterns(
            query="failed query",
            constraints=[],
            success=False,
            entities_found=[],
        )

        assert "failed query" in generator.failed_queries

    def test_updates_existing_pattern_success_rate(self):
        """Should update existing pattern success rate when query matches template."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        # Add a pattern with a template that can be extracted from the query
        generator.successful_patterns = [
            QueryPattern(
                template="{name_pattern} search",
                constraint_types=[ConstraintType.NAME_PATTERN],
                success_rate=0.6,
                example_queries=["old query"],
                discovered_entities=set(),
            ),
        ]

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="test",
            ),
        ]

        # Query with the constraint value in it - should extract pattern
        generator.update_patterns(
            query="test search",  # Contains "test" which matches constraint
            constraints=constraints,
            success=True,
            entities_found=["entity1"],
        )

        # Check that a new pattern was added with success_rate 1.0
        # Or the existing one was updated if template matched
        assert len(generator.successful_patterns) >= 1

    def test_adds_failed_query_on_failure_without_entities(self):
        """Should add query to failed_queries when success=False."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="test",
            ),
        ]

        generator.update_patterns(
            query="test query",
            constraints=constraints,
            success=False,
            entities_found=[],
        )

        assert "test query" in generator.failed_queries

    def test_does_not_add_pattern_without_entities(self):
        """Should not add pattern when no entities found even if success=True."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        initial_count = len(generator.successful_patterns)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="test",
            ),
        ]

        generator.update_patterns(
            query="test query",
            constraints=constraints,
            success=True,
            entities_found=[],  # Empty!
        )

        # Should not have added new patterns
        assert len(generator.successful_patterns) == initial_count

    def test_extracts_pattern_from_successful_query(self):
        """Should extract and add pattern from successful query with entities."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        initial_count = len(generator.successful_patterns)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="mountain",
            ),
        ]

        # Query contains the constraint value, so pattern can be extracted
        generator.update_patterns(
            query='"mountain" tall peaks',
            constraints=constraints,
            success=True,
            entities_found=["Mt. Everest"],
        )

        # Should have added a new pattern or updated existing
        new_count = len(generator.successful_patterns)
        assert new_count >= initial_count  # Pattern added or same

    def test_updates_constraint_combinations_correctly(self):
        """Should record constraint combinations on success with entities."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="mountain",
            ),
        ]

        generator.update_patterns(
            query='"mountain" search',  # Contains constraint value
            constraints=constraints,
            success=True,
            entities_found=["entity"],
        )

        # Check that some combination was recorded
        # The exact key depends on sorting behavior of ConstraintType
        assert (
            len(generator.constraint_combinations) > 0
            or len(generator.failed_queries) == 0
        )


class TestExtractPattern:
    """Tests for _extract_pattern method."""

    def test_replaces_values_with_placeholders(self):
        """Should replace values with placeholders."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="mountain",
            ),
        ]

        pattern = generator._extract_pattern(
            '"mountain" in Colorado', constraints
        )

        assert pattern is not None
        assert "{name_pattern}" in pattern.template

    def test_returns_none_without_placeholders(self):
        """Should return None if no placeholders created."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="notinquery",
            ),
        ]

        pattern = generator._extract_pattern("simple query", constraints)

        assert pattern is None

    def test_creates_pattern_with_placeholders(self):
        """Should create pattern when placeholders exist."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Test",
                value="mountain",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.LOCATION,
                description="Location",
                value="Colorado",
            ),
        ]

        pattern = generator._extract_pattern(
            '"mountain" in Colorado', constraints
        )

        assert pattern is not None
        assert pattern.success_rate == 1.0
        assert len(pattern.constraint_types) == 2


class TestFormatConstraints:
    """Tests for _format_constraints method."""

    def test_formats_single_constraint(self):
        """Should format single constraint."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="A pattern",
                value="test",
            ),
        ]

        formatted = generator._format_constraints(constraints)

        assert "name_pattern" in formatted
        assert "A pattern" in formatted
        assert "test" in formatted

    def test_formats_multiple_constraints(self):
        """Should format multiple constraints."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="D1",
                value="v1",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.LOCATION,
                description="D2",
                value="v2",
            ),
        ]

        formatted = generator._format_constraints(constraints)

        assert "name_pattern" in formatted
        assert "location" in formatted


class TestGenerateFallbackQueries:
    """Tests for generate_fallback_queries method."""

    def test_generates_simplified_query(self):
        """Should generate simplified query with fewer constraints."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "alt query 1\nalt query 2"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="a",
                weight=1.0,
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="b",
                weight=0.8,
            ),
            Constraint(
                id="c3",
                type=ConstraintType.LOCATION,
                description="Location",
                value="c",
                weight=0.5,
            ),
        ]

        fallbacks = generator.generate_fallback_queries(
            "original query", constraints
        )

        assert isinstance(fallbacks, list)
        assert len(fallbacks) > 0

    def test_generates_broadened_query(self):
        """Should generate broadened query with OR."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = ""
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="a",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="b",
            ),
        ]

        fallbacks = generator.generate_fallback_queries("original", constraints)

        # Should have OR query
        has_or = any("OR" in q for q in fallbacks)
        assert has_or

    def test_generates_single_constraint_queries(self):
        """Should generate queries for single constraints."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = ""
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="test1",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.LOCATION,
                description="Location",
                value="test2",
            ),
        ]

        fallbacks = generator.generate_fallback_queries("original", constraints)

        # Should have single-constraint queries
        assert len(fallbacks) > 0

    def test_removes_duplicates(self):
        """Should remove duplicate fallback queries."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = (
            '"test1" names list\n"test1" names list'  # Duplicates
        )
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="test1",
            ),
        ]

        fallbacks = generator.generate_fallback_queries("original", constraints)

        # Should not have duplicates
        assert len(fallbacks) == len(set(fallbacks))

    def test_removes_failed_queries(self):
        """Should remove queries in failed set."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = ""
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)
        generator.failed_queries.add('"test1" names list')

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="test1",
            ),
        ]

        fallbacks = generator.generate_fallback_queries("original", constraints)

        # Should not contain failed query
        assert '"test1" names list' not in fallbacks

    def test_limits_to_five(self):
        """Should limit to 5 fallback queries."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "q1\nq2\nq3\nq4\nq5\nq6\nq7"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="test",
                weight=1.0,
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="prop",
                weight=0.8,
            ),
            Constraint(
                id="c3",
                type=ConstraintType.LOCATION,
                description="Location",
                value="loc",
                weight=0.6,
            ),
        ]

        fallbacks = generator.generate_fallback_queries("original", constraints)

        assert len(fallbacks) <= 5


class TestGenerateSingleConstraintQuery:
    """Tests for _generate_single_constraint_query method."""

    def test_name_pattern_template(self):
        """Should use name pattern template."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraint = Constraint(
            id="c1",
            type=ConstraintType.NAME_PATTERN,
            description="Pattern",
            value="mountain",
        )

        query = generator._generate_single_constraint_query(constraint)

        assert '"mountain"' in query
        assert "names list" in query

    def test_location_template(self):
        """Should use location template."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraint = Constraint(
            id="c1",
            type=ConstraintType.LOCATION,
            description="Location",
            value="Colorado",
        )

        query = generator._generate_single_constraint_query(constraint)

        assert '"Colorado"' in query
        assert "places locations" in query

    def test_event_template(self):
        """Should use event template."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraint = Constraint(
            id="c1",
            type=ConstraintType.EVENT,
            description="Event",
            value="earthquake",
        )

        query = generator._generate_single_constraint_query(constraint)

        assert '"earthquake"' in query
        assert "incidents accidents" in query

    def test_temporal_template(self):
        """Should use temporal template."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraint = Constraint(
            id="c1",
            type=ConstraintType.TEMPORAL,
            description="Temporal",
            value="2020",
        )

        query = generator._generate_single_constraint_query(constraint)

        assert "2020" in query
        assert "events in" in query

    def test_default_template(self):
        """Should use default template for unknown types."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraint = Constraint(
            id="c1",
            type=ConstraintType.EXISTENCE,
            description="Existence",
            value="feature",
        )

        query = generator._generate_single_constraint_query(constraint)

        assert "feature" in query


class TestOptimizeConstraintCombinations:
    """Tests for optimize_constraint_combinations method."""

    def test_sorts_by_success(self):
        """Should sort combinations by success."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        # Add some combinations with different success rates
        generator.constraint_combinations[
            (ConstraintType.NAME_PATTERN, ConstraintType.PROPERTY)
        ] = 5
        generator.constraint_combinations[
            (ConstraintType.LOCATION, ConstraintType.EVENT)
        ] = 10
        generator.constraint_combinations[(ConstraintType.TEMPORAL,)] = 3

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="v1",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="v2",
            ),
            Constraint(
                id="c3",
                type=ConstraintType.LOCATION,
                description="Location",
                value="v3",
            ),
            Constraint(
                id="c4",
                type=ConstraintType.EVENT,
                description="Event",
                value="v4",
            ),
        ]

        combinations = generator.optimize_constraint_combinations(constraints)

        # Should return list of constraint lists
        assert isinstance(combinations, list)
        assert len(combinations) > 0

    def test_adds_individual_constraints(self):
        """Should add individual constraints."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="v1",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="v2",
            ),
        ]

        combinations = generator.optimize_constraint_combinations(constraints)

        # Should have individual constraints
        singles = [c for c in combinations if len(c) == 1]
        assert len(singles) >= 2

    def test_adds_pairs(self):
        """Should add constraint pairs."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id="c1",
                type=ConstraintType.NAME_PATTERN,
                description="Pattern",
                value="a",
            ),
            Constraint(
                id="c2",
                type=ConstraintType.PROPERTY,
                description="Property",
                value="b",
            ),
            Constraint(
                id="c3",
                type=ConstraintType.LOCATION,
                description="Location",
                value="c",
            ),
        ]

        combinations = generator.optimize_constraint_combinations(constraints)

        # Should have pairs
        pairs = [c for c in combinations if len(c) == 2]
        assert len(pairs) >= 1

    def test_limits_to_ten(self):
        """Should limit to 10 combinations."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraints = [
            Constraint(
                id=f"c{i}",
                type=ConstraintType.PROPERTY,
                description=f"Prop {i}",
                value=f"v{i}",
            )
            for i in range(10)
        ]

        combinations = generator.optimize_constraint_combinations(constraints)

        assert len(combinations) <= 10


class TestEdgeCases:
    """Tests for edge cases."""

    def test_empty_constraints(self):
        """Should handle empty constraints."""
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "query"
        mock_model.invoke = MagicMock(return_value=mock_response)

        generator = AdaptiveQueryGenerator(model=mock_model)

        query = generator.generate_query([])

        # Should fall back to LLM
        assert isinstance(query, str)

    def test_special_characters_in_value(self):
        """Should handle special characters in value."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraint = Constraint(
            id="c1",
            type=ConstraintType.NAME_PATTERN,
            description="Pattern",
            value="test (with) special [chars]",
        )

        query = generator._generate_single_constraint_query(constraint)

        assert "test" in query
        assert "(" in query

    def test_unicode_in_value(self):
        """Should handle unicode in value."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraint = Constraint(
            id="c1",
            type=ConstraintType.NAME_PATTERN,
            description="Pattern",
            value="\u5bcc\u58eb\u5c71",  # Mt. Fuji in Japanese
        )

        query = generator._generate_single_constraint_query(constraint)

        assert "\u5bcc\u58eb\u5c71" in query

    def test_very_long_value(self):
        """Should handle very long value."""
        mock_model = MagicMock()
        generator = AdaptiveQueryGenerator(model=mock_model)

        constraint = Constraint(
            id="c1",
            type=ConstraintType.NAME_PATTERN,
            description="Pattern",
            value="x" * 1000,
        )

        query = generator._generate_single_constraint_query(constraint)

        assert len(query) > 0
