"""Tests for rbsp.query.plan -- QueryPlanner, QueryPlan, PlanStep."""

import os
import sys
import unittest

sys.path.insert(0, os.path.dirname(__file__))

from rbsp.query.plan import PlanStep, QueryPlan, QueryPlanner


# ---------------------------------------------------------------------------
# PlanStep / QueryPlan dataclass tests
# ---------------------------------------------------------------------------

class TestPlanStep(unittest.TestCase):
    """Test the PlanStep dataclass."""

    def test_defaults(self):
        step = PlanStep(operation="lookup")
        self.assertEqual(step.operation, "lookup")
        self.assertEqual(step.term, "")
        self.assertEqual(step.estimated_docs, 0)
        self.assertEqual(step.cost, 0.0)

    def test_all_fields(self):
        step = PlanStep(operation="intersect", term="foo", estimated_docs=42, cost=3.5)
        self.assertEqual(step.operation, "intersect")
        self.assertEqual(step.term, "foo")
        self.assertEqual(step.estimated_docs, 42)
        self.assertAlmostEqual(step.cost, 3.5)

    def test_frozen(self):
        step = PlanStep(operation="lookup", term="x")
        with self.assertRaises(AttributeError):
            step.term = "y"

    def test_slots(self):
        step = PlanStep(operation="lookup")
        self.assertTrue(hasattr(step, "__slots__") or hasattr(type(step), "__slots__"))


class TestQueryPlan(unittest.TestCase):
    """Test the QueryPlan dataclass."""

    def test_empty_plan(self):
        plan = QueryPlan(steps=[], estimated_cost=0.0, strategy="empty")
        self.assertEqual(len(plan.steps), 0)
        self.assertAlmostEqual(plan.estimated_cost, 0.0)
        self.assertEqual(plan.strategy, "empty")

    def test_plan_with_steps(self):
        s1 = PlanStep(operation="lookup", term="a", estimated_docs=10, cost=10.0)
        s2 = PlanStep(operation="intersect", term="b", estimated_docs=1, cost=10.0)
        plan = QueryPlan(steps=[s1, s2], estimated_cost=20.0, strategy="and_rarest_first")
        self.assertEqual(len(plan.steps), 2)
        self.assertEqual(plan.steps[0].term, "a")
        self.assertEqual(plan.steps[1].term, "b")

    def test_frozen(self):
        plan = QueryPlan(steps=[], estimated_cost=0.0, strategy="empty")
        with self.assertRaises(AttributeError):
            plan.strategy = "other"


# ---------------------------------------------------------------------------
# QueryPlanner construction
# ---------------------------------------------------------------------------

class TestQueryPlannerInit(unittest.TestCase):
    """Test QueryPlanner construction."""

    def test_default_construction(self):
        planner = QueryPlanner()
        self.assertEqual(planner.doc_freqs, {})
        self.assertEqual(planner.total_docs, 0)

    def test_with_doc_freqs(self):
        freqs = {"alpha": 100, "beta": 50}
        planner = QueryPlanner(doc_freqs=freqs, total_docs=1000)
        self.assertEqual(planner.total_docs, 1000)
        self.assertEqual(planner.doc_freqs["alpha"], 100)

    def test_doc_freqs_returns_copy(self):
        freqs = {"x": 5}
        planner = QueryPlanner(doc_freqs=freqs)
        copy = planner.doc_freqs
        copy["y"] = 99
        self.assertNotIn("y", planner.doc_freqs)


# ---------------------------------------------------------------------------
# plan() — empty and single-term
# ---------------------------------------------------------------------------

class TestPlanEmptyAndSingle(unittest.TestCase):
    """Test edge cases: empty queries and single-term queries."""

    def setUp(self):
        self.planner = QueryPlanner(
            doc_freqs={"rare": 5, "common": 500, "medium": 100},
            total_docs=1000,
        )

    def test_empty_query(self):
        plan = self.planner.plan([])
        self.assertEqual(plan.strategy, "empty")
        self.assertEqual(len(plan.steps), 0)
        self.assertAlmostEqual(plan.estimated_cost, 0.0)

    def test_single_term_known(self):
        plan = self.planner.plan(["rare"])
        self.assertEqual(plan.strategy, "single_term")
        self.assertEqual(len(plan.steps), 1)
        self.assertEqual(plan.steps[0].operation, "lookup")
        self.assertEqual(plan.steps[0].term, "rare")
        self.assertEqual(plan.steps[0].estimated_docs, 5)
        self.assertAlmostEqual(plan.estimated_cost, 5.0)

    def test_single_term_unknown(self):
        plan = self.planner.plan(["missing"])
        self.assertEqual(plan.strategy, "single_term")
        self.assertEqual(plan.steps[0].estimated_docs, 0)
        self.assertAlmostEqual(plan.estimated_cost, 0.0)


