"""
Test suite for TextSampler with visualization and distribution comparison
"""

import os
import random
from collections import Counter
from typing import List, Optional

import matplotlib.pyplot as plt
import numpy as np

from samplers import HistogramBucket, InverseTransformSampler, TextSampler, TopNItem
from samplers.tidb_types import TiDBStringBucket, TiDBStringTopNItem
from utils.bytes_scalar_converter import common_prefix_length
from utils.scalar_to_string import convert_bytes_to_scalar, scalar_to_string


def get_validate_seed():
    """Get seed from environment or use default"""
    return int(os.environ.get("ITS_VALIDATION_SEED", 42))


def generate_validate_string_data(
    size: int = 1000, seed: Optional[int] = None
) -> List[Optional[str]]:
    """Generate test string data with different patterns based on seed"""
    if seed is None:
        seed = get_validate_seed()
    random.seed(seed)
    np.random.seed(seed)

    # Use seed to create different distribution patterns
    pattern = seed % 4  # Create 4 different distribution patterns
    data = []

    if pattern == 0:  # Short strings distribution
        print(f"  🎲 Generating short-string distribution (seed {seed})")
        common_strings = ["a", "b", "ab", "abc", "test"]
        # Focus on very short strings
        for _ in range(size // 2):
            data.append(random.choice(common_strings))
        for _ in range(size // 3):
            length = random.randint(1, 4)  # Very short
            chars = "abcdefghijklmnopqrstuvwxyz"
            data.append("".join(random.choice(chars) for _ in range(length)))

    elif pattern == 1:  # Names distribution
        print(f"  🎲 Generating names distribution (seed {seed})")
        common_strings = ["Alice", "Bob", "Charlie", "Diana", "Edward"]
        # Focus on person names
        names = [
            "Alice",
            "Bob",
            "Charlie",
            "David",
            "Emma",
            "Frank",
            "Grace",
            "Henry",
            "Isabella",
            "Jack",
            "Katherine",
            "Liam",
            "Mia",
            "Noah",
            "Olivia",
            "Peter",
            "Quinn",
            "Rachel",
            "Samuel",
            "Taylor",
            "Uma",
            "Victor",
            "Wendy",
            "Xavier",
        ]
        for _ in range(size // 2):
            data.append(random.choice(common_strings))
        for _ in range(size // 3):
            data.append(random.choice(names))

    elif pattern == 2:  # Long strings distribution
        print(f"  🎲 Generating long-string distribution (seed {seed})")
        common_strings = [
            "longteststring",
            "verylongstringdata",
            "thisisverylong",
            "muchlongertext",
            "extralongstring",
        ]
        # Focus on longer strings
        for _ in range(size // 2):
            data.append(random.choice(common_strings))
        for _ in range(size // 3):
            length = random.randint(10, 25)  # Long strings
            chars = "abcdefghijklmnopqrstuvwxyz0123456789"
            data.append("".join(random.choice(chars) for _ in range(length)))

    else:  # Mixed pattern (emails, URLs, etc.)
        print(f"  🎲 Generating mixed-format string distribution (seed {seed})")
        common_strings = [
            "user@test.com",
            "http://example.com",
            "file.txt",
            "data.json",
            "config.yaml",
        ]
        # Mix of different string formats
        domains = ["gmail.com", "test.org", "example.net", "demo.co"]
        names = ["user", "admin", "test", "demo", "sample"]

        for _ in range(size // 4):
            data.append(random.choice(common_strings))
        for _ in range(size // 4):
            # Generate emails
            name = random.choice(names)
            domain = random.choice(domains)
            data.append(f"{name}@{domain}")
        for _ in range(size // 4):
            # Generate file names
            base = random.choice(["file", "data", "test", "config", "backup"])
            ext = random.choice(["txt", "json", "csv", "log", "xml"])
            num = random.randint(1, 999)
            data.append(f"{base}{num}.{ext}")

    # Add some nulls (consistent across patterns)
    null_count = size // 20
    for _ in range(null_count):
        data.append(None)

    # Fill remaining with more random data
    remaining = size - len(data)
    # Default values for variables that might not be set in all patterns
    if "names" not in locals():
        names = []
    if "chars" not in locals():
        chars = "abcdefghijklmnopqrstuvwxyz0123456789"

    for _ in range(remaining):
        if random.random() < 0.7 and (common_strings + names):
            data.append(random.choice(common_strings + names))
        else:
            length = random.randint(2, 10)
            random_str = "".join(random.choice(chars) for _ in range(length))
            data.append(random_str)

    # Shuffle the data
    random.shuffle(data)
    return data[:size]


def create_validate_text_sampler_from_data(string_data: List[Optional[str]], **kwargs):
    """Helper function to create TextSampler from string data using TiDB types for testing"""
    charset = kwargs.get("charset", "utf8mb4")
    num_buckets = kwargs.get("num_buckets", 10)
    top_n_size = kwargs.get("top_n_size", 5)
    seed = kwargs.get("seed", None)
    character_maximum_length = kwargs.get("character_maximum_length", None)

    # Process string data
    clean_strings = [s for s in string_data if s is not None]
    null_count = len(string_data) - len(clean_strings)

    if not clean_strings:
        return TextSampler([], [], null_count=null_count), {}

    # Find TopN most frequent strings
    string_counts = Counter(clean_strings)
    most_common_strings = string_counts.most_common(top_n_size)

    # Create TiDB string TopN items directly from strings
    tidb_top_n = []
    for string_val, count in most_common_strings:
        tidb_topn_item = TiDBStringTopNItem(
            data_type="VARCHAR",
            encoded=string_val,  # Store the original string
            count=count,
        )
        tidb_top_n.append(tidb_topn_item)

    # Remove TopN strings for histogram creation
    topn_strings = {s for s, _ in most_common_strings}
    histogram_strings = [s for s in clean_strings if s not in topn_strings]

    # Create TiDB string buckets directly from string ranges
    tidb_buckets = []
    if histogram_strings:
        # Sort strings for bucketing
        sorted_strings = sorted(set(histogram_strings))

        if len(sorted_strings) == 1:
            # All remaining strings are the same
            single_string = sorted_strings[0]
            bucket_count = len(histogram_strings)
            tidb_bucket = TiDBStringBucket(
                data_type="VARCHAR",
                lower_bound=single_string,
                upper_bound=single_string,
                count=bucket_count,
                repeats=bucket_count,
                ndv=1,
                character_maximum_length=character_maximum_length,
                bucket_id=1,
                column_name="1",
            )
            tidb_buckets.append(tidb_bucket)
        else:
            # Create buckets from string ranges
            bucket_size = max(1, len(sorted_strings) // num_buckets)

            for i in range(0, len(sorted_strings), bucket_size):
                bucket_strings = sorted_strings[i : i + bucket_size]
                if not bucket_strings:
                    continue

                lower_bound = bucket_strings[0]
                upper_bound = bucket_strings[-1]

                # Count occurrences in this bucket
                bucket_count = sum(histogram_strings.count(s) for s in bucket_strings)
                repeats = histogram_strings.count(upper_bound)
                ndv = len(bucket_strings)

                if bucket_count > 0:
                    tidb_bucket = TiDBStringBucket(
                        data_type="VARCHAR",
                        lower_bound=lower_bound,
                        upper_bound=upper_bound,
                        count=bucket_count,
                        repeats=repeats,
                        ndv=ndv,
                        character_maximum_length=character_maximum_length,
                        bucket_id=i,
                        column_name=str(i),
                    )
                    tidb_buckets.append(tidb_bucket)

    # Create topn mappings for TextSampler (scalar → string mapping)
    topn_mappings = {}
    for tidb_item in tidb_top_n:
        original_string = tidb_item.encoded
        # Convert string to scalar for mapping key
        scalar_value = convert_bytes_to_scalar(original_string.encode("utf-8"))
        topn_mappings[scalar_value] = original_string

    # Also create scalar_to_string mapping for return value
    scalar_to_string = {}
    for s in clean_strings:
        scalar = convert_bytes_to_scalar(s.encode("utf-8"))
        scalar_to_string[scalar] = s

    # Create TextSampler with TiDB types and absolute counts (not cumulative)
    return (
        TextSampler(
            buckets=tidb_buckets,
            top_n=tidb_top_n,
            null_count=null_count,
            charset=charset,
            topn_string_mappings=topn_mappings,
            is_cumulative_count=False,  # Our test data uses absolute counts
            seed=seed,
        ),
        scalar_to_string,
    )


# Remove the old _generate_tidb_histogram_from_scalars function since we're using direct string bounds now


def validate_text_sampler_basic():
    """Test basic TextSampler functionality"""
    print("=== Test Basic TextSampler Functionality ===")

    # Generate test data
    validate_data = generate_validate_string_data(1000, seed=get_validate_seed())

    print(f"Generated {len(validate_data)} test strings")
    print(f"Null count: {sum(1 for x in validate_data if x is None)}")
    print(f"Unique strings: {len(set(x for x in validate_data if x is not None))}")

    # Create TextSampler using test helper
    sampler, original_mappings = create_validate_text_sampler_from_data(
        validate_data,
        charset="utf8mb4",
        num_buckets=10,
        top_n_size=5,
        seed=get_validate_seed(),
    )

    # Get statistics from underlying ITS
    stats = sampler.its.get_statistics()
    print("\nTextSampler Statistics:")
    print(f"  Number of buckets: {stats['num_buckets']}")
    print(f"  Number of TopN: {stats['num_topn']}")
    print(f"  Total count: {stats['total_count']}")
    print(f"  Null probability: {stats['null_probability']:.3f}")
    print(f"  Original strings: {len(original_mappings)}")

    # Sample texts
    samples = sampler.sample_batch(100)
    non_null_samples = [s for s in samples if s is not None]

    print("\nSampling Results:")
    print(f"  Generated {len(samples)} samples")
    print(f"  Non-null samples: {len(non_null_samples)}")
    print(f"  Sample examples: {non_null_samples[:10]}")

    return validate_data, sampler, samples


def create_text_sampler_from_tidb_stats(
    buckets_data: List[
        tuple
    ],  # [(lower_bound_str, upper_bound_str, count, repeats), ...]
    topn_data: List[tuple],  # [(value_str, count), ...]
    null_count: int = 0,
    charset: str = "utf8mb4",
    seed: Optional[int] = None,
):
    """
    Helper function for testing: Create TextSampler directly from TiDB statistics format

    Args:
        buckets_data: List of (lower_bound, upper_bound, count, repeats) tuples
        topn_data: List of (value, count) tuples
        null_count: Number of null values
        charset: Character set for text generation
        seed: Random seed

    Returns:
        TextSampler-like instance
    """
    # Convert string bounds to scalar bounds
    buckets = []
    for lower_str, upper_str, count, repeats in buckets_data:
        lower_scalar = convert_bytes_to_scalar(str(lower_str).encode("utf-8"))
        upper_scalar = convert_bytes_to_scalar(str(upper_str).encode("utf-8"))

        bucket = HistogramBucket(
            lower_bound=lower_scalar,
            upper_bound=upper_scalar,
            count=count,
            repeats=repeats,
        )
        buckets.append(bucket)

    # Convert TopN strings to scalars
    top_n = []
    scalar_to_string_map = {}

    for value_str, count in topn_data:
        value_scalar = convert_bytes_to_scalar(str(value_str).encode("utf-8"))
        top_n_item = TopNItem(value=value_scalar, count=count)
        top_n.append(top_n_item)
        scalar_to_string_map[value_scalar] = str(value_str)

    # Create ITS
    its = InverseTransformSampler(
        buckets=buckets, top_n=top_n, null_count=null_count, seed=seed
    )

    # Create a simplified TextSampler-like class for testing
    class TestTextSampler:
        def __init__(self, its, scalar_to_string_map, charset):
            self.its = its
            self.scalar_to_string_map = scalar_to_string_map
            self.charset = charset

        def sample(self) -> Optional[str]:
            sample_result = self.its.sample()
            scalar_value = sample_result.value

            if scalar_value is None:
                return None

            # Check if this is a known TopN value
            if scalar_value in self.scalar_to_string_map:
                return self.scalar_to_string_map[scalar_value]

            # Generate new text using scalar_to_string
            margin = abs(scalar_value) * 0.01 + 1000
            result, _ = scalar_to_string(
                scalar_value,
                scalar_value - margin,
                scalar_value + margin,
                charset=charset,
            )

            return (
                result
                if result is not None
                else f"gen_{abs(hash(str(scalar_value))) % 10000}"
            )

        def sample_batch(self, n: int) -> List[Optional[str]]:
            return [self.sample() for _ in range(n)]

        def get_statistics(self) -> dict:
            return {
                **self.its.get_statistics(),
                "charset": self.charset,
                "topn_mappings": len(self.scalar_to_string_map),
            }

    return TestTextSampler(its, scalar_to_string_map, charset)


def validate_tidb_stats_sampler():
    """Test create_text_sampler_from_tidb_stats function"""
    print("\n=== Test TiDB Stats TextSampler ===")

    # Mock TiDB statistics
    buckets_data = [
        ("a", "hello", 100, 10),
        ("hello", "python", 150, 15),
        ("python", "test", 120, 12),
        ("test", "world", 80, 8),
        ("world", "z", 50, 5),
    ]

    topn_data = [
        ("hello", 50),
        ("world", 40),
        ("python", 35),
        ("test", 30),
        ("data", 25),
    ]

    # Create sampler from TiDB stats
    tidb_sampler = create_text_sampler_from_tidb_stats(
        buckets_data=buckets_data,
        topn_data=topn_data,
        null_count=20,
        charset="utf8mb4",
        seed=get_validate_seed(),
    )

    # Get statistics
    stats = tidb_sampler.get_statistics()
    print("TiDB Sampler Statistics:")
    print(f"  Number of buckets: {stats['num_buckets']}")
    print(f"  Number of TopN: {stats['num_topn']}")
    print(f"  Total count: {stats['total_count']}")
    print(f"  TopN mappings: {stats['topn_mappings']}")

    # Sample texts
    samples = tidb_sampler.sample_batch(50)
    non_null_samples = [s for s in samples if s is not None]

    print("\nSampling Results:")
    print(f"  Generated {len(samples)} samples")
    print(f"  Non-null samples: {len(non_null_samples)}")
    print(f"  Sample examples: {non_null_samples[:15]}")

    # Count TopN occurrences
    topn_values = {"hello", "world", "python", "test", "data"}
    topn_count = sum(1 for s in non_null_samples if s in topn_values)
    print(
        f"  TopN values in samples: {topn_count}/{len(non_null_samples)} ({topn_count / len(non_null_samples) * 100:.1f}%)"
    )

    return tidb_sampler, samples


def visualize_text_distributions(
    original_data: List[Optional[str]],
    generated_samples: List[Optional[str]],
    sampler: TextSampler,
    output_dir: str = "tests",
):
    """Create visualization comparing text distributions"""
    print("\n=== Creating Text Distribution Visualizations ===")

    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)

    # Filter out nulls
    original_clean = [x for x in original_data if x is not None]
    generated_clean = [x for x in generated_samples if x is not None]

    # Convert to scalars for comparison
    original_scalars = [
        convert_bytes_to_scalar(s.encode("utf-8")) for s in original_clean
    ]
    generated_scalars = [
        convert_bytes_to_scalar(s.encode("utf-8"))
        for s in generated_clean
        if isinstance(s, str)
    ]

    # Create figure with subplots
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle("TextSampler - Distribution Comparison", fontsize=16)

    # 1. Scalar distribution comparison
    ax1 = axes[0, 0]
    bins = np.linspace(
        min(min(original_scalars), min(generated_scalars)),
        max(max(original_scalars), max(generated_scalars)),
        50,
    )

    ax1.hist(
        original_scalars,
        bins=bins,
        alpha=0.7,
        label="Original",
        density=True,
        color="blue",
    )
    ax1.hist(
        generated_scalars,
        bins=bins,
        alpha=0.7,
        label="Generated",
        density=True,
        color="red",
    )
    ax1.set_xlabel("Scalar Value")
    ax1.set_ylabel("Density")
    ax1.set_title("Scalar Distribution Comparison")
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # 2. String length distribution
    ax2 = axes[0, 1]
    original_lengths = [len(s) for s in original_clean]
    generated_lengths = [len(s) for s in generated_clean if isinstance(s, str)]

    max_len = max(max(original_lengths), max(generated_lengths))
    length_bins = np.arange(0, max_len + 2) - 0.5

    ax2.hist(
        original_lengths,
        bins=length_bins,
        alpha=0.7,
        label="Original",
        density=True,
        color="blue",
    )
    ax2.hist(
        generated_lengths,
        bins=length_bins,
        alpha=0.7,
        label="Generated",
        density=True,
        color="red",
    )
    ax2.set_xlabel("String Length")
    ax2.set_ylabel("Density")
    ax2.set_title("String Length Distribution")
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    # 3. Top values comparison
    ax3 = axes[1, 0]
    original_counts = Counter(original_clean)
    generated_counts = Counter(generated_clean)

    # Get top 10 values from each
    top_original = dict(original_counts.most_common(10))
    top_generated = dict(generated_counts.most_common(10))

    # Combine all top values
    all_top_values = set(top_original.keys()) | set(top_generated.keys())

    if all_top_values:
        x_pos = np.arange(len(all_top_values))
        orig_counts = [top_original.get(val, 0) for val in all_top_values]
        gen_counts = [top_generated.get(val, 0) for val in all_top_values]

        width = 0.35
        ax3.bar(
            x_pos - width / 2,
            orig_counts,
            width,
            label="Original",
            alpha=0.7,
            color="blue",
        )
        ax3.bar(
            x_pos + width / 2,
            gen_counts,
            width,
            label="Generated",
            alpha=0.7,
            color="red",
        )

        ax3.set_xlabel("Top Values")
        ax3.set_ylabel("Count")
        ax3.set_title("Top Values Comparison")
        ax3.set_xticks(x_pos)
        ax3.set_xticklabels(list(all_top_values), rotation=45, ha="right")
        ax3.legend()
        ax3.grid(True, alpha=0.3)

    # 4. Sampler buckets visualization
    ax4 = axes[1, 1]
    its_stats = sampler.its.get_statistics()

    if its_stats["bucket_ranges"]:
        bucket_centers = []
        bucket_counts = []
        bucket_widths = []

        for lower, upper, count in its_stats["bucket_ranges"]:
            center = (lower + upper) / 2
            bucket_centers.append(center)
            bucket_counts.append(count)
            bucket_widths.append(abs(upper - lower))

        ax4.bar(
            bucket_centers,
            bucket_counts,
            width=bucket_widths,
            alpha=0.7,
            color="green",
            edgecolor="black",
        )
        ax4.set_xlabel("Scalar Value")
        ax4.set_ylabel("Count")
        ax4.set_title("Histogram Buckets")
        ax4.grid(True, alpha=0.3)

    plt.tight_layout()

    # Save the plot
    output_path = os.path.join(output_dir, "text_sampler_comparison.png")
    plt.savefig(output_path, dpi=300, bbox_inches="tight")
    print(f"Visualization saved to: {output_path}")

    # Print statistics
    print("\nDistribution Comparison:")
    print(f"  Original strings: {len(original_clean)}")
    print(f"  Generated strings: {len(generated_clean)}")
    print(f"  Original avg length: {np.mean(original_lengths):.2f}")
    print(f"  Generated avg length: {np.mean(generated_lengths):.2f}")
    print(
        f"  Original scalar range: {min(original_scalars):.0f} - {max(original_scalars):.0f}"
    )
    print(
        f"  Generated scalar range: {min(generated_scalars):.0f} - {max(generated_scalars):.0f}"
    )

    return output_path


def validate_charset_variations():
    """Test different charset options"""
    print("\n=== Test Different Charsets ===")

    # Generate test data with unicode characters
    validate_strings = [
        "hello",
        "world",
        "test",
        "你好",
        "世界",
        "café",
        "naïve",
        "résumé",
        "Москва",
        "Tokyo",
        "東京",
        "مرحبا",
        "🌟",
        "💻",
        "🚀",
    ]

    # Test with different charsets
    charsets = ["utf8mb4", "utf8"]

    for charset in charsets:
        print(f"\n--- Testing charset: {charset} ---")

        try:
            sampler, _ = create_validate_text_sampler_from_data(
                validate_strings * 10,  # Repeat for more data
                charset=charset,
                num_buckets=5,
                top_n_size=3,
                seed=get_validate_seed(),
            )

            samples = sampler.sample_batch(20)
            non_null_samples = [s for s in samples if s is not None]

            print(f"  Generated {len(non_null_samples)} samples")
            print(f"  Sample examples: {non_null_samples[:8]}")

        except Exception as e:
            print(f"  Error with {charset}: {e}")


def validate_common_prefix_scenarios():
    """Test common prefix handling scenarios (TiDB-style)"""
    print("\n=== Test Common Prefix Scenarios ===")

    # Test 1: Strings with common prefix
    print("\n--- Test 1: Common Prefix Strings ---")
    common_prefix_data = [
        "user_001",
        "user_002",
        "user_003",
        "user_004",
        "user_005",
        "user_010",
        "user_011",
        "user_012",
        "user_013",
        "user_014",
        "user_100",
        "user_101",
        "user_102",
        "user_103",
        "user_104",
        "user_admin",
        "user_test",
        "user_demo",
    ] * 10  # Repeat for more data

    try:
        prefix_sampler, prefix_mappings = create_validate_text_sampler_from_data(
            common_prefix_data,
            charset="utf8mb4",
            num_buckets=8,
            top_n_size=3,
            seed=get_validate_seed(),
        )

        prefix_samples = prefix_sampler.sample_batch(30)
        prefix_non_null = [s for s in prefix_samples if s is not None]

        print(f"  Generated {len(prefix_non_null)} samples")
        print(f"  Sample examples: {prefix_non_null[:15]}")

        # Check if generated samples preserve common pattern
        user_prefix_count = sum(1 for s in prefix_non_null if s.startswith("user_"))
        print(
            f"  Samples with 'user_' prefix: {user_prefix_count}/{len(prefix_non_null)} ({user_prefix_count / len(prefix_non_null) * 100:.1f}%)"
        )

    except Exception as e:
        print(f"  Common prefix test error: {e}")
        import traceback

        traceback.print_exc()

    # Test 2: URL-like strings with common prefix
    print("\n--- Test 2: URL-like Common Prefixes ---")
    url_data = [
        "https://example.com/page1",
        "https://example.com/page2",
        "https://example.com/page3",
        "https://example.com/user/1",
        "https://example.com/user/2",
        "https://example.com/user/3",
        "https://test.org/home",
        "https://test.org/about",
        "https://test.org/contact",
        "http://demo.net/api/v1",
        "http://demo.net/api/v2",
        "http://demo.net/docs",
    ] * 8

    try:
        url_sampler, url_mappings = create_validate_text_sampler_from_data(
            url_data,
            charset="utf8mb4",
            num_buckets=10,
            top_n_size=4,
            seed=get_validate_seed(),
        )

        url_samples = url_sampler.sample_batch(25)
        url_non_null = [s for s in url_samples if s is not None]

        print(f"  Generated {len(url_non_null)} samples")
        print(f"  Sample examples: {url_non_null[:12]}")

        # Check URL-like pattern preservation
        http_count = sum(
            1 for s in url_non_null if s.startswith(("http://", "https://"))
        )
        print(
            f"  Samples with http(s):// prefix: {http_count}/{len(url_non_null)} ({http_count / len(url_non_null) * 100:.1f}%)"
        )

    except Exception as e:
        print(f"  URL prefix test error: {e}")
        import traceback

        traceback.print_exc()

    # Test 3: Email-like strings with domain prefix
    print("\n--- Test 3: Email Domain Common Prefixes ---")
    email_data = [
        "john@company.com",
        "jane@company.com",
        "bob@company.com",
        "alice@company.com",
        "admin@company.com",
        "test@company.com",
        "demo@company.com",
        "support@company.com",
        "user1@test.org",
        "user2@test.org",
        "user3@test.org",
        "user4@test.org",
        "contact@demo.net",
        "info@demo.net",
        "help@demo.net",
        "sales@demo.net",
    ] * 6

    try:
        email_sampler, email_mappings = create_validate_text_sampler_from_data(
            email_data,
            charset="utf8mb4",
            num_buckets=12,
            top_n_size=5,
            seed=get_validate_seed(),
        )

        email_samples = email_sampler.sample_batch(20)
        email_non_null = [s for s in email_samples if s is not None]

        print(f"  Generated {len(email_non_null)} samples")
        print(f"  Sample examples: {email_non_null[:10]}")

        # Check email pattern preservation
        at_count = sum(1 for s in email_non_null if "@" in s)
        print(
            f"  Samples with '@' (email-like): {at_count}/{len(email_non_null)} ({at_count / len(email_non_null) * 100:.1f}%)"
        )

    except Exception as e:
        print(f"  Email prefix test error: {e}")
        import traceback

        traceback.print_exc()


def validate_common_prefix_accuracy():
    """Test the accuracy of common prefix handling implementation"""
    print("\n=== Test Common Prefix Accuracy ===")

    # Test the internal common prefix calculation directly
    from samplers import TextSampler

    # Create a dummy sampler to test internal methods
    dummy_buckets = []
    TextSampler(dummy_buckets, seed=42)

    # Test 1: Basic common prefix calculation
    print("\n--- Test 1: Common Prefix Length Calculation ---")
    validate_cases = [
        ([b"user_001", b"user_002"], 7),  # 'user_00' = 7 bytes
        ([b"hello", b"help"], 3),  # 'hel' = 3 bytes
        ([b"same", b"same"], 4),  # 'same' = 4 bytes
        ([b"a", b"b"], 0),  # no common prefix
        ([b"prefix_a", b"prefix_b"], 7),  # 'prefix_' = 7 bytes
        ([b"", b""], 0),  # empty string
        (
            ["你好呀".encode("utf-8"), "你好".encode("utf-8")],
            6,
        ),  # '你好' = 6 bytes in UTF-8
    ]

    for byte_strings, expected_len in validate_cases:
        try:
            calculated_len = common_prefix_length(byte_strings[0], byte_strings[1])
            status = "✅" if calculated_len == expected_len else "❌"
            print(
                f"  {status} {byte_strings} -> {calculated_len} (expected {expected_len})"
            )
        except Exception as e:
            print(f"  ❌ {byte_strings} -> Error: {e}")

    # Test 2: Scalar conversion consistency
    print("\n--- Test 2: Scalar Conversion Consistency ---")
    from utils.scalar_to_string import convert_bytes_to_scalar

    validate_strings = [b"hello", b"world", b"test", b"python", b"data"]

    for original_bytes in validate_strings:
        try:
            # Convert to scalar and check consistency
            scalar = convert_bytes_to_scalar(original_bytes)

            print(f"  Original: {original_bytes} -> Scalar: {scalar:.0f}")

            # Test if scalar can be used for sampling
            if scalar > 0:
                print("    ✅ Valid scalar for sampling")
            else:
                print("    ⚠️ Zero scalar value")

        except Exception as e:
            print(f"  ❌ {original_bytes} -> Error: {e}")

    # Test 3: Integration test with real data
    print("\n--- Test 3: Integration Test ---")
    prefix_data = [
        "prefix_" + str(i).zfill(3) for i in range(1, 21)
    ]  # prefix_001 to prefix_020

    try:
        prefix_sampler, _ = create_validate_text_sampler_from_data(
            prefix_data,
            charset="utf8mb4",
            num_buckets=5,
            top_n_size=2,
            seed=get_validate_seed(),
        )

        # Check if bucket common prefixes were computed (stored in TiDB buckets)
        if hasattr(prefix_sampler, "tidb_buckets") and prefix_sampler.tidb_buckets:
            print(
                f"  TiDB bucket common prefixes computed: {len(prefix_sampler.tidb_buckets)} buckets"
            )

            for bucket_idx, bucket in enumerate(prefix_sampler.tidb_buckets):
                if hasattr(bucket, "common_prefix_length"):
                    common_len = bucket.common_prefix_length
                    common_bytes = bucket.common_prefix_bytes
                    print(
                        f"    TiDB Bucket {bucket_idx}: prefix_len={common_len}, prefix_bytes={common_bytes}"
                    )
                else:
                    print(f"    TiDB Bucket {bucket_idx}: no common prefix info")
        else:
            print("  ❌ No TiDB buckets found")

        # Test sampling
        samples = prefix_sampler.sample_batch(15)
        non_null_samples = [s for s in samples if s is not None]

        print(f"  Generated {len(non_null_samples)} samples:")
        for i, sample in enumerate(non_null_samples[:10]):
            print(f"    Sample {i + 1}: '{sample}'")

        # Check prefix preservation
        prefix_preserved = sum(1 for s in non_null_samples if s.startswith("prefix_"))
        print(
            f"  Samples preserving 'prefix_': {prefix_preserved}/{len(non_null_samples)} ({prefix_preserved / len(non_null_samples) * 100:.1f}%)"
        )

    except Exception as e:
        print(f"  Integration test error: {e}")
        import traceback

        traceback.print_exc()


def validate_edge_cases():
    """Test edge cases"""
    print("\n=== Test Edge Cases ===")

    # Test with empty data
    try:
        empty_sampler, _ = create_validate_text_sampler_from_data([], seed=42)
        empty_samples = empty_sampler.sample_batch(5)
        print(f"Empty data samples: {empty_samples}")
    except Exception as e:
        print(f"Empty data error: {e}")

    # Test with only nulls
    try:
        null_sampler, _ = create_validate_text_sampler_from_data([None] * 10, seed=42)
        null_samples = null_sampler.sample_batch(5)
        print(f"Only nulls samples: {null_samples}")
    except Exception as e:
        print(f"Only nulls error: {e}")

    # Test with single unique string
    try:
        single_sampler, _ = create_validate_text_sampler_from_data(
            ["same"] * 100, seed=42
        )
        single_samples = single_sampler.sample_batch(10)
        print(f"Single string samples: {single_samples}")
    except Exception as e:
        print(f"Single string error: {e}")

    # Test with very long strings
    try:
        long_strings = ["x" * i for i in range(1, 21)] * 5
        long_sampler, _ = create_validate_text_sampler_from_data(long_strings, seed=42)
        long_samples = long_sampler.sample_batch(10)
        print(
            f"Long strings samples (lengths): {[len(s) if s else 0 for s in long_samples]}"
        )
    except Exception as e:
        print(f"Long strings error: {e}")


def run_comprehensive_text_sampler_test():
    """Run comprehensive test suite for TextSampler"""
    print("🚀 Starting Comprehensive TextSampler Test Suite")
    print("=" * 60)

    # Set random seed for reproducibility
    seed = get_validate_seed()
    print(f"🎲 Using test seed: {seed}")
    random.seed(seed)
    np.random.seed(seed)

    # 1. Basic functionality test
    original_data, sampler, generated_samples = validate_text_sampler_basic()

    # 2. TiDB stats test
    tidb_sampler, tidb_samples = validate_tidb_stats_sampler()

    # 3. Create visualizations
    try:
        visualize_text_distributions(original_data, generated_samples, sampler)
        print("✅ Visualizations created successfully")
    except Exception as e:
        print(f"❌ Visualization failed: {e}")
        import traceback

        traceback.print_exc()

    # 4. Test common prefix scenarios
    validate_common_prefix_scenarios()

    # 5. Test common prefix accuracy
    validate_common_prefix_accuracy()

    # 6. Test different charsets
    validate_charset_variations()

    # 7. Test edge cases
    validate_edge_cases()

    print("\n" + "=" * 60)
    print("🎉 Comprehensive TextSampler Test Suite Completed!")


if __name__ == "__main__":
    run_comprehensive_text_sampler_test()