"""Tests for rbsp.query.rewrite -- RewrittenQuery, all RewriteRules, QueryRewriter."""

import os
import sys
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
from rbsp.query.rewrite import (
    RewrittenQuery,
    LowercaseRule,
    StripPunctuationRule,
    CollapseWhitespaceRule,
    StopWordRemovalRule,
    StemRule,
    FieldExtractionRule,
    QuotedPhraseRule,
    OperatorNormalizationRule,
    QueryRewriter,
    RewriteRule,
)


def _make_query(text, **kwargs):
    """Create a RewrittenQuery with the given text."""
    return RewrittenQuery(text=text, **kwargs)


# ---------------------------------------------------------------------------
# RewrittenQuery
# ---------------------------------------------------------------------------

class TestRewrittenQuery(unittest.TestCase):
    """Tests for the RewrittenQuery dataclass."""

    def test_basic_fields(self):
        rq = RewrittenQuery(text="hello world", original="Hello World!")
        self.assertEqual(rq.text, "hello world")
        self.assertEqual(rq.original, "Hello World!")

    def test_default_values(self):
        rq = RewrittenQuery(text="test")
        self.assertEqual(rq.phrases, [])
        self.assertEqual(rq.field_filters, {})
        self.assertEqual(rq.original, "")

    def test_mutable_phrases(self):
        rq = RewrittenQuery(text="test")
        rq.phrases.append("hello world")
        self.assertEqual(rq.phrases, ["hello world"])

    def test_mutable_field_filters(self):
        rq = RewrittenQuery(text="test")
        rq.field_filters["type"] = "python"
        self.assertEqual(rq.field_filters, {"type": "python"})


# ---------------------------------------------------------------------------
# LowercaseRule
# ---------------------------------------------------------------------------

class TestLowercaseRule(unittest.TestCase):
    """Tests for LowercaseRule."""

    def test_lowercases_text(self):
        rule = LowercaseRule()
        rq = rule.rewrite(_make_query("Hello WORLD"))
        self.assertEqual(rq.text, "hello world")

    def test_already_lowercase(self):
        rule = LowercaseRule()
        rq = rule.rewrite(_make_query("already lower"))
        self.assertEqual(rq.text, "already lower")

    def test_empty_string(self):
        rule = LowercaseRule()
        rq = rule.rewrite(_make_query(""))
        self.assertEqual(rq.text, "")


# ---------------------------------------------------------------------------
# StripPunctuationRule
# ---------------------------------------------------------------------------

class TestStripPunctuationRule(unittest.TestCase):
    """Tests for StripPunctuationRule."""

    def test_strips_punctuation(self):
        rule = StripPunctuationRule()
        rq = rule.rewrite(_make_query("hello, world! foo.bar"))
        # Comma, exclamation, period are stripped
        self.assertNotIn(",", rq.text)
        self.assertNotIn("!", rq.text)
        self.assertNotIn(".", rq.text)

    def test_preserves_colons(self):
        rule = StripPunctuationRule()
        rq = rule.rewrite(_make_query("type:python"))
        self.assertIn(":", rq.text)

    def test_preserves_quotes(self):
        rule = StripPunctuationRule()
        rq = rule.rewrite(_make_query('"exact match"'))
        self.assertIn('"', rq.text)

    def test_preserves_alphanumeric(self):
        rule = StripPunctuationRule()
        rq = rule.rewrite(_make_query("abc123"))
        self.assertEqual(rq.text, "abc123")


# ---------------------------------------------------------------------------
# CollapseWhitespaceRule
# ---------------------------------------------------------------------------

class TestCollapseWhitespaceRule(unittest.TestCase):
    """Tests for CollapseWhitespaceRule."""

    def test_collapses_whitespace(self):
        rule = CollapseWhitespaceRule()
        rq = rule.rewrite(_make_query("hello   world"))
        self.assertEqual(rq.text, "hello world")

    def test_strips_edges(self):
        rule = CollapseWhitespaceRule()
        rq = rule.rewrite(_make_query("  hello world  "))
        self.assertEqual(rq.text, "hello world")

    def test_tabs_and_newlines(self):
        rule = CollapseWhitespaceRule()
        rq = rule.rewrite(_make_query("hello\t\nworld"))
        self.assertEqual(rq.text, "hello world")

    def test_empty_string(self):
        rule = CollapseWhitespaceRule()
        rq = rule.rewrite(_make_query(""))
        self.assertEqual(rq.text, "")


