from __future__ import annotations

from datetime import datetime

from streamer.localize import LocalizedShard
from streamer.query.find import find_shard, find_shard_by_position


def generate_localized_shard(
    *,
    location: dict[str, str] | None = None,
    children: list[LocalizedShard] | None = None,
) -> LocalizedShard:
    return LocalizedShard(
        start_line=1,
        end_line=1,
        moment=datetime(2020, 1, 1),
        location=location or {},
        children=children or [],
        markers=[],
        tags=[],
    )


class TestFindShard:
    def test_returns_empty_when_no_match(self) -> None:
        root = generate_localized_shard(location={"file": "a.md"})
        shards = [root]

        result = find_shard(shards, lambda s: "missing" in s.location)

        assert result == []

    def test_finds_matches_depth_first_and_preserves_order(self) -> None:
        grandchild = generate_localized_shard(location={"k": "match"})
        child1 = generate_localized_shard(
            location={"k": "match"}, children=[grandchild]
        )
        child2 = generate_localized_shard(location={"k": "nope"})
        root = generate_localized_shard(
            location={"k": "nope"}, children=[child1, child2]
        )

        result = find_shard([root], lambda s: s.location.get("k") == "match")

        assert result == [child1, grandchild]

    def test_includes_root_if_it_matches(self) -> None:
        root = generate_localized_shard(
            location={"k": "match"},
            children=[generate_localized_shard(location={"k": "match"})],
        )

        result = find_shard([root], lambda s: s.location.get("k") == "match")

        assert result[0] is root
        assert len(result) == 2

    def test_multiple_roots_keeps_left_to_right_order(self) -> None:
        a = generate_localized_shard(location={"k": "match"})
        b = generate_localized_shard(location={"k": "match"})
        c = generate_localized_shard(location={"k": "nope"})

        result = find_shard([a, b, c], lambda s: s.location.get("k") == "match")

        assert result == [a, b]

    def test_query_function_can_use_arbitrary_logic(self) -> None:
        # Ensures typing/behavior supports any callable that returns bool.
        a = generate_localized_shard(location={"x": "1"})
        b = generate_localized_shard(location={"x": "2"})
        c = generate_localized_shard(location={"x": "3"})
        root = generate_localized_shard(location={}, children=[a, b, c])

        def is_even_x(shard: LocalizedShard) -> bool:
            x = shard.location.get("x")
            return x is not None and int(x) % 2 == 0

        result = find_shard([root], is_even_x)

        assert result == [b]


class TestFindShardByPosition:
    def test_matches_only_when_dimension_present_and_equal(self) -> None:
        match = generate_localized_shard(location={"file": "a.md", "line": "10"})
        wrong_value = generate_localized_shard(location={"file": "a.md", "line": "11"})
        missing_dim = generate_localized_shard(location={"file": "a.md"})
        root = generate_localized_shard(
            location={"root": "x"}, children=[match, wrong_value, missing_dim]
        )

        result = find_shard_by_position([root], "line", "10")

        assert result == [match]

    def test_recurses_through_children(self) -> None:
        deep = generate_localized_shard(location={"section": "s1"})
        mid = generate_localized_shard(location={"section": "s0"}, children=[deep])
        root = generate_localized_shard(location={}, children=[mid])

        result = find_shard_by_position([root], "section", "s1")

        assert result == [deep]