import logging
import math
import random
import warnings
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, field
from datetime import datetime
from typing import Dict, List, Optional, Union

import avro.io
import avro.schema
from pydantic import BaseModel, ConfigDict, Field

from datahub.ingestion.source.kafka.kafka_config import ProfilerConfig
from datahub.ingestion.source.kafka.kafka_constants import (
    DEFAULT_HISTOGRAM_BUCKETS,
    DEFAULT_MAX_FIELDS_TO_PROFILE,
    MAX_DISTINCT_VALUE_FREQUENCIES,
    MAX_SAMPLE_VALUE_STR_LENGTH,
    PROFILER_BOOLEAN_TYPES,
    PROFILER_DATETIME_TYPES,
    PROFILER_NUMERIC_TYPES,
    PROFILER_STRING_TYPES,
    PROFILER_UNKNOWN_TYPES,
    ProfilerFieldType,
)
from datahub.ingestion.source.kafka.kafka_profiler_utils import (
    AGGREGATION_OVERFLOW_THRESHOLD,
    VALUE_OVERFLOW_THRESHOLD,
    calculate_numeric_stats,
    filter_numeric_values,
)
from datahub.ingestion.source.kafka.kafka_report import KafkaSourceReport
from datahub.ingestion.source.kafka.kafka_utils import MessageValue
from datahub.metadata.schema_classes import (
    ArrayTypeClass,
    BooleanTypeClass,
    BytesTypeClass,
    CalendarIntervalClass,
    DatasetFieldProfileClass,
    DatasetProfileClass,
    DateTypeClass,
    EnumTypeClass,
    FixedTypeClass,
    HistogramClass,
    KafkaSchemaClass,
    MapTypeClass,
    NullTypeClass,
    NumberTypeClass,
    PartitionSpecClass,
    PartitionTypeClass,
    QuantileClass,
    RecordTypeClass,
    SchemaMetadataClass,
    StringTypeClass,
    TimeTypeClass,
    TimeWindowSizeClass,
    UnionTypeClass,
    ValueFrequencyClass,
)

logger = logging.getLogger(__name__)

# Suppress Avro logical type warnings - these are expected when schemas use custom logical types
warnings.filterwarnings(
    "ignore",
    category=UserWarning,
    module="avro.schema",
    message=".*Unknown .*, using .*",
)

# Profiled field values are exactly Kafka message values (nested dict/list/bytes
# included) — the stats helpers handle the full set, so they share one type.
FieldValue = MessageValue
NumericValue = Union[int, float]


@dataclass
class FieldMappings:
    paths: Dict[str, str] = field(default_factory=dict)  # clean name -> full field path
    values: Dict[str, List[FieldValue]] = field(
        default_factory=dict
    )  # full field path -> collected values


class KafkaFieldStatistics(BaseModel):
    field_path: str
    data_type: Optional[str] = None
    sample_values: List[str] = Field(default_factory=list)
    unique_count: int = Field(default=0, ge=0)
    unique_proportion: float = Field(default=0.0, ge=0.0, le=1.0)
    null_count: int = Field(default=0, ge=0)
    null_proportion: float = Field(default=0.0, ge=0.0, le=1.0)
    min_value: Optional[float] = None
    max_value: Optional[float] = None
    mean_value: Optional[float] = None
    median_value: Optional[float] = None
    stdev: Optional[float] = None
    quantiles: List[QuantileClass] = Field(default_factory=list)
    distinct_value_frequencies: Dict[str, int] = Field(default_factory=dict)

    # validate_assignment enforces the count/proportion constraints on
    # post-construction assignment too.
    model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)


def is_overflow_value(v: FieldValue) -> bool:
    # Guards numeric aggregation against float64 overflow only. NOT a sentinel filter:
    # legitimate values like -1 and INT_MAX must be profiled normally.
    if isinstance(v, (int, float)):
        return abs(v) > VALUE_OVERFLOW_THRESHOLD
    return False


