"""
PAMOLA.CORE - Privacy-Preserving AI Fake Data Generators
Module:        Fake Data Generator Operation
Package:       pamola_core.fake_data
Version:       3.0.0
Status:        stable
Author:        PAMOLA Core Team
Created:       2025-03
Updated:       2025-06-15
License:       BSD 3-Clause

Description:
   Base classes and utilities for fake / synthetic data generation operations.
   Provides a standardized lifecycle, execution model, metrics collection,
   visualization support, memory-efficient processing and optional Dask/joblib
   execution paths for large-scale synthetic data workloads.

Purpose:
   Centralize generator and synthetic-data operation logic to ensure consistent
   behavior, deterministic mapping support, parameter validation, progress
   reporting and artifact management across all fake-data generators.

Key Features:
   - Standardized operation lifecycle (validation, execution, result handling)
   - Support for REPLACE and ENRICH modes for synthetic value injection
   - Configurable null handling strategies and conditional processing
   - Chunked processing for memory efficiency with Dask and joblib support
   - Comprehensive metrics collection and visualization generation
   - Deterministic mapping store support for repeatable generation
   - Result caching and artifact restoration for faster re-runs

Design Principles:
   - Type safety and explicit validation for inputs and parameters
   - Extensibility: subclasses override batch/value processing and cache params
   - Compatibility with PAMOLA operation framework and reporter interfaces
   - Performance: efficient batching, optional parallelism, and cached validators
   - Robustness: best-effort artifact generation and graceful degradation

Dependencies:
   - dask, pandas, numpy, joblib, psutil
   - pamola_core.* utilities (io, ops, fake_data.commons, utils.progress, etc.)
   - hashlib, json, datetime, pathlib, threading, contextvars
"""

import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional, Union
import pandas as pd

from pamola_core.errors.codes import ErrorCode
from pamola_core.errors.error_handler import ErrorHandler
from pamola_core.errors.exceptions import (
    TypeValidationError,
    ValidationError,
    FeatureNotImplementedError,
)
from pamola_core.common.constants import Constants
from pamola_core.fake_data.commons.base import BaseGenerator
from pamola_core.fake_data.commons.mapping_store import MappingStore
from pamola_core.fake_data.commons.metrics import (
    create_metrics_collector,
    generate_metrics_report,
)
from pamola_core.fake_data.commons.processing_utils import (
    process_dataframe_using_chunk,
    process_dataframe_using_dask,
    process_dataframe_using_joblib,
)
from pamola_core.utils.io import (
    load_data_operation,
    load_settings_operation,
)
from pamola_core.utils.io_helpers.crypto_utils import get_encryption_mode
from pamola_core.utils.ops.op_base import FieldOperation

# Import framework utilities
from pamola_core.utils.ops.op_cache import OperationCache
from pamola_core.utils.ops.op_data_source import DataSource
from pamola_core.utils.ops.op_data_writer import DataWriter
from pamola_core.utils.ops.op_result import OperationResult, OperationStatus
from pamola_core.utils.progress import HierarchicalProgressTracker
from pamola_core.utils.helpers import (
    build_base_cache,
    filter_used_kwargs,
    get_cache_result,
)
import pamola_core.utils.helpers as helpers


