"""Tests for ``shard_for_document`` (#000065).

The routing helper picks a target shard index from a content hash.
Identical answer on every peer; uniform distribution across shards.
"""
from __future__ import annotations

import hashlib

import pytest

from arborist.document import shard_for_document


def _root(text: str) -> str:
    """Return a synthetic hex sha256 to use as a document_root."""
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


class TestDeterminism:
    def test_same_input_same_output(self):
        root = _root("the rain in spain")
        assert shard_for_document(root, 4) == shard_for_document(root, 4)

    def test_independent_of_prior_calls(self):
        a = _root("a"), 4
        b = _root("b"), 4
        first = (shard_for_document(*a), shard_for_document(*b))
        second = (shard_for_document(*a), shard_for_document(*b))
        assert first == second


class TestBounds:
    @pytest.mark.parametrize("M", [1, 2, 4, 8, 16, 64])
    def test_within_range(self, M):
        for i in range(200):
            root = _root(f"doc-{i}")
            idx = shard_for_document(root, M)
            assert 0 <= idx < M

    def test_M_equals_one_always_zero(self):
        for i in range(50):
            assert shard_for_document(_root(f"doc-{i}"), 1) == 0

    def test_M_zero_rejected(self):
        with pytest.raises(ValueError):
            shard_for_document(_root("x"), 0)

    def test_M_negative_rejected(self):
        with pytest.raises(ValueError):
            shard_for_document(_root("x"), -1)


class TestUniformity:
    """First-32-bit prefix of SHA-256 is uniform; check that distribution
    is near-uniform across shards at meaningful sample sizes."""

    @pytest.mark.parametrize("M", [2, 4, 8])
    def test_distribution_within_5pct(self, M):
        N = 20_000
        counts = [0] * M
        for i in range(N):
            counts[shard_for_document(_root(f"doc-{i}"), M)] += 1
        expected = N / M
        for i, c in enumerate(counts):
            ratio = c / expected
            assert 0.95 <= ratio <= 1.05, (
                f"shard {i}: {c} hits, expected ~{expected:.0f} "
                f"({ratio:.3f}× — outside 5%)"
            )

    def test_M_changes_partition(self):
        """When M doubles, docs can move shard. Confirm at least one
        sampled doc lands in different shards under different M values."""
        roots = [_root(s) for s in ("YouTube", "SQLite", "arborist", "Wikipedia")]
        any_varied = False
        for r in roots:
            seen = {shard_for_document(r, m) for m in (2, 4, 8, 16)}
            if len(seen) >= 2:
                any_varied = True
                break
        assert any_varied, "no sampled root produced different shards under varying M"


class TestKnownRoots:
    """Lock-in fixtures: if these change, peers will disagree."""

    @pytest.mark.parametrize(
        "root,M,expected",
        [
            ("00000000" + "0" * 56, 4, 0),
            ("ffffffff" + "0" * 56, 4, 3),
            ("00000004" + "0" * 56, 4, 0),
            ("00000003" + "0" * 56, 4, 3),
            ("12345678" + "0" * 56, 4, 0),
            ("12345679" + "0" * 56, 4, 1),
            ("ffffffff" + "0" * 56, 1, 0),
        ],
    )
    def test_fixture(self, root, M, expected):
        assert shard_for_document(root, M) == expected