def clean_field_path(field_path: str) -> str:
    # Return the innermost field name, ignoring Avro path tokens like
    # [version=...]/[type=...] and stripping any trailing [key=...] suffix.
    parts = field_path.split(".")
    for part in reversed(parts):
        if part and not (part.startswith("[version=") or part.startswith("[type=")):
            return part.split("[key=")[0]

    return field_path


class KafkaProfiler:
    def __init__(
        self,
        profiler_config: ProfilerConfig,
        report: Optional[KafkaSourceReport] = None,
    ) -> None:
        self.profiler_config = profiler_config
        self.report = report
        self._expensive_profiling_disabled = (
            profiler_config.turn_off_expensive_profiling_metrics
        )
        self._processed_field_count = 0
        self.profiling_errors: List[str] = []
        self.samples_skipped = 0

    @staticmethod
    def profile_topic(
        topic_name: str,
        samples: List[Dict[str, MessageValue]],
        schema_metadata: Optional[SchemaMetadataClass],
        config: ProfilerConfig,
        report: Optional[KafkaSourceReport] = None,
        report_lock: Optional[AbstractContextManager] = None,
    ) -> Optional[DatasetProfileClass]:
        if not samples:
            logger.warning(f"No samples available for topic {topic_name}")
            return None

        profiler = KafkaProfiler(config, report=report)

        try:
            profile = profiler.generate_dataset_profile(samples, schema_metadata)

            if profile:
                profile.partitionSpec = PartitionSpecClass(
                    partition=f"SAMPLE ({len(samples)} messages)",
                    type=PartitionTypeClass.QUERY,
                )

                if profiler.profiling_errors:
                    logger.warning(
                        f"Profiling completed with {len(profiler.profiling_errors)} errors for topic {topic_name}. "
                        f"Skipped {profiler.samples_skipped} samples. First error: {profiler.profiling_errors[0]}"
                    )
                    if report is not None:
                        with report_lock or nullcontext():
                            report.profiling_samples_skipped += profiler.samples_skipped
                            report.warning(
                                "profiling",
                                f"Profiling of topic {topic_name} skipped "
                                f"{profiler.samples_skipped} sample(s) due to processing "
                                f"errors. First error: {profiler.profiling_errors[0]}",
                                log=False,
                            )

            return profile

        except RecursionError as e:
            error_msg = (
                f"Maximum recursion depth exceeded while profiling topic {topic_name}. "
                f"This is likely due to deeply nested or circular JSON structures. "
                f"Consider reducing the message complexity or adjusting flatten_json max_depth. Error: {e}"
            )
            logger.error(error_msg)
            raise RuntimeError(error_msg) from e
        except Exception as e:
            error_msg = f"Failed to profile topic {topic_name}: {e}"
            logger.error(error_msg)
            raise

    def _get_sample_values(
        self, values: List[FieldValue], max_samples: int
    ) -> List[str]:
        if not values:
            return []

        try:
            sample_size = min(max_samples, len(values))
            indices = sorted(random.sample(range(len(values)), sample_size))
            # Cap each rendered value so a huge payload can't bloat the profile.
            return [str(values[i])[:MAX_SAMPLE_VALUE_STR_LENGTH] for i in indices]
        except Exception as e:
            logger.debug(f"Failed to build sample values: {e}")
            return []

    def _create_histogram(
        self, values: List[NumericValue], max_buckets: int = DEFAULT_HISTOGRAM_BUCKETS
    ) -> Optional[HistogramClass]:
        try:
            if not values or len(values) < 2:
                return None

            # Only create histograms for numeric values
            numeric_values = [
                v
                for v in filter_numeric_values(values, exclude_special=True)
                if abs(v) < AGGREGATION_OVERFLOW_THRESHOLD
            ]

            if len(numeric_values) < 2:
                return None

            # Create buckets
            min_val = min(numeric_values)
            max_val = max(numeric_values)
            if min_val == max_val:
                return None

            bucket_size = (max_val - min_val) / max_buckets
            buckets = [min_val + i * bucket_size for i in range(max_buckets + 1)]
            counts = [0] * max_buckets

            # Count values in buckets
            for val in numeric_values:
                bucket = int((val - min_val) / bucket_size)
                if bucket >= max_buckets:
                    bucket = max_buckets - 1
                counts[bucket] += 1

            return HistogramClass(
                boundaries=[str(b) for b in buckets],
                heights=[float(c) / len(numeric_values) for c in counts],
            )
        except Exception as e:
            logger.warning(f"Error creating histogram: {e}")
            return None

    def _init_schema_fields(
        self, schema_metadata: Optional[SchemaMetadataClass]
    ) -> FieldMappings:
        field_mappings = FieldMappings()

        if not schema_metadata or not isinstance(
            schema_metadata.platformSchema, KafkaSchemaClass
        ):
            return field_mappings

        # Handle key schema if present
        if schema_metadata.platformSchema.keySchema:
            try:
                key_schema = avro.schema.parse(
                    schema_metadata.platformSchema.keySchema,
                    validate_names=False,
                )
                if hasattr(key_schema, "fields"):
                    for schema_field in key_schema.fields:
                        clean_name = clean_field_path(schema_field.name)
                        field_mappings.paths[clean_name] = schema_field.name
                        field_mappings.values[schema_field.name] = []
            except Exception as e:
                logger.warning(f"Failed to parse key schema: {e}")

        # Map all schema fields
        for schema_field in schema_metadata.fields or []:
            clean_name = clean_field_path(schema_field.fieldPath)
            field_mappings.paths[clean_name] = schema_field.fieldPath
            field_mappings.values[schema_field.fieldPath] = []

        return field_mappings

    def _get_field_path(
        self,
        field_name: str,
        schema_metadata: Optional[SchemaMetadataClass],
        field_paths: Dict[str, str],
    ) -> Optional[str]:
        if field_name in ("offset", "timestamp"):
            return None

        # Handle key prefix specially
        if field_name.startswith("key."):
            base_field_name = field_name[4:]  # Remove "key." prefix
            if schema_metadata and schema_metadata.fields:
                # Look for field with [key=True] annotation and the matching base name
                for schema_field in schema_metadata.fields:
                    if (
                        "[key=True]" in schema_field.fieldPath
                        and clean_field_path(schema_field.fieldPath) == base_field_name
                    ):
                        return schema_field.fieldPath

            # If no match found in schema, keep the key prefix
            return field_name

        # Special case for simple key field (not a complex key)
        if field_name == "key" and schema_metadata and schema_metadata.fields:
            # Look for schema fields with [key=True] in fieldPath
            # We should match on the complete field path, not just the field name
            key_schema_fields = [
                schema_field
                for schema_field in schema_metadata.fields
                if "[key=True]" in schema_field.fieldPath
            ]

            # If we found exactly one key field, use its full fieldPath
            if len(key_schema_fields) == 1:
                return key_schema_fields[0].fieldPath
            # If we found multiple key fields, try to find the best match
            elif len(key_schema_fields) > 1:
                # Look for a simple type like long, string, etc.
                for schema_field in key_schema_fields:
                    if "[type=" in schema_field.fieldPath:
                        return schema_field.fieldPath
                # If no simple type found, just use the first one
                return key_schema_fields[0].fieldPath

            return field_name  # No match found, keep as-is

        # Try to find matching schema field
        clean_sample = clean_field_path(field_name)
        if schema_metadata and schema_metadata.fields:
            for schema_field in schema_metadata.fields:
                # Skip key fields when looking for regular field matches
                if "[key=True]" in schema_field.fieldPath:
                    continue

                if clean_field_path(schema_field.fieldPath) == clean_sample:
                    return schema_field.fieldPath

        return field_paths.get(clean_sample, field_name)

    def _process_sample(
        self,
        sample: Dict[str, MessageValue],
        schema_metadata: Optional[SchemaMetadataClass],
        field_mappings: FieldMappings,
    ) -> None:
        for field_name, value in sample.items():
            field_path = self._get_field_path(
                field_name, schema_metadata, field_mappings.paths
            )
            if not field_path:
                continue

            values = field_mappings.values.setdefault(field_path, [])
            # Skip only values that would overflow numeric aggregation.
            if not is_overflow_value(value):
                values.append(value)

    def _calculate_field_stats(
        self,
        field_values: Dict[str, List[FieldValue]],
        schema_metadata: Optional[SchemaMetadataClass] = None,
    ) -> Dict[str, KafkaFieldStatistics]:
        field_stats = {}
        for field_path, values in field_values.items():
            if values:  # Only process fields that have values
                try:
                    field_stats[field_path] = self._process_field_statistics(
                        field_path=field_path,
                        values=values,
                        schema_metadata=schema_metadata,
                    )
                except Exception as e:
                    logger.warning(
                        f"Failed to process statistics for field {field_path}: {e}"
                    )
                    continue
        return field_stats

    def generate_dataset_profile(
        self,
        samples: List[Dict[str, MessageValue]],
        schema_metadata: Optional[SchemaMetadataClass] = None,
    ) -> DatasetProfileClass:
        try:
            # Initialize field mappings from schema
            field_mappings = self._init_schema_fields(schema_metadata)

            # Process all samples with error handling for each sample
            for i, sample in enumerate(samples):
                try:
                    self._process_sample(sample, schema_metadata, field_mappings)
                except RecursionError as e:
                    error_msg = f"Skipping sample {i} due to recursion error (likely circular/deep JSON): {e}"
                    logger.warning(error_msg)
                    self.profiling_errors.append(error_msg)
                    self.samples_skipped += 1
                    continue
                except Exception as e:
                    error_msg = f"Error processing sample {i}: {e}"
                    logger.warning(error_msg)
                    self.profiling_errors.append(error_msg)
                    self.samples_skipped += 1
                    continue

            # Calculate statistics for all fields
            field_stats = self._calculate_field_stats(
                field_mappings.values, schema_metadata
            )

            # Create and return profile
            return self.create_profile_data(field_stats, len(samples))

        except RecursionError as e:
            error_msg = f"Maximum recursion depth exceeded in dataset profiling: {e}"
            logger.error(error_msg)
            self.profiling_errors.append(error_msg)
            raise RuntimeError(error_msg) from e
        except Exception as e:
            error_msg = f"Unexpected error in dataset profiling: {e}"
            logger.error(error_msg)
            self.profiling_errors.append(error_msg)
            raise

    def _get_data_type_from_schema_field_type(
        self,
        type_class: Union[
            NumberTypeClass,
            BooleanTypeClass,
            DateTypeClass,
            TimeTypeClass,
            StringTypeClass,
            EnumTypeClass,
            ArrayTypeClass,
            MapTypeClass,
            RecordTypeClass,
            UnionTypeClass,
            BytesTypeClass,
            FixedTypeClass,
            NullTypeClass,
        ],
    ) -> ProfilerFieldType:
        if isinstance(type_class, NumberTypeClass):
            return ProfilerFieldType.NUMERIC
        elif isinstance(type_class, BooleanTypeClass):
            return ProfilerFieldType.BOOLEAN
        elif isinstance(type_class, (DateTypeClass, TimeTypeClass)):
            return ProfilerFieldType.DATETIME
        elif isinstance(type_class, (StringTypeClass, EnumTypeClass)):
            return ProfilerFieldType.STRING
        return ProfilerFieldType.UNKNOWN

    def _get_data_type_from_native_type(self, native_type: str) -> ProfilerFieldType:
        if native_type in PROFILER_NUMERIC_TYPES:
            return ProfilerFieldType.NUMERIC
        elif native_type in PROFILER_BOOLEAN_TYPES:
            return ProfilerFieldType.BOOLEAN
        elif native_type in PROFILER_DATETIME_TYPES:
            return ProfilerFieldType.DATETIME
        elif native_type in PROFILER_STRING_TYPES:
            return ProfilerFieldType.STRING
        elif native_type in PROFILER_UNKNOWN_TYPES:
            return ProfilerFieldType.UNKNOWN
        else:
            return ProfilerFieldType.STRING

    def _process_field_statistics(
        self,
        field_path: str,
        values: List[FieldValue],
        schema_metadata: Optional[SchemaMetadataClass] = None,
    ) -> KafkaFieldStatistics:
        # Determine data type and filter values
        data_type = self._determine_field_data_type(field_path, values, schema_metadata)
        non_null_values = self._filter_non_null_values(values)

        # Check field limits for expensive profiling
        if self._should_skip_field_processing():
            return self._create_minimal_stats(field_path, data_type)

        # Create base statistics
        stats = self._create_base_field_statistics(
            field_path, data_type, values, non_null_values
        )

        # Add distinct value statistics
        self._add_distinct_value_statistics(stats, values)

        # Add numeric statistics if applicable
        if data_type == ProfilerFieldType.NUMERIC:
            self._add_numeric_field_statistics(stats, non_null_values)

        return stats

    def _determine_field_data_type(
        self,
        field_path: str,
        values: List[FieldValue],
        schema_metadata: Optional[SchemaMetadataClass],
    ) -> ProfilerFieldType:
        data_type = ProfilerFieldType.STRING  # Default fallback

        # Try to determine type from schema metadata first
        if schema_metadata and schema_metadata.fields:
            data_type = self._get_data_type_from_schema(field_path, schema_metadata)
            if data_type != ProfilerFieldType.STRING:
                return data_type

        # Fallback to detecting from sample values
        return self._detect_data_type_from_values(values, data_type)

    def _filter_non_null_values(self, values: List[FieldValue]) -> List[FieldValue]:
        return [
            v
            for v in values
            if v is not None
            and v != ""
            and not (isinstance(v, float) and (math.isnan(v) or math.isinf(v)))
        ]

    def _get_data_type_from_schema(
        self, field_path: str, schema_metadata: SchemaMetadataClass
    ) -> ProfilerFieldType:
        for schema_field in schema_metadata.fields:
            if schema_field.fieldPath == field_path:
                # Check the actual schema field type (more reliable than nativeDataType)
                if schema_field.type and schema_field.type.type:
                    data_type = self._get_data_type_from_schema_field_type(
                        schema_field.type.type
                    )
                    if data_type != ProfilerFieldType.STRING:  # Found a specific type
                        return data_type

                # Fallback to native data type if type class is not available
                native_type = getattr(schema_field, "nativeDataType", "").lower()
                data_type = self._get_data_type_from_native_type(native_type)
                if data_type != ProfilerFieldType.STRING:  # Found a specific type
                    return data_type
        return ProfilerFieldType.STRING

    def _detect_data_type_from_values(
        self, values: List[FieldValue], default_type: ProfilerFieldType
    ) -> ProfilerFieldType:
        non_null_values = self._filter_non_null_values(values)
        if not non_null_values:
            return default_type

        sample_value = non_null_values[0]
        # bool is a subclass of int, so check it first.
        if isinstance(sample_value, bool):
            return ProfilerFieldType.BOOLEAN
        elif isinstance(sample_value, (int, float)):
            return ProfilerFieldType.NUMERIC
        elif isinstance(sample_value, (dict, list)):
            return (
                ProfilerFieldType.UNKNOWN
            )  # Complex types -> UNKNOWN (aligned with GE)
        return default_type

    def _should_skip_field_processing(self) -> bool:
        if not self._expensive_profiling_disabled:
            return False

        max_fields = (
            self.profiler_config.max_number_of_fields_to_profile
            or DEFAULT_MAX_FIELDS_TO_PROFILE
        )
        self._processed_field_count += 1
        return self._processed_field_count > max_fields

    def _create_minimal_stats(
        self, field_path: str, data_type: str
    ) -> KafkaFieldStatistics:
        return KafkaFieldStatistics(
            field_path=field_path,
            sample_values=[],
            data_type=data_type,
        )

    def _create_base_field_statistics(
        self,
        field_path: str,
        data_type: str,
        values: List[FieldValue],
        non_null_values: List[FieldValue],
    ) -> KafkaFieldStatistics:
        total_count = len(values)

        # Get sample values
        sample_values = (
            self._get_sample_values(
                [
                    v
                    for v in non_null_values
                    if not (isinstance(v, float) and math.isnan(v))
                    and not is_overflow_value(v)
                ],
                max_samples=self.profiler_config.field_sample_values_limit,
            )
            if self.profiler_config.include_field_sample_values
            and not self._expensive_profiling_disabled
            else []
        )

        stats = KafkaFieldStatistics(
            field_path=field_path,
            sample_values=sample_values,
            data_type=data_type,
        )

        # Calculate null statistics - treat NaN as null
        stats.null_count = sum(
            1
            for v in values
            if v is None
            or v == ""
            or (isinstance(v, float) and (math.isnan(v) or math.isinf(v)))
        )
        stats.null_proportion = stats.null_count / total_count if total_count > 0 else 0

        return stats

    def _add_distinct_value_statistics(
        self, stats: KafkaFieldStatistics, values: List[FieldValue]
    ) -> None:
        if self._expensive_profiling_disabled:
            return

        need_distinct = self.profiler_config.include_field_distinct_count
        # Histograms are rebuilt from this frequency map, so also build it for them.
        need_histogram = (
            self.profiler_config.include_field_histogram
            and stats.data_type == ProfilerFieldType.NUMERIC
        )
        if not (need_distinct or need_histogram):
            return

        # Convert to strings only for counting distinct values, skip NaN
        value_counts: Dict[str, int] = {}
        for value in values:
            if not (
                isinstance(value, float) and (math.isnan(value) or math.isinf(value))
            ):
                str_value = str(value)
                value_counts[str_value] = value_counts.get(str_value, 0) + 1

        stats.distinct_value_frequencies = value_counts
        if need_distinct:
            stats.unique_count = len(value_counts)
            stats.unique_proportion = stats.unique_count / len(values) if values else 0

    def _add_numeric_field_statistics(
        self, stats: KafkaFieldStatistics, non_null_values: List[FieldValue]
    ) -> None:
        numeric_values = self._extract_numeric_values(non_null_values)
        if not numeric_values:
            return

        # Calculate basic numeric statistics
        numeric_stats = calculate_numeric_stats(numeric_values)
        stats.min_value = numeric_stats.min
        stats.max_value = numeric_stats.max
        stats.mean_value = numeric_stats.mean
        stats.median_value = numeric_stats.median
        stats.stdev = numeric_stats.stdev

        # Calculate quantiles (skip for expensive profiling)
        if (
            self.profiler_config.include_field_quantiles
            and len(numeric_values) >= 4
            and not self._expensive_profiling_disabled
        ):
            self._add_quantile_statistics(stats, numeric_values)

    def _extract_numeric_values(self, non_null_values: List[FieldValue]) -> List[float]:
        numeric_values = []
        for v in non_null_values:
            try:
                # Handle both numeric types and string representations of numbers
                if isinstance(v, (int, float)):
                    num_val = float(v)
                elif isinstance(v, str):
                    num_val = float(v)  # Convert string to float
                else:
                    continue  # Skip non-numeric types

                # Skip invalid float values and special values
                if (
                    math.isnan(num_val)
                    or math.isinf(num_val)
                    or is_overflow_value(num_val)
                ):
                    continue

                numeric_values.append(num_val)
            except (ValueError, TypeError):
                # Skip values that cannot be converted to numbers
                continue
        return numeric_values

    def _add_quantile_statistics(
        self, stats: KafkaFieldStatistics, numeric_values: List[float]
    ) -> None:
        sorted_values = sorted(numeric_values)
        stats.quantiles = [
            QuantileClass(
                quantile=str(0.25),
                value=str(sorted_values[int(len(sorted_values) * 0.25)]),
            ),
            QuantileClass(
                quantile=str(0.5),
                value=str(sorted_values[int(len(sorted_values) * 0.5)]),
            ),
            QuantileClass(
                quantile=str(0.75),
                value=str(sorted_values[int(len(sorted_values) * 0.75)]),
            ),
        ]

    def create_profile_data(
        self, field_stats: Dict[str, KafkaFieldStatistics], sample_count: int
    ) -> DatasetProfileClass:
        timestamp_millis = int(datetime.now().timestamp() * 1000)
        field_profiles = []

        for field_path, stats in field_stats.items():
            histogram = None
            if (
                self.profiler_config.include_field_histogram
                and stats.data_type == ProfilerFieldType.NUMERIC
                and not self._expensive_profiling_disabled
            ):
                # Get values from distinct_value_frequencies
                values = []
                if stats.distinct_value_frequencies:
                    for value_str, freq in stats.distinct_value_frequencies.items():
                        try:
                            value = float(value_str)
                            if (
                                not math.isnan(value)
                                and not math.isinf(value)
                                and not is_overflow_value(value)
                            ):
                                values.extend([value] * freq)
                        except ValueError:
                            continue

                histogram = self._create_histogram(values)

            field_profile = DatasetFieldProfileClass(
                fieldPath=field_path,
                sampleValues=stats.sample_values
                if self.profiler_config.include_field_sample_values
                else None,
                uniqueCount=stats.unique_count
                if self.profiler_config.include_field_distinct_count
                else None,
                uniqueProportion=stats.unique_proportion
                if self.profiler_config.include_field_distinct_count
                else None,
                nullCount=stats.null_count
                if self.profiler_config.include_field_null_count
                else None,
                nullProportion=stats.null_proportion
                if self.profiler_config.include_field_null_count
                else None,
                min=str(stats.min_value)
                if self.profiler_config.include_field_min_value
                and stats.min_value is not None
                else None,
                max=str(stats.max_value)
                if self.profiler_config.include_field_max_value
                and stats.max_value is not None
                else None,
                mean=str(stats.mean_value)
                if self.profiler_config.include_field_mean_value
                and stats.mean_value is not None
                else None,
                median=str(stats.median_value)
                if self.profiler_config.include_field_median_value
                and stats.median_value is not None
                else None,
                stdev=str(stats.stdev)
                if self.profiler_config.include_field_stddev_value
                and stats.stdev is not None
                else None,
                quantiles=stats.quantiles
                if self.profiler_config.include_field_quantiles and stats.quantiles
                else None,
                distinctValueFrequencies=[
                    ValueFrequencyClass(value=str(value), frequency=freq)
                    for value, freq in sorted(
                        stats.distinct_value_frequencies.items(),
                        key=lambda x: x[1],
                        reverse=True,
                    )[:MAX_DISTINCT_VALUE_FREQUENCIES]
                ]
                if self.profiler_config.include_field_distinct_value_frequencies
                and stats.distinct_value_frequencies
                and not self._expensive_profiling_disabled
                else None,
                histogram=histogram
                if self.profiler_config.include_field_histogram
                and not self._expensive_profiling_disabled
                else None,
            )
            field_profiles.append(field_profile)

        return DatasetProfileClass(
            timestampMillis=timestamp_millis,
            columnCount=len(field_profiles),
            eventGranularity=TimeWindowSizeClass(
                unit=CalendarIntervalClass.SECOND,
                multiple=self.profiler_config.max_sample_time_seconds,
            ),
            # Add partition specification
            partitionSpec=PartitionSpecClass(
                partition=f"SAMPLE ({str(sample_count)} samples)",
                type=PartitionTypeClass.QUERY,
            ),
            fieldProfiles=field_profiles,
        )