# ---------------------------------------------------------------------------
# plan() — AND queries
# ---------------------------------------------------------------------------

class TestPlanAnd(unittest.TestCase):
    """Test AND query planning: rarest term first."""

    def setUp(self):
        self.planner = QueryPlanner(
            doc_freqs={"rare": 5, "common": 500, "medium": 100},
            total_docs=1000,
        )

    def test_and_orders_rarest_first(self):
        plan = self.planner.plan(["common", "rare", "medium"], "AND")
        self.assertEqual(plan.strategy, "and_rarest_first")
        terms = [s.term for s in plan.steps]
        self.assertEqual(terms, ["rare", "medium", "common"])

    def test_and_first_step_is_lookup(self):
        plan = self.planner.plan(["common", "rare"], "AND")
        self.assertEqual(plan.steps[0].operation, "lookup")
        self.assertEqual(plan.steps[0].term, "rare")
        self.assertEqual(plan.steps[0].estimated_docs, 5)

    def test_and_subsequent_steps_are_intersect(self):
        plan = self.planner.plan(["common", "rare", "medium"], "AND")
        self.assertEqual(plan.steps[1].operation, "intersect")
        self.assertEqual(plan.steps[2].operation, "intersect")

    def test_and_cost_accumulates(self):
        plan = self.planner.plan(["common", "rare"], "AND")
        # First step: cost = 5 (df of "rare")
        # Second step: cost = min(5.0, 500.0) = 5.0
        self.assertAlmostEqual(plan.steps[0].cost, 5.0)
        self.assertAlmostEqual(plan.steps[1].cost, 5.0)
        self.assertAlmostEqual(plan.estimated_cost, 10.0)

    def test_and_selectivity_shrinks_docs(self):
        plan = self.planner.plan(["common", "rare"], "AND")
        # After first lookup: running_docs = 5
        # After intersect: running_docs = 5 * 0.1 = 0.5 -> int(0.5) = 0
        self.assertEqual(plan.steps[1].estimated_docs, 0)

    def test_and_three_terms_progressive_shrinkage(self):
        planner = QueryPlanner(
            doc_freqs={"a": 1000, "b": 500, "c": 100},
            total_docs=2000,
        )
        plan = planner.plan(["a", "b", "c"], "AND")
        # Order: c(100), b(500), a(1000)
        self.assertEqual([s.term for s in plan.steps], ["c", "b", "a"])
        # Step 1: lookup c, docs=100, cost=100
        self.assertEqual(plan.steps[0].estimated_docs, 100)
        # Step 2: intersect b, docs=100*0.1=10, cost=min(100,500)=100
        self.assertEqual(plan.steps[1].estimated_docs, 10)
        self.assertAlmostEqual(plan.steps[1].cost, 100.0)
        # Step 3: intersect a, docs=10*0.1=1, cost=min(10,1000)=10
        self.assertEqual(plan.steps[2].estimated_docs, 1)
        self.assertAlmostEqual(plan.steps[2].cost, 10.0)

    def test_and_default_type(self):
        plan = self.planner.plan(["common", "rare"])
        self.assertEqual(plan.strategy, "and_rarest_first")

    def test_and_case_insensitive_type(self):
        plan = self.planner.plan(["common", "rare"], "and")
        self.assertEqual(plan.strategy, "and_rarest_first")


# ---------------------------------------------------------------------------
# plan() — OR queries
# ---------------------------------------------------------------------------

class TestPlanOr(unittest.TestCase):
    """Test OR query planning: most common term first."""

    def setUp(self):
        self.planner = QueryPlanner(
            doc_freqs={"rare": 5, "common": 500, "medium": 100},
            total_docs=1000,
        )

    def test_or_orders_most_common_first(self):
        plan = self.planner.plan(["rare", "medium", "common"], "OR")
        self.assertEqual(plan.strategy, "or_largest_first")
        terms = [s.term for s in plan.steps]
        self.assertEqual(terms, ["common", "medium", "rare"])

    def test_or_first_step_is_lookup(self):
        plan = self.planner.plan(["rare", "common"], "OR")
        self.assertEqual(plan.steps[0].operation, "lookup")
        self.assertEqual(plan.steps[0].term, "common")

    def test_or_subsequent_steps_are_union(self):
        plan = self.planner.plan(["rare", "common", "medium"], "OR")
        self.assertEqual(plan.steps[1].operation, "union")
        self.assertEqual(plan.steps[2].operation, "union")

    def test_or_cost_is_sum_of_dfs(self):
        plan = self.planner.plan(["rare", "common"], "OR")
        self.assertAlmostEqual(plan.estimated_cost, 505.0)

    def test_or_three_terms(self):
        plan = self.planner.plan(["rare", "medium", "common"], "OR")
        self.assertAlmostEqual(plan.estimated_cost, 605.0)

    def test_or_case_insensitive_type(self):
        plan = self.planner.plan(["rare"], "or")
        self.assertEqual(plan.strategy, "single_term")


