"""Tests for key-range partitioning (_shard_key) and partition-aware dry runs."""
from __future__ import annotations

from collections import Counter

from universe_reads import run
from universe_reads.parallel_read import _shard_key


# ---------------------------------------------------------------------------
# Unit tests for _shard_key
# ---------------------------------------------------------------------------

def test_shard_key_prefix_deterministic():
    """Same key always maps to the same shard."""
    for _ in range(10):
        assert _shard_key("ATT-CS101-20260310-0001", 5, "prefix") == \
               _shard_key("ATT-CS101-20260310-0001", 5, "prefix")


def test_shard_key_hash_deterministic():
    for _ in range(10):
        assert _shard_key("ATT-CS101-20260310-0001", 5, "hash") == \
               _shard_key("ATT-CS101-20260310-0001", 5, "hash")


def test_shard_key_prefix_range():
    """All shards should be in [0, sessions)."""
    sessions = 5
    for c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789":
        shard = _shard_key(c + "-rest", sessions, "prefix")
        assert 0 <= shard < sessions, f"shard {shard} out of range for key starting with {c}"


def test_shard_key_hash_range():
    sessions = 5
    for i in range(100):
        shard = _shard_key(f"KEY-{i:06d}", sessions, "hash")
        assert 0 <= shard < sessions


def test_shard_key_empty_key_prefix():
    """Empty key should not crash."""
    shard = _shard_key("", 5, "prefix")
    assert 0 <= shard < 5


def test_shard_key_hash_distribution():
    """Hash partitioning should produce a reasonably even distribution."""
    sessions = 5
    counts: Counter[int] = Counter()
    n = 5000
    for i in range(n):
        shard = _shard_key(f"KEY-{i:09d}", sessions, "hash")
        counts[shard] += 1

    # Each shard should get ~20% ± 10%
    for s in range(sessions):
        pct = counts[s] / n
        assert 0.10 < pct < 0.30, f"shard {s} got {pct:.1%} — too skewed"


def test_shard_key_prefix_groups_similar_keys():
    """Keys starting with the same character should land on the same shard."""
    sessions = 5
    shard_a1 = _shard_key("A-001", sessions, "prefix")
    shard_a2 = _shard_key("A-999", sessions, "prefix")
    assert shard_a1 == shard_a2, "Same-prefix keys should share a shard"


# ---------------------------------------------------------------------------
# Integration: dry-run with partition modes
# ---------------------------------------------------------------------------

def test_dry_run_partition_hash():
    """All records should be returned regardless of partition strategy."""
    n = 500
    results = run(
        file_name="",
        dry_run=n,
        sessions=3,
        count_only=True,
        partition="hash",
    )
    total = sum(r.records for r in results)
    assert total == n, f"Expected {n} records, got {total}"


def test_dry_run_partition_prefix():
    n = 500
    results = run(
        file_name="",
        dry_run=n,
        sessions=3,
        count_only=True,
        partition="prefix",
    )
    total = sum(r.records for r in results)
    assert total == n, f"Expected {n} records, got {total}"


def test_dry_run_no_partition_default():
    """Default (None) partition should still work (round-robin modulo)."""
    n = 500
    results = run(
        file_name="",
        dry_run=n,
        sessions=3,
        count_only=True,
        partition=None,
    )
    total = sum(r.records for r in results)
    assert total == n, f"Expected {n} records, got {total}"


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    test_shard_key_prefix_deterministic()
    print("PASS: _shard_key prefix deterministic")

    test_shard_key_hash_deterministic()
    print("PASS: _shard_key hash deterministic")

    test_shard_key_prefix_range()
    print("PASS: _shard_key prefix range")

    test_shard_key_hash_range()
    print("PASS: _shard_key hash range")

    test_shard_key_empty_key_prefix()
    print("PASS: _shard_key empty key prefix")

    test_shard_key_hash_distribution()
    print("PASS: _shard_key hash distribution")

    test_shard_key_prefix_groups_similar_keys()
    print("PASS: _shard_key prefix groups similar keys")

    test_dry_run_partition_hash()
    print("PASS: dry-run partition=hash")

    test_dry_run_partition_prefix()
    print("PASS: dry-run partition=prefix")

    test_dry_run_no_partition_default()
    print("PASS: dry-run partition=None (default)")

    print("\nAll partitioning tests passed.")


if __name__ == "__main__":
    main()