# ---------------------------------------------------------------------------
# StopWordRemovalRule
# ---------------------------------------------------------------------------

class TestStopWordRemovalRule(unittest.TestCase):
    """Tests for StopWordRemovalRule."""

    def test_removes_stop_words(self):
        rule = StopWordRemovalRule()
        rq = rule.rewrite(_make_query("the quick brown fox"))
        self.assertNotIn("the", rq.text.split())
        self.assertIn("quick", rq.text.split())
        self.assertIn("brown", rq.text.split())
        self.assertIn("fox", rq.text.split())

    def test_preserves_boolean_operators(self):
        rule = StopWordRemovalRule()
        rq = rule.rewrite(_make_query("python AND java"))
        self.assertIn("AND", rq.text.split())

    def test_preserves_field_value_tokens(self):
        rule = StopWordRemovalRule()
        rq = rule.rewrite(_make_query("type:python the error"))
        self.assertIn("type:python", rq.text.split())
        self.assertNotIn("the", rq.text.split())

    def test_preserves_quoted_regions(self):
        rule = StopWordRemovalRule()
        rq = rule.rewrite(_make_query('"the quick" fox'))
        # "the" inside quotes should be preserved
        self.assertIn('"the', rq.text)

    def test_custom_stop_words(self):
        custom = frozenset({"foo", "bar"})
        rule = StopWordRemovalRule(stop_words=custom)
        rq = rule.rewrite(_make_query("foo bar baz the"))
        tokens = rq.text.split()
        self.assertNotIn("foo", tokens)
        self.assertNotIn("bar", tokens)
        self.assertIn("baz", tokens)
        # "the" is not in our custom list, so it should stay
        self.assertIn("the", tokens)

    def test_all_stop_words(self):
        rule = StopWordRemovalRule()
        rq = rule.rewrite(_make_query("the a an"))
        self.assertEqual(rq.text, "")


# ---------------------------------------------------------------------------
# StemRule
# ---------------------------------------------------------------------------

class TestStemRule(unittest.TestCase):
    """Tests for StemRule."""

    def test_stems_words(self):
        rule = StemRule()
        rq = rule.rewrite(_make_query("running"))
        # Porter stemmer: "running" -> "run"
        self.assertEqual(rq.text, "run")

    def test_preserves_boolean_operators(self):
        rule = StemRule()
        rq = rule.rewrite(_make_query("AND running OR jumping"))
        tokens = rq.text.split()
        self.assertIn("AND", tokens)
        self.assertIn("OR", tokens)

    def test_preserves_field_value_tokens(self):
        rule = StemRule()
        rq = rule.rewrite(_make_query("type:running testing"))
        tokens = rq.text.split()
        self.assertIn("type:running", tokens)

    def test_preserves_quoted_regions(self):
        rule = StemRule()
        rq = rule.rewrite(_make_query('"running fast" testing'))
        # "running" inside quotes should not be stemmed
        self.assertIn('"running', rq.text)


# ---------------------------------------------------------------------------
# FieldExtractionRule
# ---------------------------------------------------------------------------