# ---------------------------------------------------------------------------
# plan_boolean() — complex boolean queries
# ---------------------------------------------------------------------------

class TestPlanBoolean(unittest.TestCase):
    """Test boolean query planning with mixed AND/OR/NOT."""

    def setUp(self):
        self.planner = QueryPlanner(
            doc_freqs={"rare": 5, "common": 500, "medium": 100, "noise": 800},
            total_docs=1000,
        )

    def test_empty_clauses(self):
        plan = self.planner.plan_boolean([])
        self.assertEqual(plan.strategy, "empty")
        self.assertEqual(len(plan.steps), 0)

    def test_and_only_boolean(self):
        clauses = [
            {"op": "AND", "term": "common"},
            {"op": "AND", "term": "rare"},
        ]
        plan = self.planner.plan_boolean(clauses)
        self.assertEqual(plan.strategy, "and_rarest_first")
        self.assertEqual(plan.steps[0].term, "rare")
        self.assertEqual(plan.steps[1].term, "common")

    def test_or_only_boolean(self):
        clauses = [
            {"op": "OR", "term": "rare"},
            {"op": "OR", "term": "common"},
        ]
        plan = self.planner.plan_boolean(clauses)
        self.assertEqual(plan.strategy, "or_largest_first")
        self.assertEqual(plan.steps[0].term, "common")

    def test_mixed_and_or(self):
        clauses = [
            {"op": "AND", "term": "common"},
            {"op": "AND", "term": "rare"},
            {"op": "OR", "term": "medium"},
        ]
        plan = self.planner.plan_boolean(clauses)
        self.assertEqual(plan.strategy, "boolean_mixed")
        # AND terms processed first (rarest first)
        self.assertEqual(plan.steps[0].term, "rare")
        self.assertEqual(plan.steps[0].operation, "lookup")
        self.assertEqual(plan.steps[1].term, "common")
        self.assertEqual(plan.steps[1].operation, "intersect")
        # Then OR
        self.assertEqual(plan.steps[2].term, "medium")
        self.assertEqual(plan.steps[2].operation, "union")

    def test_not_clauses_at_end(self):
        clauses = [
            {"op": "NOT", "term": "noise"},
            {"op": "AND", "term": "rare"},
            {"op": "AND", "term": "common"},
        ]
        plan = self.planner.plan_boolean(clauses)
        # NOT should come last
        self.assertEqual(plan.steps[-1].operation, "negate")
        self.assertEqual(plan.steps[-1].term, "noise")

    def test_not_only(self):
        clauses = [
            {"op": "NOT", "term": "noise"},
        ]
        plan = self.planner.plan_boolean(clauses)
        self.assertEqual(len(plan.steps), 1)
        self.assertEqual(plan.steps[0].operation, "negate")
        self.assertEqual(plan.steps[0].term, "noise")

    def test_all_three_ops(self):
        clauses = [
            {"op": "AND", "term": "rare"},
            {"op": "OR", "term": "medium"},
            {"op": "NOT", "term": "noise"},
        ]
        plan = self.planner.plan_boolean(clauses)
        self.assertEqual(plan.strategy, "boolean_mixed")
        ops = [s.operation for s in plan.steps]
        # lookup (AND first term), then union (OR), then negate (NOT)
        self.assertEqual(ops, ["lookup", "union", "negate"])

    def test_default_op_is_and(self):
        clauses = [{"term": "rare"}, {"term": "common"}]
        plan = self.planner.plan_boolean(clauses)
        self.assertEqual(plan.steps[0].term, "rare")
        self.assertEqual(plan.steps[0].operation, "lookup")
        self.assertEqual(plan.steps[1].operation, "intersect")

    def test_empty_term_skipped(self):
        clauses = [
            {"op": "AND", "term": ""},
            {"op": "AND", "term": "rare"},
        ]
        plan = self.planner.plan_boolean(clauses)
        self.assertEqual(len(plan.steps), 1)
        self.assertEqual(plan.steps[0].term, "rare")

    def test_boolean_cost_positive(self):
        clauses = [
            {"op": "AND", "term": "common"},
            {"op": "AND", "term": "rare"},
            {"op": "OR", "term": "medium"},
            {"op": "NOT", "term": "noise"},
        ]
        plan = self.planner.plan_boolean(clauses)
        self.assertGreater(plan.estimated_cost, 0.0)


