"""
Test text sampler compatibility with prepare.py format
"""

from samplers import TextSampler
from samplers.tidb_types import TiDBStringBucket


def validate_text_sampler_with_bytes():
    """Test TextSampler with byte string format from prepare.py"""
    print("Testing TextSampler with byte string format...")

    # Test text bucket format with byte strings using TiDB types (like in prepare.py)
    text_buckets_tidb_format = [
        TiDBStringBucket(
            data_type="VARCHAR",
            lower_bound=b"0004hbc6BXsDWbYI",
            upper_bound=b"0009abc7CYtEXcZJ",
            count=50,
            repeats=5,
            ndv=25,
            bucket_id=1,
            column_name="1",
        ),
        TiDBStringBucket(
            data_type="VARCHAR",
            lower_bound="hello",
            upper_bound="world",
            count=100,
            repeats=10,
            ndv=50,
            bucket_id=2,
            column_name="2",
        ),
        TiDBStringBucket(
            data_type="VARCHAR",
            lower_bound=b"test",
            upper_bound=b"zzz",
            count=75,
            repeats=8,
            ndv=38,
            bucket_id=3,
            column_name="3",
        ),
    ]

    try:
        text_sampler = TextSampler(text_buckets_tidb_format)
        text_samples = text_sampler.sample_batch(10)
        print(
            f"✅ TextSampler with byte strings: {len(text_samples)} samples generated"
        )
        print(f"   Sample values: {[s for s in text_samples if s is not None][:5]}")
        return True

    except Exception as e:
        print(f"❌ TextSampler byte string test failed: {e}")
        import traceback

        traceback.print_exc()
        return False


if __name__ == "__main__":
    success = validate_text_sampler_with_bytes()
    print(f"\nTextSampler compatibility test: {'PASSED' if success else 'FAILED'}")
    exit(0 if success else 1)