class TestFieldExtractionRule(unittest.TestCase):
    """Tests for FieldExtractionRule."""

    def test_extracts_field_value(self):
        rule = FieldExtractionRule()
        rq = rule.rewrite(_make_query("type:python error message"))
        self.assertEqual(rq.field_filters, {"type": "python"})
        self.assertNotIn("type:python", rq.text)
        self.assertIn("error", rq.text)
        self.assertIn("message", rq.text)

    def test_multiple_fields(self):
        rule = FieldExtractionRule()
        rq = rule.rewrite(_make_query("type:python lang:en search"))
        self.assertEqual(rq.field_filters["type"], "python")
        self.assertEqual(rq.field_filters["lang"], "en")
        self.assertIn("search", rq.text)

    def test_last_value_wins(self):
        rule = FieldExtractionRule()
        rq = rule.rewrite(_make_query("type:python type:java"))
        self.assertEqual(rq.field_filters["type"], "java")

    def test_no_fields(self):
        rule = FieldExtractionRule()
        rq = rule.rewrite(_make_query("just a query"))
        self.assertEqual(rq.field_filters, {})
        self.assertIn("just", rq.text)

    def test_colon_at_start_not_treated_as_field(self):
        rule = FieldExtractionRule()
        rq = rule.rewrite(_make_query(":value something"))
        # ":value" should not be extracted as a field
        self.assertEqual(rq.field_filters, {})


# ---------------------------------------------------------------------------
# QuotedPhraseRule
# ---------------------------------------------------------------------------

class TestQuotedPhraseRule(unittest.TestCase):
    """Tests for QuotedPhraseRule."""

    def test_extracts_phrase(self):
        rule = QuotedPhraseRule()
        rq = rule.rewrite(_make_query('"hello world" other terms'))
        self.assertEqual(rq.phrases, ["hello world"])
        self.assertNotIn('"hello world"', rq.text)
        self.assertIn("other", rq.text)

    def test_multiple_phrases(self):
        rule = QuotedPhraseRule()
        rq = rule.rewrite(_make_query('"first phrase" middle "second phrase"'))
        self.assertEqual(len(rq.phrases), 2)
        self.assertIn("first phrase", rq.phrases)
        self.assertIn("second phrase", rq.phrases)
        self.assertIn("middle", rq.text)

    def test_empty_quotes_ignored(self):
        rule = QuotedPhraseRule()
        rq = rule.rewrite(_make_query('"" other'))
        self.assertEqual(rq.phrases, [])

    def test_no_quotes(self):
        rule = QuotedPhraseRule()
        rq = rule.rewrite(_make_query("no quotes here"))
        self.assertEqual(rq.phrases, [])

    def test_preserves_existing_phrases(self):
        rule = QuotedPhraseRule()
        rq = _make_query('"new phrase" text', phrases=["old phrase"])
        rq = rule.rewrite(rq)
        self.assertIn("old phrase", rq.phrases)
        self.assertIn("new phrase", rq.phrases)


# ---------------------------------------------------------------------------
# OperatorNormalizationRule
# ---------------------------------------------------------------------------

class TestOperatorNormalizationRule(unittest.TestCase):
    """Tests for OperatorNormalizationRule."""

    def test_lowercase_and(self):
        rule = OperatorNormalizationRule()
        rq = rule.rewrite(_make_query("foo and bar"))
        self.assertEqual(rq.text, "foo AND bar")

    def test_double_ampersand(self):
        rule = OperatorNormalizationRule()
        rq = rule.rewrite(_make_query("foo && bar"))
        self.assertEqual(rq.text, "foo AND bar")

    def test_lowercase_or(self):
        rule = OperatorNormalizationRule()
        rq = rule.rewrite(_make_query("foo or bar"))
        self.assertEqual(rq.text, "foo OR bar")

    def test_double_pipe(self):
        rule = OperatorNormalizationRule()
        rq = rule.rewrite(_make_query("foo || bar"))
        self.assertEqual(rq.text, "foo OR bar")

    def test_lowercase_not(self):
        rule = OperatorNormalizationRule()
        rq = rule.rewrite(_make_query("not foo"))
        self.assertEqual(rq.text, "NOT foo")

    def test_exclamation_mark(self):
        rule = OperatorNormalizationRule()
        rq = rule.rewrite(_make_query("! foo"))
        self.assertEqual(rq.text, "NOT foo")

    def test_already_uppercase(self):
        rule = OperatorNormalizationRule()
        rq = rule.rewrite(_make_query("foo AND bar"))
        self.assertEqual(rq.text, "foo AND bar")

    def test_non_operator_tokens_unchanged(self):
        rule = OperatorNormalizationRule()
        rq = rule.rewrite(_make_query("python java"))
        self.assertEqual(rq.text, "python java")