# ---------------------------------------------------------------------------
# explain()
# ---------------------------------------------------------------------------

class TestExplain(unittest.TestCase):
    """Test human-readable plan explanation."""

    def setUp(self):
        self.planner = QueryPlanner(
            doc_freqs={"rare": 5, "common": 500},
            total_docs=1000,
        )

    def test_explain_contains_strategy(self):
        plan = self.planner.plan(["rare", "common"], "AND")
        explanation = self.planner.explain(plan)
        self.assertIn("and_rarest_first", explanation)

    def test_explain_contains_cost(self):
        plan = self.planner.plan(["rare", "common"], "AND")
        explanation = self.planner.explain(plan)
        self.assertIn("Estimated total cost:", explanation)

    def test_explain_contains_steps(self):
        plan = self.planner.plan(["rare", "common"], "AND")
        explanation = self.planner.explain(plan)
        self.assertIn("lookup", explanation)
        self.assertIn("intersect", explanation)
        self.assertIn("'rare'", explanation)
        self.assertIn("'common'", explanation)

    def test_explain_empty_plan(self):
        plan = self.planner.plan([])
        explanation = self.planner.explain(plan)
        self.assertIn("empty", explanation)
        self.assertIn("Steps (0):", explanation)

    def test_explain_line_count(self):
        plan = self.planner.plan(["rare", "common"], "AND")
        explanation = self.planner.explain(plan)
        lines = explanation.strip().split("\n")
        # Header (strategy, cost, "Steps (N):") + one line per step
        self.assertEqual(len(lines), 3 + len(plan.steps))

    def test_explain_or_plan(self):
        plan = self.planner.plan(["rare", "common"], "OR")
        explanation = self.planner.explain(plan)
        self.assertIn("or_largest_first", explanation)
        self.assertIn("union", explanation)


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------

class TestPlannerEdgeCases(unittest.TestCase):
    """Miscellaneous edge cases."""

    def test_unknown_terms_have_zero_df(self):
        planner = QueryPlanner(doc_freqs={"known": 10}, total_docs=100)
        plan = planner.plan(["unknown1", "unknown2"], "AND")
        for step in plan.steps:
            self.assertEqual(step.estimated_docs, 0)

    def test_all_same_frequency(self):
        planner = QueryPlanner(
            doc_freqs={"a": 50, "b": 50, "c": 50},
            total_docs=200,
        )
        plan = planner.plan(["a", "b", "c"], "AND")
        self.assertEqual(plan.strategy, "and_rarest_first")
        self.assertEqual(len(plan.steps), 3)

    def test_unknown_query_type_defaults_to_and(self):
        planner = QueryPlanner(doc_freqs={"x": 10, "y": 20})
        plan = planner.plan(["x", "y"], "XNOR")
        self.assertEqual(plan.strategy, "and_rarest_first")

    def test_no_doc_freqs(self):
        planner = QueryPlanner()
        plan = planner.plan(["foo", "bar"], "AND")
        self.assertEqual(plan.strategy, "and_rarest_first")
        self.assertAlmostEqual(plan.estimated_cost, 0.0)

    def test_single_term_or(self):
        planner = QueryPlanner(doc_freqs={"x": 42})
        plan = planner.plan(["x"], "OR")
        self.assertEqual(plan.strategy, "single_term")
        self.assertAlmostEqual(plan.estimated_cost, 42.0)

    def test_large_posting_lists(self):
        planner = QueryPlanner(
            doc_freqs={"big": 1_000_000, "small": 10},
            total_docs=2_000_000,
        )
        plan = planner.plan(["big", "small"], "AND")
        self.assertEqual(plan.steps[0].term, "small")
        self.assertLess(plan.estimated_cost, 1_000_000)

    def test_plan_boolean_or_only_strategy(self):
        planner = QueryPlanner(doc_freqs={"a": 10, "b": 20})
        plan = planner.plan_boolean([
            {"op": "OR", "term": "a"},
            {"op": "OR", "term": "b"},
        ])
        self.assertEqual(plan.strategy, "or_largest_first")

    def test_plan_boolean_multiple_not(self):
        planner = QueryPlanner(
            doc_freqs={"keep": 100, "noise1": 50, "noise2": 30},
            total_docs=200,
        )
        clauses = [
            {"op": "AND", "term": "keep"},
            {"op": "NOT", "term": "noise1"},
            {"op": "NOT", "term": "noise2"},
        ]
        plan = planner.plan_boolean(clauses)
        not_steps = [s for s in plan.steps if s.operation == "negate"]
        self.assertEqual(len(not_steps), 2)


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