class GeneratorOperation(FieldOperation):
    """
    Base class for all generator-based generator or synthetic data operations.

    This class provides a consistent interface for operations that use
    generator-based approaches (e.g., fake names, fake emails, synthetic IDs),
    handling consistency, mapping stores, and metadata propagation.
    """

    def __init__(
        self,
        field_name: str,
        generator: BaseGenerator,
        generator_params: Optional[Dict[str, Any]] = None,
        consistency_mechanism: str = "prgn",
        id_field: Optional[str] = None,
        mapping_store_path: Optional[str] = None,
        save_mapping: bool = False,
        **kwargs,
    ):
        """
        Initialize a generator-based operation.

        Parameters
        ----------
        field_name : str
            Target field for generation/generator.
        generator : BaseGenerator
            Generator instance used to create synthetic data.
        generator_params : dict, optional
            Additional generator configuration parameters.
        consistency_mechanism : str, optional
            Consistency preservation mechanism (default "prgn").
        id_field : str, optional
            Unique identifier field for mapping consistency.
        mapping_store_path : str, optional
            Path to persistent mapping store (for deterministic reuse).
        mapping_store : MappingStore, optional
            Pre-loaded mapping store instance (used if provided).
        save_mapping : bool, optional
            Whether to persist mapping results after processing.
        **kwargs : dict
            Additional arguments forwarded to FieldOperation/BaseOperation.
        """

        # Default metadata
        kwargs.setdefault("name", f"{field_name}_generation")
        kwargs.setdefault(
            "description", f"Generator-based operation applied to '{field_name}'"
        )

        # Initialize parent FieldOperation
        super().__init__(field_name=field_name, **kwargs)

        # === Core Generator Configuration ===
        self.generator: BaseGenerator = generator
        self.generator_params = generator_params or {}

        # === Consistency and Mapping ===
        self.consistency_mechanism = consistency_mechanism.lower()
        self.mapping_store = MappingStore()
        self.mapping_store_path = mapping_store_path
        self.save_mapping = save_mapping

        # === id_field===
        self.id_field = id_field

        # === Descriptive Metadata ===
        self.name = f"{generator.__class__.__name__}_operation"
        self.description = (
            f"Generates synthetic '{field_name}' data using {generator.__class__.__name__} "
            f"(consistency='{self.consistency_mechanism}')"
        )

        # Create a metrics collector
        self._metrics_collector = create_metrics_collector()

        # === Optional Mapping Store Setup ===
        if mapping_store_path:
            self._initialize_mapping_store(mapping_store_path)

    def execute(
        self,
        data_source: DataSource,
        task_dir: Path,
        reporter: Any,
        progress_tracker: Optional[HierarchicalProgressTracker] = None,
        **kwargs,
    ) -> OperationResult:
        """
        Execute the generator operation with enhanced features including Dask support.

        Parameters
        -----------
        data_source : DataSource
            Source of data for the operation
        task_dir : Path
            Directory where task artifacts should be saved
        reporter : Any
            Reporter object for tracking progress and artifacts
        progress_tracker : Optional[HierarchicalProgressTracker]
            Progress tracker for the operation
        **kwargs : dict
            Additional parameters for the operation

        Returns
        --------
        OperationResult
            Results of the operation
        """
        try:
            # Start timing
            self.start_time = time.time()
            self.logger = kwargs.get("logger", self.logger)
            self.logger.info(
                f"Starting: {self.operation_name} operation at {self.start_time}"
            )

            # Initialize result object
            result = OperationResult(status=OperationStatus.PENDING)

            # Extract dataset name from kwargs (default to "main")
            dataset_name = kwargs.get("dataset_name", "main")

            # Generate single timestamp for all artifacts
            operation_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

            # Prepare directories for artifacts
            dirs = self._prepare_directories(task_dir)

            # Initialize operation cache
            self.operation_cache = OperationCache(
                cache_dir=dirs["cache"],
            )

            # Initialize error handler
            self.error_handler = ErrorHandler(
                logger=self.logger,
                operation_name=self.operation_name,
            )

            # Create DataWriter for consistent file operations
            writer = DataWriter(
                task_dir=task_dir, logger=self.logger, progress_tracker=progress_tracker
            )

            # Save operation configuration
            self.save_config(task_dir)

            self.logger.info(
                f"Visualization settings: theme={self.visualization_theme}, backend={self.visualization_backend}, strict={self.visualization_strict}, timeout={self.visualization_timeout}s"
            )

            # Load settings operation
            settings_operation = load_settings_operation(
                data_source, dataset_name, **kwargs
            )

            # Set up progress tracking with proper steps
            # Main steps: 1. Data Loading & Validation, 2. Cache check, 3. Prepare output field, 4. Processing, 5. Metrics, 6. Visualization, 7. Save output
            TOTAL_MAIN_STEPS = 6 + (
                1 if self.use_cache and not self.force_recalculation else 0
            )
            main_progress = progress_tracker
            current_steps = 0
            if main_progress:
                self.logger.info(
                    f"Setting up progress tracker with {TOTAL_MAIN_STEPS} main steps"
                )
                try:
                    main_progress.total = TOTAL_MAIN_STEPS
                    main_progress.update(
                        current_steps,
                        {
                            "step": f"Starting {self.name}",
                            "field": self.field_name,
                        },
                    )
                except Exception as e:
                    self.logger.warning(f"Could not update progress tracker: {e}")

            # Step 1: Data Loading & Validation
            if main_progress:
                current_steps += 1
                main_progress.update(
                    current_steps, {"step": "Data Loading", "field": self.field_name}
                )

            # Validate and get dataframe
            try:

                self.logger.info(
                    f"Operation: {self.operation_name}, Load data and validate input parameters"
                )
                self._original_df = self._validate_and_get_dataframe(
                    data_source, dataset_name, **settings_operation
                )

            except Exception as e:
                return self.error_handler.handle_error(
                    error=e,
                    error_code=ErrorCode.DATA_LOAD_FAILED,
                    context={"dataset": dataset_name, "operation": self.operation_name},
                    message_kwargs={"source": dataset_name, "reason": str(e)},
                )

            # Step 2: Check Cache (if enabled and not forced to recalculate)
            if self.use_cache and not self.force_recalculation:
                if main_progress:
                    current_steps += 1
                    main_progress.update(
                        current_steps,
                        {"step": "Checking cache", "field": self.field_name},
                    )

                # Load data for cache check
                self.logger.info("Checking operation cache...")
                cache_result = self._check_cache(self._original_df, reporter)

                if cache_result:
                    self.logger.info(
                        f"Using cached result for {self.field_name} gereration operation"
                    )

                    # Update progress
                    if main_progress:
                        main_progress.update(
                            current_steps,
                            {"step": "Complete (cached)", "field": self.field_name},
                        )

                    # Report cache hit to reporter
                    if reporter:
                        reporter.add_operation(
                            f"Generator of {self.field_name} (from cache)",
                            details={"cached": True},
                        )

                    return cache_result

            # Step 3: Prepare output field
            if main_progress:
                current_steps += 1
                main_progress.update(
                    current_steps,
                    {"step": "Preparing output field", "field": self.field_name},
                )

            try:
                self.output_field_name = self._prepare_output_field(self._original_df)
                self.logger.info(f"Prepared output_field: '{self.output_field_name}'")
                self._report_operation_details(reporter, self.output_field_name)
            except Exception as e:
                return self.error_handler.handle_error(
                    error=e,
                    error_code=ErrorCode.PROCESSING_FAILED,
                    context={"step": "prepare_output_field", "field": self.field_name},
                    message_kwargs={
                        "field_name": self.field_name,
                        "operation": self.operation_name,
                        "reason": str(e),
                    },
                )

            # Step 4: Processing
            if main_progress:
                current_steps += 1
                main_progress.update(
                    current_steps, {"step": "Processing", "field": self.field_name}
                )

            try:
                # Copy original data for processing
                original_data = self._original_df[self.field_name].copy(deep=True)

                # Create child progress tracker for Chunk processing
                data_tracker = None
                if main_progress and hasattr(main_progress, "create_subtask"):
                    try:
                        data_tracker = main_progress.create_subtask(
                            total=3,
                            description="Processing dataframe",
                            unit="steps",
                        )
                    except Exception as e:
                        self.logger.debug(
                            f"Could not create child progress tracker: {e}"
                        )

                # Process the filtered data
                processed_df = self._process_data_with_config(
                    df=self._original_df,
                    progress_tracker=data_tracker,
                    **kwargs,
                )

                # Get the generated data
                generated_data = processed_df[self.output_field_name]

                # Close child progress tracker
                if data_tracker:
                    try:
                        data_tracker.close()
                    except:
                        pass
            except Exception as e:
                return self.error_handler.handle_error(
                    error=e,
                    error_code=ErrorCode.PROCESSING_FAILED,
                    context={"step": "processing", "field": self.field_name},
                    message_kwargs={
                        "field_name": self.field_name,
                        "operation": self.operation_name,
                        "reason": str(e),
                    },
                )

            # Record end time after processing metrics
            self.end_time = time.time()
            if self.end_time and self.start_time:
                self.execution_time = self.end_time - self.start_time

            # Step 5: Metrics Calculation
            if main_progress:
                current_steps += 1
                main_progress.update(
                    current_steps,
                    {"step": "Metrics Calculation", "field": self.field_name},
                )

            # Initialize metrics in scope
            metrics = {}

            try:
                metrics = self._collect_metrics(processed_df)

                self._save_metrics(
                    metrics=metrics,
                    writer=writer,
                    result=result,
                    reporter=reporter,
                    progress_tracker=progress_tracker,
                    operation_timestamp=operation_timestamp,
                )
            except Exception as e:
                error_message = f"Error calculating metrics: {str(e)}"
                self.logger.warning(error_message)
                # Continue execution - metrics failure is not critical

            try:
                report_dir = task_dir / "reports"
                generate_metrics_report(
                    metrics,
                    report_dir,
                    op_type=self.operation_name,
                    field_name=self.field_name,
                    operation_timestamp=operation_timestamp,
                )

                if reporter:
                    reporter.add_operation(
                        f"Operation {self.operation_name}",
                        status="info",
                        details={
                            "message": "Generate and save metrics report successfully"
                        },
                    )
            except Exception as e:
                error_message = f"Operation: {self.operation_name}, Error generating metrics report: {str(e)}"
                self.logger.warning(error_message)
                # Continue execution - metrics failure is not critical

            # Step 6: Visualizations
            if main_progress:
                current_steps += 1
                main_progress.update(
                    current_steps,
                    {"step": "Generating Visualizations", "field": self.field_name},
                )

            # Generate visualizations if required
            # Initialize visualization paths dictionary
            if self.generate_visualization and self.visualization_backend is not None:
                try:
                    kwargs_encryption = {
                        "use_encryption": self.use_encryption,
                        "encryption_key": self.encryption_key,
                    }
                    self._handle_visualizations(
                        df=processed_df,
                        metrics=metrics,
                        task_dir=task_dir,
                        result=result,
                        reporter=reporter,
                        progress_tracker=main_progress,
                        vis_theme=self.visualization_theme,
                        vis_backend=self.visualization_backend,
                        vis_strict=self.visualization_strict,
                        vis_timeout=self.visualization_timeout,
                        operation_timestamp=operation_timestamp,
                        **kwargs_encryption,
                    )
                except Exception as e:
                    error_message = f"Error generating visualizations: {str(e)}"
                    self.logger.warning(error_message)
                    # Continue execution - visualization failure is not critical
            else:
                self.logger.info(
                    "Skipping visualizations as generate_visualization is False or backend is not set"
                )

            # Step 7: Save Output Data
            if main_progress:
                current_steps += 1
                main_progress.update(
                    current_steps,
                    {"step": "Save Output Data", "field": self.field_name},
                )

            # Save output data if required
            if self.save_output:
                try:
                    safe_kwargs = filter_used_kwargs(kwargs, self._save_output_data)
                    # Save output data
                    self._save_output_data(
                        result_df=processed_df,
                        writer=writer,
                        result=result,
                        reporter=reporter,
                        progress_tracker=main_progress,
                        timestamp=operation_timestamp,
                        **safe_kwargs,
                    )
                except Exception as e:
                    return self.error_handler.handle_error(
                        error=e,
                        error_code=ErrorCode.ARTIFACT_WRITE_FAILED,
                        context={"step": "save_output", "field": self.field_name},
                        message_kwargs={
                            "path": str(task_dir / "output"),
                            "reason": str(e),
                        },
                    )

            # Cache the result if caching is enabled
            if self.use_cache:
                try:
                    self._save_to_cache(
                        original_data=original_data,
                        generated_data=generated_data,
                        result=result,
                        task_dir=task_dir,
                    )
                except Exception as e:
                    # Failure to cache is non-critical
                    self.logger.warning(f"Failed to cache results: {str(e)}")

            # Clean up memory AFTER all write operations are complete
            self.logger.info("Cleaning up memory after all file operations")
            self._cleanup_memory(
                processed_df=processed_df,
                original_data=original_data,
                generated_data=generated_data,
            )

            # Report completion
            if reporter:
                reporter.add_operation(
                    f"Generator of {self.field_name} completed",
                    details={
                        "records_processed": self.process_count,
                        "execution_time": self.execution_time,
                    },
                )

            # Set success status
            result.status = OperationStatus.SUCCESS
            result.execution_time = self.execution_time
            self.logger.info(
                f"Processing completed {self.operation_name} operation in {self.execution_time:.2f} seconds"
            )
            return result

        except Exception as e:
            self.logger.exception(f"Error in {self.operation_name}: {str(e)}")
            return self.error_handler.handle_error(
                error=e,
                error_code=ErrorCode.PROCESSING_FAILED,
                context={"operation": self.operation_name, "field": self.field_name},
                message_kwargs={
                    "field_name": self.field_name,
                    "operation": self.operation_name,
                    "reason": str(e),
                },
            )

    def _initialize_mapping_store(self, path: Union[str, Path]) -> None:
        """
        Initialize the mapping store if needed.

        Parameters
        ----------
            path: Path to mapping store file
        """
        try:
            from pamola_core.fake_data.commons.mapping_store import MappingStore

            self.mapping_store = MappingStore()

            # Load existing mappings if the file exists
            path_obj = Path(path)
            if path_obj.exists():
                self.mapping_store.load(path_obj)
                self.logger.info(f"Loaded mapping store from {path_obj.name}")
        except Exception as e:
            self.logger.warning(f"Failed to initialize mapping store: {str(e)}")
            self.mapping_store = None

    def _validate_and_get_dataframe(
        self, data_source: DataSource, dataset_name: str, **kwargs: Any
    ) -> pd.DataFrame:
        """
        Validate data source and retrieve the main dataframe.

        Parameters
        -----------
        data_source : DataSource
            The data source to validate
        dataset_name : str
            The name of the dataset to retrieve
        **kwargs : Any
            Additional keyword arguments to pass to the data loading function

        Returns
        --------
        pd.DataFrame
            The validated dataframe

        Raises
        -------
        ValueError
            If no valid dataframe is found or the field is missing
        """
        # Get DataFrame from the data source
        df = load_data_operation(data_source, dataset_name, **kwargs)
        if df is None:
            error_message = "Failed to load input data!"
            self.logger.error(error_message)
            raise ValidationError(error_message)

        if self.field_name not in df.columns:
            error_message = f"Field {self.field_name} not found in DataFrame"
            self.logger.error(error_message)
            raise ValidationError(error_message)

        # Apply data types from data source
        try:
            df = data_source.apply_data_types(df, dataset_name)
        except (ValidationError, ValueError) as e:
            error_msg = (
                f"Failed to apply data types for dataset '{dataset_name}': {str(e)}"
            )
            self.logger.error(error_msg)
            raise ValidationError(error_msg) from e

        except TypeError as e:
            error_msg = f"Invalid dataframe type for dataset '{dataset_name}': {str(e)}"
            self.logger.error(error_msg)
            raise TypeValidationError(error_msg) from e

        return df

    def _prepare_output_field(self, df: pd.DataFrame) -> str:
        """
        Validate and generate the output field name.

        Parameters
        -----------
        df : pd.DataFrame
            The dataframe to check field names against

        Returns
        --------
        str
            The validated output field name
        """
        # Determine output field name based on mode
        if self.mode == "REPLACE":
            output_field = self.field_name
        else:  # ENRICH mode
            if self.output_field_name:
                output_field = self.output_field_name
            else:
                output_field = f"{self.column_prefix}{self.field_name}"

            # Check if output field already exists in DataFrame
            if output_field in df.columns:
                self.logger.warning(
                    f"Output field '{output_field}' already exists and will be overwritten"
                )
        return output_field

    def _report_operation_details(self, reporter: Any, output_field: str) -> None:
        """
        Report details of the operation to the reporter.

        Parameters
        -----------
        reporter : Any
            The reporter to log details to
        output_field : str
            The name of the output field
        """
        if reporter:
            reporter.add_operation(
                f"Generator field: {self.field_name}",
                details={
                    "field_name": self.field_name,
                    "output_field": output_field,
                    "mode": self.mode,
                    "null_strategy": self.null_strategy,
                    "operation_type": self.operation_name,
                },
            )

    def _process_data_with_config(
        self,
        df: pd.DataFrame,
        progress_tracker: Optional[HierarchicalProgressTracker] = None,
        **kwargs,
    ) -> pd.DataFrame:
        """
        Handle processing of the dataframe, including chunk-wise or full processing.

        Parameters
        -----------
        df : pd.DataFrame
            The dataframe to process
        progress_tracker : Optional[HierarchicalProgressTracker]
            Optional progress tracker

        Returns
        --------
        pd.DataFrame
            The processed dataframe
        """
        # Check if dataframe is empty
        if len(df) == 0:
            self.logger.warning("Empty DataFrame provided, returning as is")
            return df

        # Handle null values based on strategy
        if self.null_strategy != "PRESERVE":
            from pamola_core.anonymization.commons.data_utils import process_nulls

            df[self.field_name] = process_nulls(
                df[self.field_name], strategy=self.null_strategy.upper()
            )

        processed_df = None
        flag_processed = False
        # Backup and clear operation cache during processing
        cache_backup = self.operation_cache
        self.operation_cache = None
        self.logger.info("Process with config")

        # For larger dataframes, check if we should use parallel processing
        if not flag_processed and self.use_dask:
            try:
                self.logger.info("Parallel Enabled")
                self.logger.info("Parallel Engine: Dask")
                self.logger.info(f"Parallel Workers: {self.npartitions}")
                self.logger.info(
                    f"Using dask processing with chunk size {self.chunk_size}"
                )
                if progress_tracker:
                    progress_tracker.update(0, {"step": "Setting up dask processing"})

                self.logger.info("Process using Dask")

                # Process with Dask - delegate to subclass method
                processed_df, flag_processed = process_dataframe_using_dask(
                    df=df,
                    process_function=self.process_batch,
                    progress_tracker=progress_tracker,
                    task_logger=self.logger,
                    **kwargs,
                )

                if flag_processed:
                    self.logger.info("Completed using Dask")

            except Exception as e:
                self.logger.warning(
                    f"Error in dask processing: {e}, falling back to chunk processing"
                )

        if not flag_processed and self.use_vectorization:
            try:
                self.logger.info("Parallel Enabled")
                self.logger.info("Parallel Engine: Joblib")
                self.logger.info(f"Parallel Workers: {self.parallel_processes}")
                self.logger.info(
                    f"Using vectorized processing with chunk size {self.chunk_size}"
                )
                if progress_tracker:
                    progress_tracker.update(
                        0, {"step": "Setting up vectorized processing"}
                    )

                self.logger.info("Process using Joblib")

                processed_df, flag_processed = process_dataframe_using_joblib(
                    df=df,
                    process_function=self.process_batch,
                    progress_tracker=progress_tracker,
                    task_logger=self.logger,
                    **kwargs,
                )

                if flag_processed:
                    self.logger.info("Completed using Joblib")

            except Exception as e:
                self.logger.warning(
                    f"Error in vectorized processing: {e}, falling back to chunk processing"
                )

        if not flag_processed and self.chunk_size > 1:
            try:
                # Regular chunk processing
                self.logger.info(
                    f"Processing in chunks with chunk size {self.chunk_size}"
                )
                total_chunks = (len(df) + self.chunk_size - 1) // self.chunk_size
                self.logger.info(f"Total chunks to process: {total_chunks}")
                if progress_tracker:
                    progress_tracker.update(
                        0,
                        {
                            "step": "Processing in chunks",
                            "total_chunks": total_chunks,
                        },
                    )
                self.logger.info("Process using chunk")
                processed_df, flag_processed = process_dataframe_using_chunk(
                    df=df,
                    process_function=self.process_batch,
                    progress_tracker=progress_tracker,
                    task_logger=self.logger,
                    **kwargs,
                )
                if flag_processed:
                    self.logger.info("Completed using chunk")
            except Exception as e:
                self.logger.warning(f"Error in chunk processing: {e}")

        if not flag_processed:
            self.logger.info("Fallback process as usual")
            processed_df = self.process_batch(df, **kwargs)
            flag_processed = True

        # Update process count
        self.process_count += len(df)
        # Restore operation cache
        self.operation_cache = cache_backup
        return processed_df

    def _collect_metrics(self, df: pd.DataFrame) -> Dict[str, Any]:
        """
        Collects metrics for the operation (including generator-specific and performance metrics).

        Parameters
        ----------
        df : pd.DataFrame
            Processed DataFrame.

        Returns
        -------
        Dict[str, Any]
            Metrics for the operation.
        """
        import psutil
        import time

        # === 1. Basic metrics ===
        total_records = len(df)
        metrics = {
            "total_records": total_records,
            "non_null_records": df[self.field_name].notna().sum(),
            "execution_time": 0,  # kept for backward compatibility
            "performance": {},
        }

        # === 2. Performance metrics ===
        if getattr(self, "start_time", None) is not None:
            try:
                execution_time = time.time() - self.start_time
                records_per_second = (
                    int(total_records / execution_time) if execution_time > 0 else 0
                )

                process = psutil.Process()
                memory_info = process.memory_info()
                memory_usage_mb = memory_info.rss / (1024 * 1024)

                metrics["performance"].update(
                    {
                        "generation_time": round(execution_time, 4),
                        "records_per_second": records_per_second,
                        "memory_usage_mb": round(memory_usage_mb, 2),
                    }
                )
            except Exception as e:
                self.logger.warning(f"Error collecting performance metrics: {str(e)}")

        # === 3. Output field metrics (for ENRICH mode) ===
        if self.mode == "ENRICH" and self.output_field_name in df.columns:
            metrics["output_field"] = {
                "name": self.output_field_name,
                "non_null_records": df[self.output_field_name].notna().sum(),
            }

        # === 4. Metrics collector (distribution comparison, utility metrics) ===
        try:
            orig_series = (
                self._original_df[self.field_name]
                if self._original_df is not None
                else None
            )
            gen_series = None
            if self.mode == "REPLACE":
                gen_series = df[self.field_name]
            elif self.mode == "ENRICH" and self.output_field_name in df.columns:
                gen_series = df[self.output_field_name]

            if orig_series is not None and gen_series is not None:
                collector_metrics = self._metrics_collector.collect_metrics(
                    orig_data=orig_series,
                    gen_data=gen_series,
                    field_name=self.field_name,
                    operation_params={"field_name": self.field_name},
                )

                for key, value in collector_metrics.items():
                    if key not in metrics:
                        metrics[key] = value
                    elif isinstance(metrics[key], dict) and isinstance(value, dict):
                        metrics[key].update(value)
        except Exception as e:
            self.logger.warning(f"Error collecting detailed metrics: {str(e)}")

        # === 5. Generator-specific extension (if exists) ===
        if hasattr(self, "generator"):
            metrics["generator"] = {
                "type": self.generator.__class__.__name__,
                "consistency_mechanism": getattr(self, "consistency_mechanism", None),
            }

            # Mapping mechanism details
            if getattr(self, "consistency_mechanism", None) == "mapping":
                try:
                    field_mappings = self.mapping_store.get_field_mappings(
                        self.field_name
                    )
                    metrics["mapping"] = {"total_mappings": len(field_mappings)}
                except Exception as e:
                    self.logger.warning(f"Error collecting mapping metrics: {str(e)}")

            # Dictionary info if available
            if hasattr(self.generator, "get_dictionary_info"):
                try:
                    dictionary_info = self.generator.get_dictionary_info()
                    if dictionary_info:
                        metrics.setdefault("dictionary_metrics", {}).update(
                            dictionary_info
                        )
                except Exception as e:
                    self.logger.warning(f"Error getting dictionary info: {str(e)}")

            # Transformation / distribution comparison
            if orig_series is not None and gen_series is not None:
                try:
                    compare_metrics = self._metrics_collector.compare_distributions(
                        orig_series, gen_series
                    )
                    metrics.setdefault("transformation_metrics", {}).update(
                        compare_metrics
                    )
                except Exception as e:
                    self.logger.warning(f"Error comparing distributions: {str(e)}")

        # === 6. Operation-specific extension hook ===
        if hasattr(self, "_collect_specific_metrics"):
            try:
                specific = self._collect_specific_metrics(df)
                if specific:
                    metrics.update(specific)
            except Exception as e:
                self.logger.warning(f"Error collecting specific metrics: {str(e)}")

        return metrics

    def _collect_specific_metrics(self, df: pd.DataFrame) -> Dict[str, Any]:
        """
        Collect operation-specific metrics. Should be overridden by subclasses.

        Parameters
        ----------
            df: Processed DataFrame

        Returns
        -------
            Dictionary with metrics
        """
        return {}

    def _handle_visualizations(
        self,
        df: pd.DataFrame,
        metrics: Dict[str, Any],
        task_dir: Path,
        result: OperationResult,
        reporter: Any,
        progress_tracker: Optional[HierarchicalProgressTracker],
        vis_theme: Optional[str] = None,
        vis_backend: Optional[str] = None,
        vis_strict: bool = False,
        vis_timeout: int = 120,
        operation_timestamp: Optional[str] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        """
        Generate and save visualizations with thread-safe context support.

        Parameters
        -----------
        df : pd.DataFrame
            The processed dataframe
        metrics : Dict[str, Any]
            The collected metrics
        task_dir : Path
            The task directory
        result : OperationResult
            The operation result to add artifacts to
        reporter : Any
            The reporter to log artifacts to
        progress_tracker : Optional[HierarchicalProgressTracker]
            Optional progress tracker
        vis_theme : str, optional
            Theme to use for visualizations
        vis_backend : str, optional
            Backend to use: "plotly" or "matplotlib"
        vis_strict : bool, optional
            If True, raise exceptions for configuration errors
        vis_timeout : int, optional
            Timeout for visualization generation (default: 120 seconds)
        operation_timestamp : str, optional
            Timestamp for the operation (default: current time)
        **kwargs : dict
            Additional parameters for the operation
        """
        self.logger.info(
            f"Generating visualizations with backend: {vis_backend}, timeout: {vis_timeout}s"
        )
        try:
            import threading
            import contextvars

            visualization_paths = {}
            visualization_error = None

            def generate_viz_with_diagnostics():
                nonlocal visualization_paths, visualization_error
                thread_id = threading.current_thread().ident
                thread_name = threading.current_thread().name

                self.logger.info(
                    f"[DIAG] Visualization thread started - Thread ID: {thread_id}, Name: {thread_name}"
                )
                self.logger.info(
                    f"[DIAG] Field: {self.field_name}, Strategy: {self.null_strategy}, Backend: {vis_backend}, Theme: {vis_theme}, Strict: {vis_strict}"
                )

                start_time = time.time()

                try:
                    # Log context variables
                    self.logger.info("[DIAG] Checking context variables...")
                    try:
                        current_context = contextvars.copy_context()
                        self.logger.info(
                            f"[DIAG] Context vars count: {len(list(current_context))}"
                        )
                    except Exception as ctx_e:
                        self.logger.warning(
                            f"[DIAG] Could not inspect context: {ctx_e}"
                        )

                    # Generate visualizations with visualization context parameters
                    self.logger.info("[DIAG] Calling _generate_visualizations...")
                    # Create child progress tracker for visualization if available
                    total_steps = 3  # prepare data, create viz, save
                    viz_progress = None
                    if progress_tracker and hasattr(progress_tracker, "create_subtask"):
                        try:
                            viz_progress = progress_tracker.create_subtask(
                                total=total_steps,
                                description="Generating visualizations",
                                unit="steps",
                            )
                        except Exception as e:
                            self.logger.debug(
                                f"Could not create child progress tracker: {e}"
                            )

                    # Generate visualizations with context parameters
                    visualization_paths = self._generate_visualizations(
                        df=df,
                        metrics=metrics,
                        task_dir=task_dir,
                        vis_theme=vis_theme,
                        vis_backend=vis_backend or "plotly",
                        vis_strict=vis_strict,
                        progress_tracker=viz_progress,
                        timestamp=operation_timestamp,  # Pass the same timestamp
                        **kwargs,
                    )

                    # Close visualization progress tracker
                    if viz_progress:
                        try:
                            viz_progress.close()
                        except:
                            pass

                    elapsed = time.time() - start_time
                    self.logger.info(
                        f"[DIAG] Visualization completed in {elapsed:.2f}s, generated {len(visualization_paths)} files"
                    )

                except Exception as e:
                    elapsed = time.time() - start_time
                    visualization_error = e
                    self.logger.error(
                        f"[DIAG] Visualization failed after {elapsed:.2f}s: {type(e).__name__}: {e}"
                    )
                    self.logger.error("[DIAG] Stack trace:", exc_info=True)

            # Copy context for the thread
            self.logger.info("[DIAG] Preparing to launch visualization thread...")
            ctx = contextvars.copy_context()

            # Create thread with context
            viz_thread = threading.Thread(
                target=ctx.run,
                args=(generate_viz_with_diagnostics,),
                name=f"VizThread-{self.field_name}",
                daemon=False,  # Changed from True to ensure proper cleanup
            )

            self.logger.info(
                f"[DIAG] Starting visualization thread with timeout={vis_timeout}s"
            )
            thread_start_time = time.time()
            viz_thread.start()

            # Log periodic status while waiting
            check_interval = 5  # seconds
            elapsed = 0
            while viz_thread.is_alive() and elapsed < vis_timeout:
                viz_thread.join(timeout=check_interval)
                elapsed = time.time() - thread_start_time
                if viz_thread.is_alive():
                    self.logger.info(
                        f"[DIAG] Visualization thread still running after {elapsed:.1f}s..."
                    )

            if viz_thread.is_alive():
                self.logger.error(
                    f"[DIAG] Visualization thread still alive after {vis_timeout}s timeout"
                )
                self.logger.error(
                    f"[DIAG] Thread state: alive={viz_thread.is_alive()}, daemon={viz_thread.daemon}"
                )
                visualization_paths = {}
            elif visualization_error:
                self.logger.error(
                    f"[DIAG] Visualization failed with error: {visualization_error}"
                )
                visualization_paths = {}
            else:
                total_time = time.time() - thread_start_time
                self.logger.info(
                    f"[DIAG] Visualization thread completed successfully in {total_time:.2f}s"
                )
                self.logger.info(
                    f"[DIAG] Generated visualizations: {list(visualization_paths.keys())}"
                )

        except Exception as e:
            self.logger.error(
                f"[DIAG] Error in visualization thread setup: {type(e).__name__}: {e}"
            )
            self.logger.error("[DIAG] Stack trace:", exc_info=True)
            visualization_paths = {}

        # Register visualization artifacts
        for viz_type, path in visualization_paths.items():
            # Add to result
            result.add_artifact(
                artifact_type="png",
                path=path,
                description=f"{self.field_name} {viz_type} visualization",
                category=Constants.Artifact_Category_Visualization,
            )

            # Report to reporter
            if reporter:
                reporter.add_operation(
                    f"{self.field_name} {viz_type} visualization",
                    details={"artifact_type": "png", "path": str(path)},
                )

        return visualization_paths

    def _generate_visualizations(
        self,
        df: pd.DataFrame,
        metrics: Dict[str, Any],
        task_dir: Path,
        vis_theme: Optional[str] = None,
        vis_backend: Optional[str] = None,
        vis_strict: bool = False,
        progress_tracker: Optional[HierarchicalProgressTracker] = None,
        timestamp: Optional[str] = None,
        **kwargs,
    ) -> Dict[str, Path]:
        """
        Generate visualizations using the core visualization utilities with thread-safe context support.

        This is a base implementation that provides a basic distribution comparison.
        Subclasses should override to provide operation-specific visualizations.

        Parameters
        -----------
        df : pd.DataFrame
            The processed dataframe
        metrics : Dict[str, Any]
            The collected metrics
        task_dir : Path
            Task directory for saving visualizations
        result : OperationResult
            Operation result to add artifacts to
        reporter : Any
            Reporter object for tracking artifacts
        vis_theme : str, optional
            Theme to use for visualizations
        vis_backend : str, optional
            Backend to use: "plotly" or "matplotlib"
        vis_strict : bool, optional
            If True, raise exceptions for configuration errors
        **kwargs : dict
            Additional parameters for the operation

        Returns
        --------
        Dict[str, Path]
            Dictionary with visualization types and paths
        """
        visualization_paths = {}
        viz_dir = task_dir / "visualizations"
        viz_dir.mkdir(parents=True, exist_ok=True)

        # Use provided timestamp or generate new one
        if timestamp is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # Check if visualization should be skipped
        if vis_backend is None:
            self.logger.info(
                f"Skipping visualization for {self.field_name} (backend=None)"
            )
            return visualization_paths

        self.logger.info(
            f"[VIZ] Starting visualization generation for {self.field_name} using {self.null_strategy} strategy"
        )
        self.logger.debug(
            f"[VIZ] Backend: {vis_backend}, Theme: {vis_theme}, Strict: {vis_strict}"
        )

        try:
            # Step 1: Prepare data
            if progress_tracker:
                progress_tracker.update(1, {"step": "Preparing visualization data"})

            if self.mode == "REPLACE":
                gen_series = df[self.field_name]
            elif self.mode == "ENRICH" and self.output_field_name in df.columns:
                gen_series = df[self.output_field_name]
            else:
                gen_series = None

            kwargs_visualization = {
                "use_encryption": self.use_encryption,
                "encryption_key": self.encryption_key,
                "backend": self.visualization_backend,
                "theme": self.visualization_theme,
                "strict": self.visualization_strict,
                "timestamp": timestamp,
            }

            if gen_series is not None:
                visualizations = self._metrics_collector.visualize_metrics(
                    metrics=metrics,
                    field_name=self.field_name,
                    output_dir=viz_dir,
                    op_type=self.operation_name,
                    **kwargs_visualization,
                )
                visualization_paths = {
                    name: str(path) for name, path in visualizations.items()
                }

            # Step 3: Finalize visualizations
            if progress_tracker:
                progress_tracker.update(3, {"step": "Finalizing visualizations"})

        except Exception as e:
            self.logger.warning(f"Error generating visualizations: {e}")

        return visualization_paths

    def _save_output_data(
        self,
        result_df: pd.DataFrame,
        writer: DataWriter,
        result: OperationResult,
        reporter: Any,
        progress_tracker: Optional[HierarchicalProgressTracker],
        timestamp: Optional[str] = None,
        file_name: str = None,
        **kwargs,
    ) -> str:
        """
        Save the processed output data.

        Parameters
        -----------
        result_df : pd.DataFrame
            The processed dataframe to save
        writer : DataWriter
            The writer to use for saving data
        result : OperationResult
            The operation result to add artifacts to
        reporter : Any
            The reporter to log artifacts to
        progress_tracker : Optional[HierarchicalProgressTracker]
            Optional progress tracker
        timestamp : Optional[str]
            Optional timestamp for the operation
        file_name : str
            File name for the name of output file
        **kwargs : dict
            Additional parameters for the operation
        """
        if progress_tracker:
            progress_tracker.update(0, {"step": "Saving output data"})

        # Generate standardized output filename with timestamp
        file_name = file_name or (
            f"{self.field_name}_{self.operation_name}_output_{timestamp}"
        )

        # Use the DataWriter to save the DataFrame
        safe_kwargs = filter_used_kwargs(kwargs, writer.write_dataframe)
        safe_kwargs["encryption_mode"] = get_encryption_mode(
            result_df, self.use_encryption
        )
        output_result = writer.write_dataframe(
            df=result_df,
            name=file_name,
            format=self.output_format,
            subdir="output",
            timestamp_in_name=False,
            encryption_key=self.encryption_key if self.use_encryption else None,
            overwrite=True,
            **safe_kwargs,
        )

        # Register output artifact with the result
        result.add_artifact(
            artifact_type=self.output_format,
            path=output_result.path,
            description=f"{self.field_name} {self.operation_name} generator data",
            category=Constants.Artifact_Category_Output,
        )

        # Save output data types
        self._save_dtypes_output(
            df=result_df,
            writer=writer,
            reporter=reporter,
            result=result,
            file_name=file_name,
        )

        # Report to reporter
        if reporter:
            reporter.add_operation(
                f"{self.field_name} generator data",
                details={
                    "artifact_type": self.output_format,
                    "path": str(output_result.path),
                },
            )
        return str(output_result.path)

    def _cleanup_memory(
        self,
        processed_df: Optional[pd.DataFrame] = None,
        original_data: Optional[pd.Series] = None,
        generated_data: Optional[pd.Series] = None,
    ) -> None:
        """
        Clean up memory after operation completes.

        For large datasets, explicitly free memory by deleting
        references and optionally calling garbage collection.

        Parameters
        -----------
        processed_df : pd.DataFrame, optional
            Processed DataFrame to clear from memory
        original_data : pd.Series, optional
            Original data to clear from memory
        generated_data : pd.Series, optional
            Generated data to clear from memory
        """
        # Delete references
        if processed_df is not None:
            del processed_df
        if original_data is not None:
            del original_data
        if generated_data is not None:
            del generated_data

        # cleanup memory from instance
        helpers.cleanup_memory(instance=self)

    def process_batch(self, batch: pd.DataFrame) -> pd.DataFrame:
        """
        Process a batch of data. Must be implemented by subclasses.

        Parameters
        -----------
        batch : pd.DataFrame
            DataFrame batch to process
        kwargs : dict
            Additional keyword arguments for processing

        Returns
        --------
        pd.DataFrame
            Processed DataFrame batch
        """
        raise FeatureNotImplementedError(
            "Subclasses must implement process_batch method"
        )

    def _check_cache(
        self, df: pd.DataFrame, reporter: Any
    ) -> Optional[OperationResult]:
        """
        Check if a cached result exists for this operation.

        Parameters
        ----------
        df : pd.DataFrame
            DataFrame for the operation
        reporter : Any
            Reporter object for tracking progress and artifacts

        Returns
        -------
        Optional[OperationResult]
            Cached result if found, None otherwise
        """
        if not self.use_cache:
            return None

        try:
            if self.field_name not in df.columns:
                self.logger.warning(
                    f"Field '{self.field_name}' not found in DataFrame."
                )
                return None

            cache_key = self._generate_cache_key(df[self.field_name])
            self.logger.debug(f"Checking cache for key: {cache_key}")

            cached_result = self.operation_cache.get_cache(
                cache_key=cache_key, operation_type=self.operation_name
            )

            if not cached_result:
                self.logger.info("No cached result found, proceeding with operation")
                return None

            result = get_cache_result(cached_result)

            self.logger.info(
                f"Using cached result for {self.field_name} generalization"
            )

            if reporter:
                reporter.add_operation(
                    f"Generalization of {self.field_name} (cached)",
                    details={"null_strategy": self.null_strategy, "cached": True},
                )
            return result

        except Exception as e:
            self.logger.warning(f"Error checking cache: {str(e)}")
            return None

    def _save_to_cache(
        self,
        original_data: pd.Series,
        generated_data: pd.Series,
        result: OperationResult,
        task_dir: Path,
    ) -> bool:
        """
        Save operation results to cache.

        Parameters
        -----------
        original_data : pd.Series
            Original input data
        generated_data : pd.Series
            Generated output data
        result: OperationResult
            Result object OperationResult
        task_dir : Path
            Task directory

        Returns
        --------
        bool
            True if successfully saved to cache, False otherwise
        """
        if not self.use_cache:
            return False

        try:
            # Generate cache key
            cache_key = self._generate_cache_key(original_data)

            # Prepare cache data
            cache_data = build_base_cache(
                parameters=self._get_base_parameters(), result=result
            )

            cache_data.update(
                {
                    "data_info": {
                        "original_length": len(original_data),
                        "generated_length": len(generated_data),
                        "original_null_count": int(original_data.isna().sum()),
                        "generated_null_count": int(generated_data.isna().sum()),
                    }
                }
            )

            # Save to cache
            self.logger.debug(f"Saving to cache with key: {cache_key}")

            success = self.operation_cache.save_cache(
                data=cache_data,
                cache_key=cache_key,
                operation_type=self.operation_name,
                metadata={"task_dir": str(task_dir)},
            )

            if success:
                self.logger.info(
                    f"Successfully saved {self.field_name} generator results to cache"
                )
            else:
                self.logger.warning(
                    f"Failed to save {self.field_name} generator results to cache"
                )

            return success

        except Exception as e:
            self.logger.warning(f"Error saving to cache: {str(e)}")
            return False

    def _get_base_parameters(self) -> Dict[str, str]:
        """Get the basic parameters for the cache key generation."""
        # Get basic operation parameters

        parameters = super()._get_base_parameters()

        parameters.update(
            {
                "field_name": self.field_name,
                "output_field_name": self.output_field_name,
                "null_strategy": self.null_strategy,
            }
        )

        return parameters

    def _save_dtypes_output(
        self,
        df: pd.DataFrame,
        writer: DataWriter,
        result: OperationResult,
        reporter: Any,
        file_name: str = None,
    ) -> bool:
        """
        Saves data types dataframe format to a JSON file.

        Returns
        -------
        Path or None
            Path to saved file if success, otherwise None
        """
        try:

            # Get the dtypes of the columns as a Series
            dtypes_series = df.dtypes

            # Convert the dtypes Series to a dictionary
            dtypes_dict = dtypes_series.astype(str).to_dict()

            # Generate standardized output filename with timestamp
            dtypes_filename = f"data_types_{file_name}"

            dtypes_result = writer.write_metrics(
                metrics=dtypes_dict,
                name=dtypes_filename,
                timestamp_in_name=False,
                encryption_key=self.encryption_key,
            )

            result.add_metric(dtypes_filename, dtypes_dict)

            result.add_artifact(
                artifact_type="json",
                path=dtypes_result.path,
                description=f"Data types of output {self.field_name} {self.operation_name}",
                category=Constants.Artifact_Category_Metrics,
            )

            if reporter:
                reporter.add_artifact(
                    artifact_type="json",
                    path=str(dtypes_result.path),
                    description=f"Data types of output {self.field_name} {self.operation_name}",
                )

            self.logger.info(f"Dtypes output saved to: {Path(dtypes_result.path).name}")
            return True

        except Exception as e:
            self.logger.warning(f"Failed to save dtypes format: {str(e)}")
            return False

    def _save_metrics(
        self,
        metrics: Dict[str, Any],
        writer: DataWriter,
        result: OperationResult,
        reporter: Any,
        progress_tracker: Optional[HierarchicalProgressTracker],
        file_name: str = None,
        operation_timestamp: Optional[str] = None,
    ) -> bool:
        """
        Save metrics.

        Parameters
        -----------
        metrics : dict
            The metrics of operation
        writer : DataWriter
            The writer to use for saving data
        result : OperationResult
            The operation result to add artifacts to
        reporter : Any
            The reporter to log artifacts to
        progress_tracker : Optional[HierarchicalProgressTracker]
            Optional progress tracker
        operation_timestamp : str, optional
            Timestamp of the operation, if any (for filename purposes)
        file_name : Str
            The file name

        Returns
        -------
        True or False

        """
        try:
            if progress_tracker:
                progress_tracker.update(0, {"step": "Saving metrics"})

            # Use the DataWriter to save
            metrics_filename = file_name or (
                f"{self.field_name}_{self.operation_name}_metrics_{operation_timestamp}"
            )

            # Save metrics using writer
            metrics_result = writer.write_metrics(
                metrics=metrics,
                name=metrics_filename,
                timestamp_in_name=False,
                encryption_key=(self.encryption_key if self.use_encryption else None),
            )

            # Add metrics to result
            for key, value in metrics.items():
                if isinstance(value, (int, float, str, bool)):
                    result.add_metric(key, value)

            # Register metrics artifact
            result.add_artifact(
                artifact_type="json",
                path=metrics_result.path,
                description=f"{self.operation_name.lower()} {self.field_name} generator metrics",
                category=Constants.Artifact_Category_Metrics,
            )

            # Report artifact
            if reporter:
                reporter.add_artifact(
                    artifact_type="json",
                    path=str(metrics_result.path),
                    description=f"{self.operation_name.lower()} {self.field_name} generator metrics",
                )
        except Exception as e:
            self.logger.warning(f"Failed to save metrics: {str(e)}")
            return False