# ---------------------------------------------------------------------------
# QueryRewriter
# ---------------------------------------------------------------------------

class TestQueryRewriter(unittest.TestCase):
    """Tests for QueryRewriter pipeline."""

    def test_default_pipeline_has_rules(self):
        rewriter = QueryRewriter()
        self.assertGreater(len(rewriter.rules), 0)

    def test_default_pipeline_count(self):
        rewriter = QueryRewriter()
        # 8 default rules
        self.assertEqual(len(rewriter.rules), 8)

    def test_rewrite_preserves_original(self):
        rewriter = QueryRewriter()
        result = rewriter.rewrite("Hello World!")
        self.assertEqual(result.original, "Hello World!")

    def test_rewrite_lowercases(self):
        rewriter = QueryRewriter()
        result = rewriter.rewrite("HELLO WORLD")
        # After full pipeline, text should be lowercased
        self.assertEqual(result.text, result.text.lower())

    def test_rewrite_extracts_fields(self):
        rewriter = QueryRewriter()
        result = rewriter.rewrite("type:python error message")
        self.assertIn("type", result.field_filters)

    def test_rewrite_extracts_phrases(self):
        rewriter = QueryRewriter()
        result = rewriter.rewrite('"exact match" search query')
        self.assertIn("exact match", result.phrases)

    def test_rewrite_removes_stop_words(self):
        rewriter = QueryRewriter()
        result = rewriter.rewrite("the quick brown fox")
        tokens = result.text.split()
        self.assertNotIn("the", tokens)

    def test_rewrite_normalizes_operators(self):
        rewriter = QueryRewriter()
        result = rewriter.rewrite("foo && bar")
        # After operator normalization (then lowercase), "AND" stays uppercase
        # but the StopWordRemovalRule preserves boolean operators
        # Note: LowercaseRule runs after OperatorNormalization,
        # so "AND" becomes "and" -- but StemRule processes it...
        # Let's just check the output is well-formed
        self.assertIsInstance(result.text, str)

    def test_custom_rules(self):
        rewriter = QueryRewriter(rules=[LowercaseRule()])
        result = rewriter.rewrite("HELLO WORLD")
        self.assertEqual(result.text, "hello world")
        # No other transformations
        self.assertEqual(result.phrases, [])
        self.assertEqual(result.field_filters, {})

    def test_add_rule(self):
        rewriter = QueryRewriter(rules=[])
        self.assertEqual(len(rewriter.rules), 0)
        rewriter.add_rule(LowercaseRule())
        self.assertEqual(len(rewriter.rules), 1)
        result = rewriter.rewrite("HELLO")
        self.assertEqual(result.text, "hello")

    def test_rules_returns_copy(self):
        rewriter = QueryRewriter(rules=[LowercaseRule()])
        rules = rewriter.rules
        rules.append(CollapseWhitespaceRule())
        # Modifying the copy should not affect the rewriter
        self.assertEqual(len(rewriter.rules), 1)

    def test_empty_pipeline(self):
        rewriter = QueryRewriter(rules=[])
        result = rewriter.rewrite("Hello World")
        self.assertEqual(result.text, "Hello World")
        self.assertEqual(result.original, "Hello World")

    def test_pipeline_order_matters(self):
        # Lowercase first, then collapse whitespace
        r1 = QueryRewriter(rules=[LowercaseRule(), CollapseWhitespaceRule()])
        result = r1.rewrite("  HELLO   WORLD  ")
        self.assertEqual(result.text, "hello world")


# ---------------------------------------------------------------------------
# Custom RewriteRule subclass
# ---------------------------------------------------------------------------

class TestCustomRewriteRule(unittest.TestCase):
    """Tests for custom rule subclassing."""

    def test_custom_rule(self):
        class PrefixRule(RewriteRule):
            def rewrite(self, query):
                query.text = "prefix: " + query.text
                return query

        rewriter = QueryRewriter(rules=[PrefixRule()])
        result = rewriter.rewrite("hello")
        self.assertEqual(result.text, "prefix: hello")


if __name__ == "__main__":
    unittest.main()
