"""Tests for cross-pit search functionality."""

from pydantic import BaseModel

from mosh import Pit, search_by_type, search_pits
from mosh.search import search_tags_across_pits


class Person(BaseModel):
    """Test person class."""

    name: str
    age: int


class Product(BaseModel):
    """Test product class."""

    name: str
    price: float


class TestSearchPits:
    """Tests for search_pits function."""

    def test_search_across_multiple_pits(self) -> None:
        """Test searching across multiple pits."""
        with Pit("pit1", object_types=[Person]) as pit1:
            with Pit("pit2", object_types=[Person]) as pit2:
                pit1.put(Person(name="Alice", age=30))
                pit1.put(Person(name="Bob", age=25))
                pit2.put(Person(name="Charlie", age=30))
                pit2.put(Person(name="Diana", age=35))

                results = search_pits([pit1, pit2], age=30)

                assert len(results) == 2
                names = {r.obj.name for r in results}
                assert names == {"Alice", "Charlie"}

    def test_search_empty_pits(self) -> None:
        """Test searching across empty pits."""
        with Pit("empty1") as pit1:
            with Pit("empty2") as pit2:
                results = search_pits([pit1, pit2], name="Nobody")
                assert len(results) == 0

    def test_search_single_pit(self) -> None:
        """Test searching single pit through search_pits."""
        with Pit("single", object_types=[Person]) as pit:
            pit.put(Person(name="Alice", age=30))

            results = search_pits([pit], name="Alice")

            assert len(results) == 1
            assert results[0].obj.name == "Alice"

    def test_search_no_criteria(self) -> None:
        """Test searching with no criteria returns all."""
        with Pit("pit1", object_types=[Person]) as pit1:
            with Pit("pit2", object_types=[Person]) as pit2:
                pit1.put(Person(name="A", age=1))
                pit2.put(Person(name="B", age=2))
                pit2.put(Person(name="C", age=3))

                results = search_pits([pit1, pit2])

                assert len(results) == 3


class TestSearchByType:
    """Tests for search_by_type function."""

    def test_search_by_type_single_pit(self) -> None:
        """Test searching by type in single pit."""
        with Pit("mixed", object_types=[Person, Product]) as pit:
            pit.put(Person(name="Alice", age=30))
            pit.put(Product(name="Widget", price=10.0))
            pit.put(Person(name="Bob", age=25))

            results = search_by_type([pit], Person)

            assert len(results) == 2
            assert all(isinstance(r.obj, Person) for r in results)

    def test_search_by_type_multiple_pits(self) -> None:
        """Test searching by type across multiple pits."""
        with Pit("pit1", object_types=[Person, Product]) as pit1:
            with Pit("pit2", object_types=[Person, Product]) as pit2:
                pit1.put(Person(name="Alice", age=30))
                pit1.put(Product(name="Widget", price=10.0))
                pit2.put(Person(name="Bob", age=25))
                pit2.put(Product(name="Gadget", price=20.0))

                people = search_by_type([pit1, pit2], Person)
                products = search_by_type([pit1, pit2], Product)

                assert len(people) == 2
                assert len(products) == 2

    def test_search_by_type_no_matches(self) -> None:
        """Test searching by type with no matches."""
        with Pit("products", object_types=[Product]) as pit:
            pit.put(Product(name="Widget", price=10.0))

            people = search_by_type([pit], Person)

            assert len(people) == 0


class TestSearchTagsAcrossPits:
    """Tests for search_tags_across_pits function."""

    def test_search_tags_multiple_pits(self) -> None:
        """Test searching by tags across multiple pits."""
        with Pit("pit1", object_types=[Person]) as pit1:
            with Pit("pit2", object_types=[Person]) as pit2:
                pit1.put(Person(name="A", age=1), tags=["important"])
                pit1.put(Person(name="B", age=2), tags=["normal"])
                pit2.put(Person(name="C", age=3), tags=["important", "urgent"])

                results = search_tags_across_pits([pit1, pit2], ["important"])

                assert len(results) == 2

    def test_search_tags_match_all(self) -> None:
        """Test searching for all tags across pits."""
        with Pit("pit1", object_types=[Person]) as pit1:
            with Pit("pit2", object_types=[Person]) as pit2:
                pit1.put(Person(name="A", age=1), tags=["a", "b"])
                pit1.put(Person(name="B", age=2), tags=["a"])
                pit2.put(Person(name="C", age=3), tags=["a", "b", "c"])

                results = search_tags_across_pits([pit1, pit2], ["a", "b"], match_all=True)

                assert len(results) == 2

    def test_search_tags_match_any(self) -> None:
        """Test searching for any tags across pits."""
        with Pit("pit1", object_types=[Person]) as pit1:
            with Pit("pit2", object_types=[Person]) as pit2:
                pit1.put(Person(name="A", age=1), tags=["a"])
                pit2.put(Person(name="B", age=2), tags=["b"])
                pit2.put(Person(name="C", age=3), tags=["c"])

                results = search_tags_across_pits([pit1, pit2], ["a", "b"], match_all=False)

                assert len(results) == 2
