import pandas as pd

from samplers.tidb_types import create_tidb_bucket, create_tidb_topn_item

from .datetime_sampler import DatetimeSampler
from .numeric_samplers import DecimalSampler, IntSampler
from .text_sampler import TextSampler

TEXT_SAMPLER_TYPES = [
    "char",
    "varchar",
    "text",
    "tinytext",
    "mediumtext",
    "longtext",
    "binary",
    "varbinary",
    "blob",
    "tinyblob",
    "mediumblob",
    "longblob",
]
INT_SAMPLER_TYPES = [
    "tinyint",
    "smallint",
    "mediumint",
    "int",
    "integer",
    "bigint",
]
DECIMAL_SAMPLER_TYPES = ["float", "double", "decimal", "numeric"]
DATETIME_SAMPLER_TYPES = ["datetime", "timestamp", "date", "time"]


def sampler_factory(
    data_type: str,
    stats_buckets_df: pd.DataFrame,
    stats_top_n_df: pd.DataFrame,
    constraint_type: str = "normal",  # "normal" or "unique"
    column_name: str = None,  # Column name for logging
    **kwargs,
):
    # Extract seed from kwargs, default to 42 if not provided
    seed = kwargs.get("seed", 42)
    buckets = [
        (
            (bucket["lower_bound"], bucket["upper_bound"]),
            bucket["count"],
            bucket["repeats"],
            bucket["bucket_id"],
            bucket["column_name"],
        )
        for bucket in stats_buckets_df.to_dict(orient="records")
    ]
    # Process TopN items through TiDB types (same flow as buckets: 先 golang,再 types)
    tidb_topn_items = []
    raw_top_n = [
        (top_n["value"], top_n["count"])
        for top_n in stats_top_n_df.to_dict(orient="records")
    ]

    for topn_item in raw_top_n:
        encoded_value, count = topn_item

        # Create TiDB TopN item with type-specific parsing
        try:
            tidb_topn_item = create_tidb_topn_item(
                data_type=data_type, encoded=encoded_value, count=count
            )

            # Keep as TiDB TopN item for the sampler
            tidb_topn_items.append(tidb_topn_item)
        except Exception as e:
            # Fall back to original TopN format if TiDB conversion fails
            print(
                f"Warning: Failed to create TiDB TopN item for {data_type}, falling back to original format: {e}"
            )
            tidb_topn_items.append(topn_item)

    # Use processed TopN items or fall back to raw data
    top_n = tidb_topn_items if tidb_topn_items else raw_top_n

    # setup sampler by stats_buckets and stats_top_n
    # Create TiDB-compatible buckets and convert to appropriate sampler
    tidb_buckets = []
    for bucket_data in buckets:
        (
            (lower_bound, upper_bound),
            count,
            repeats,
            bucket_id,
            column_name,
        ) = bucket_data

        # Create TiDB bucket with type-specific parsing
        try:
            # Only pass character_maximum_length for string types
            bucket_kwargs = {
                "data_type": data_type,
                "lower_bound": lower_bound,
                "upper_bound": upper_bound,
                "count": count,
                "repeats": repeats,
                "ndv": 0,  # Default for now, could be calculated
                "bucket_id": bucket_id,  # Pass bucket_id
                "column_name": column_name,  # Pass column_name
            }
            # Add character_maximum_length only for string types
            if data_type.lower() in TEXT_SAMPLER_TYPES:
                bucket_kwargs["character_maximum_length"] = kwargs[
                    "character_maximum_length"
                ]

            tidb_bucket = create_tidb_bucket(**bucket_kwargs)

            # Keep as TiDB bucket for the sampler
            tidb_buckets.append(tidb_bucket)
        except Exception as e:
            # Fall back to original bucket format if TiDB conversion fails
            print(
                f"Warning: Failed to create TiDB bucket for {data_type}, falling back to original format: {e}"
            )
            tidb_buckets.append(bucket_data)

    if data_type in TEXT_SAMPLER_TYPES:
        # Filter to ensure only TiDBStringBucket and TiDBStringTopNItem types
        string_buckets = (
            [
                b
                for b in tidb_buckets
                if hasattr(b, "data_type")
                and b.__class__.__name__ == "TiDBStringBucket"
            ]
            if tidb_buckets
            else []
        )
        string_topn = (
            [
                t
                for t in tidb_topn_items
                if hasattr(t, "data_type")
                and t.__class__.__name__ == "TiDBStringTopNItem"
            ]
            if tidb_topn_items
            else []
        )

        return TextSampler(
            string_buckets or buckets,
            string_topn or top_n,
            character_maximum_length=kwargs["character_maximum_length"],
            constraint_type=constraint_type,
            column_name=column_name,
            seed=seed,
        )
    elif data_type in INT_SAMPLER_TYPES:
        # Filter to ensure only TiDBNumericBucket and TiDBNumericTopNItem types
        numeric_buckets = (
            [
                b
                for b in tidb_buckets
                if hasattr(b, "data_type")
                and b.__class__.__name__ == "TiDBNumericBucket"
            ]
            if tidb_buckets
            else []
        )
        numeric_topn = (
            [
                t
                for t in tidb_topn_items
                if hasattr(t, "data_type")
                and t.__class__.__name__ == "TiDBNumericTopNItem"
            ]
            if tidb_topn_items
            else []
        )

        return IntSampler(
            numeric_buckets or buckets,
            numeric_topn or top_n,
            constraint_type=constraint_type,
            column_name=column_name,
            seed=seed,
        )
    elif data_type in DECIMAL_SAMPLER_TYPES:
        # Filter to ensure only TiDBNumericBucket and TiDBNumericTopNItem types
        numeric_buckets = (
            [
                b
                for b in tidb_buckets
                if hasattr(b, "data_type")
                and b.__class__.__name__ == "TiDBNumericBucket"
            ]
            if tidb_buckets
            else []
        )
        numeric_topn = (
            [
                t
                for t in tidb_topn_items
                if hasattr(t, "data_type")
                and t.__class__.__name__ == "TiDBNumericTopNItem"
            ]
            if tidb_topn_items
            else []
        )

        return DecimalSampler(
            numeric_buckets or buckets,
            numeric_topn or top_n,
            constraint_type=constraint_type,
            column_name=column_name,
            seed=seed,
        )
    elif data_type in DATETIME_SAMPLER_TYPES:
        # Filter to ensure only TiDBTimeBucket and TiDBTimeTopNItem types
        time_buckets = (
            [
                b
                for b in tidb_buckets
                if hasattr(b, "data_type") and b.__class__.__name__ == "TiDBTimeBucket"
            ]
            if tidb_buckets
            else []
        )
        time_topn = (
            [
                t
                for t in tidb_topn_items
                if hasattr(t, "data_type")
                and t.__class__.__name__ == "TiDBTimeTopNItem"
            ]
            if tidb_topn_items
            else []
        )

        return DatetimeSampler(
            time_buckets or [],
            time_topn or [],
            datetime_type=data_type.upper(),
            constraint_type=constraint_type,
            column_name=column_name,
            seed=seed,
        )
    else:
        raise NotImplementedError(f"column type `{data_type}` not supported yet")