"""
Test tuple format compatibility for prepare.py
"""

from samplers import DatetimeSampler, FloatSampler, IntSampler, TextSampler
from samplers.tidb_types import TiDBNumericBucket, TiDBStringBucket, TiDBTimeBucket


def validate_tuple_format():
    """Test that samplers work with prepare.py tuple format"""
    print("Testing tuple format compatibility...")

    # Test the format using TiDB types (converted from prepare.py lines 71-76)
    buckets_tidb_format = [
        TiDBNumericBucket(
            data_type="INT",
            lower_bound=1,
            upper_bound=10,
            count=100,
            repeats=10,
            ndv=50,
        ),
        TiDBNumericBucket(
            data_type="INT",
            lower_bound=10,
            upper_bound=50,
            count=200,
            repeats=20,
            ndv=100,
        ),
        TiDBNumericBucket(
            data_type="INT",
            lower_bound=50,
            upper_bound=100,
            count=150,
            repeats=15,
            ndv=75,
        ),
    ]

    try:
        # Test IntSampler with TiDB format
        int_sampler = IntSampler(buckets_tidb_format, is_cumulative_count=False)
        int_samples = int_sampler.sample_batch(10)
        print(f"✅ IntSampler with TiDB types: {len(int_samples)} samples generated")
        print(f"   Sample values: {[s for s in int_samples if s is not None][:5]}")

        # Test FloatSampler with TiDB format
        float_buckets = [
            TiDBNumericBucket(
                data_type="FLOAT",
                lower_bound=1.0,
                upper_bound=10.0,
                count=100,
                repeats=10,
                ndv=50,
            ),
            TiDBNumericBucket(
                data_type="FLOAT",
                lower_bound=10.0,
                upper_bound=50.0,
                count=200,
                repeats=20,
                ndv=100,
            ),
            TiDBNumericBucket(
                data_type="FLOAT",
                lower_bound=50.0,
                upper_bound=100.0,
                count=150,
                repeats=15,
                ndv=75,
            ),
        ]
        float_sampler = FloatSampler(float_buckets, is_cumulative_count=False)
        float_samples = float_sampler.sample_batch(10)
        print(
            f"✅ FloatSampler with TiDB types: {len(float_samples)} samples generated"
        )
        print(f"   Sample values: {[s for s in float_samples if s is not None][:5]}")

        # Test DatetimeSampler with TiDB format (should use datetime bounds)
        datetime_buckets_tidb = [
            TiDBTimeBucket(
                data_type="DATETIME",
                lower_bound="2024-01-01 00:00:00",
                upper_bound="2024-06-30 23:59:59",
                count=100,
                repeats=10,
                ndv=50,
            ),
            TiDBTimeBucket(
                data_type="DATETIME",
                lower_bound=b"2024-07-01 00:00:00",
                upper_bound=b"2024-12-31 23:59:59",
                count=200,
                repeats=20,
                ndv=100,
            ),
        ]
        datetime_sampler = DatetimeSampler(
            datetime_buckets_tidb, is_cumulative_count=False
        )
        datetime_samples = datetime_sampler.sample_batch(10)
        print(
            f"✅ DatetimeSampler with TiDB types: {len(datetime_samples)} samples generated"
        )
        print(f"   Sample values: {[s for s in datetime_samples if s is not None][:3]}")

        # Test TextSampler with TiDB format (should use string bounds)
        text_buckets_tidb = [
            TiDBStringBucket(
                data_type="VARCHAR",
                lower_bound="apple",
                upper_bound="banana",
                count=100,
                repeats=10,
                ndv=50,
                bucket_id=1,
                column_name="1",
            ),
            TiDBStringBucket(
                data_type="VARCHAR",
                lower_bound=b"banana",
                upper_bound=b"orange",
                count=200,
                repeats=20,
                ndv=100,
                bucket_id=2,
                column_name="2",
            ),
            TiDBStringBucket(
                data_type="VARCHAR",
                lower_bound="orange",
                upper_bound="zebra",
                count=150,
                repeats=15,
                ndv=75,
                bucket_id=3,
                column_name="3",
            ),
        ]
        text_sampler = TextSampler(text_buckets_tidb, is_cumulative_count=False)
        text_samples = text_sampler.sample_batch(5)
        print(f"✅ TextSampler with TiDB types: {len(text_samples)} samples generated")
        print(f"   Sample values: {[s for s in text_samples if s is not None][:3]}")

        return True

    except Exception as e:
        print(f"❌ Tuple format test failed: {e}")
        import traceback

        traceback.print_exc()
        return False


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