"""
PAMOLA.CORE - Privacy-Preserving AI Data Processors
Module: Base Transformation Operation
Description: Base class for all data transformation operations in PAMOLA Core.

Author: PAMOLA Core Team
Created: 2025
License: BSD 3-Clause

This module provides the foundational class for all transformation operations,
defining common interfaces, behaviors, and utilities for data transformation tasks.

Key features:
- Standardized operation lifecycle: validation, execution, result handling
- Support for in-place (REPLACE) and enrichment (ENRICH) transformation modes
- Configurable batch processing for large datasets
- Optional caching, Dask-based distributed processing, and output encryption
- Integrated progress tracking and artifact reporting
- Flexible handling of output field naming and operation metadata
- Hooks for metrics calculation and visualization generation
"""

import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
import pandas as pd
from pamola_core.common.constants import Constants
from pamola_core.utils.io import load_data_operation, load_settings_operation
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_base import BaseOperation
from pamola_core.utils.ops.op_result import OperationResult, OperationStatus
from pamola_core.utils.progress import HierarchicalProgressTracker
from pamola_core.errors.codes import ErrorCode
from pamola_core.errors.error_handler import ErrorHandler
from pamola_core.errors.exceptions import (
    FeatureNotImplementedError,
    TypeValidationError,
    ValidationError,
)
from pamola_core.transformations.commons.processing_utils import (
    process_dataframe_with_config,
)
from pamola_core.transformations.commons.visualization_utils import (
    generate_dataset_overview_vis,
    generate_data_distribution_comparison_vis,
    generate_field_count_comparison_vis,
    generate_record_count_comparison_vis,
    sample_large_dataset,
)
from pamola_core.utils.io_helpers.crypto_utils import get_encryption_mode
import pamola_core.utils.helpers as helpers


class TransformationOperation(BaseOperation):
    """Base class for all transformation operations.

    This class provides common functionality for all transformation operations,
    including data source handling, result processing, and metric generation.
    """

    def __init__(
        self,
        name: str = "transformation_operation",
        field_name: Optional[str] = None,
        **kwargs,
    ):
        """Initialize the transformation operation.

        Parameters
        -----------
        name : str
            Name of the operation (default: "transformation_operation")
        field_name : str, optional
            Field name to transform (default: None, meaning all applicable fields)
        **kwargs : dict
            Additional parameters passed to :class:`BaseOperation`.
        """
        # --- Determine operation naming ---
        field_label = field_name or "all_fields"
        if "operation" not in name:
            operation_name = f"{field_label}_{name}" if field_name else name
        else:
            operation_name = name

        # Set default metadata if not provided
        kwargs.setdefault("name", operation_name)
        kwargs.setdefault(
            "description",
            f"Transformation operation applied to '{field_label}'",
        )

        # Initialize BaseOperation with remaining parameters
        super().__init__(**kwargs)

        # --- Core transformation attributes ---
        self.field_name = field_name
        self.field_label = field_label

    def execute(
        self,
        data_source: DataSource,
        task_dir: Path,
        reporter: Any,
        progress_tracker: Optional[HierarchicalProgressTracker] = None,
        **kwargs,
    ) -> OperationResult:
        """
        Execute the transformation operation.

        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)

            # Initialize dataframe
            df = None

            # Create a field List for display purposes
            field_info = (
                {"field": self.field_name}
                if self.field_name
                else {"field_label": self.field_label}
            )

            # 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, 2. Validation, 3. Cache check, 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.operation_name}",
                            **field_info,
                        },
                    )
                except Exception as e:
                    self.logger.warning(f"Could not update progress tracker: {e}")

            # Step 1: Data Loading
            if main_progress:
                current_steps += 1
                main_progress.update(
                    current_steps, {"step": "Data Loading", **field_info}
                )

            # Validate and get dataframe
            try:
                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: Validation
            if main_progress:
                current_steps += 1
                main_progress.update(
                    current_steps, {"step": "Validation", **field_info}
                )

            try:
                output_fields = self._prepare_output_fields(df)
                self._report_operation_details(reporter, output_fields)

                original_data = (
                    df[[self.field_name]].copy(deep=True)
                    if self.field_name
                    else df.copy(deep=True)
                )

            except Exception as e:
                return self.error_handler.handle_error(
                    error=e,
                    error_code=ErrorCode.PROCESSING_FAILED,
                    context={
                        "step": "_prepare_output_fields",
                        "field": self.field_label,
                    },
                    message_kwargs={
                        "field_name": self.field_label,
                        "operation": self.operation_name,
                        "reason": str(e),
                    },
                )

            # Step 3: 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_info},
                    )

                self.logger.info("Checking operation cache...")
                cache_result = self._check_cache(df, reporter)

                if cache_result:
                    self.logger.info("Cache hit! Using cached results.")

                    # Update progress
                    if main_progress:
                        main_progress.update(
                            current_steps,
                            {"step": "Complete (cached)", **field_info},
                        )

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

                    return cache_result

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

            try:
                # 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}"
                        )

                result_df = self._process_dataframe(
                    df=df, progress_tracker=data_tracker
                )
                transformed_data = result_df[output_fields]

            except Exception as e:
                return self.error_handler.handle_error(
                    error=e,
                    error_code=ErrorCode.PROCESSING_FAILED,
                    context={"step": "processing", "field": self.field_label},
                    message_kwargs={
                        "field_name": self.field_label,
                        "operation": self.operation_name,
                        "reason": str(e),
                    },
                )

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

            # Record end time after processing
            self.end_time = time.time()

            # Initialize metrics dictionary
            metrics = {}

            try:
                # Calculate all relevant metrics based on original and transformed data
                metrics = self._calculate_all_metrics(original_data, transformed_data)

                # Generate metrics file name
                metrics_file_name = f"{self.field_label}_{self.operation_name}_metrics_{operation_timestamp}"

                self._save_metrics(
                    metrics=metrics,
                    writer=writer,
                    result=result,
                    reporter=reporter,
                    progress_tracker=progress_tracker,
                    operation_timestamp=operation_timestamp,
                    file_name=metrics_file_name,
                )

            except Exception as e:
                error_message = f"Error calculating metrics: {str(e)}"
                self.logger.warning(error_message)
                # Continue execution since metrics failure is not critical

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

            # 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(
                        original_data=original_data,
                        transformed_data=transformed_data,
                        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_info},
                )
            # Save output data if required
            if self.save_output:
                try:
                    self._save_output_data(
                        result_df=result_df,
                        writer=writer,
                        result=result,
                        reporter=reporter,
                        progress_tracker=main_progress,
                        timestamp=operation_timestamp,
                        **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_label},
                        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,
                        transformed_data=transformed_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(
                result_df=result_df,
                original_data=original_data,
                transformed_data=transformed_data,
            )

            # Report completion
            if reporter:
                # Create the details dictionary with checks for all values
                details = {
                    "records_processed": self.process_count,
                    "execution_time": self.execution_time,
                }

                # Only add generalization_ratio if metrics exists and has this key
                if metrics and isinstance(metrics, dict):
                    generalization_ratio = metrics.get("generalization_ratio")
                    if generalization_ratio is not None:
                        details["generalization_ratio"] = generalization_ratio

                # Add the operation to the reporter
                reporter.add_operation(
                    f"Transformation of {field_info} completed",
                    details=details,
                )

            # 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_label},
                message_kwargs={
                    "field_name": self.field_label,
                    "operation": self.operation_name,
                    "reason": str(e),
                },
            )

    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 and load data from.
        dataset_name : str
            The name of the dataset to load.

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

        Raises
        ------
        ValueError
            If no valid dataframe is found or the specified field is missing.
        """
        # Attempt to get the main DataFrame from the data source
        df = load_data_operation(data_source, dataset_name, **kwargs)

        if df is None:
            error_msg = (
                f"Data source '{data_source}' does not contain a valid DataFrame"
            )
            raise ValidationError(error_msg)

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

        # 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_fields(self, df: pd.DataFrame) -> List[str]:
        """
        Prepare output field names based on whether `field_name` is specified.

        Parameters
        ----------
        df : pd.DataFrame
            The input DataFrame whose columns are used to determine output fields.

        Returns
        -------
        List[str]
            List of output field names:
            - If `field_name` is empty, returns all columns (for processing all).
            - If `mode` is "REPLACE", returns a list containing the `field_name`.
            - Otherwise (mode "ENRICH"), returns a list with the output field name,
            either specified or generated by prefixing `field_name`.
        """
        if not self.field_name:
            self.logger.debug(
                "No field_name provided; returning all DataFrame columns as output fields."
            )
            return df.columns.tolist()

        if self.mode == "REPLACE":
            self.logger.debug(
                f"Mode REPLACE: output fields set to ['{self.field_name}']."
            )
            return [self.field_name]

        output_field = (
            self.output_field_name or f"{self.column_prefix}{self.field_name}"
        )
        if output_field in df.columns:
            self.logger.warning(
                f"Output field '{output_field}' already exists in DataFrame and will be overwritten."
            )

        self.logger.debug(f"Mode ENRICH: output field set to ['{output_field}'].")
        return [output_field]

    def _report_operation_details(
        self, reporter: Optional[Any], output_fields: List[str]
    ) -> None:
        """
        Report details of the operation to the reporter.

        Parameters
        ----------
        reporter : Optional[Any]
            The reporter to log details to. If None, reporting is skipped.
        output_fields : List[str]
            The list of output field names.
        """
        if not reporter:
            self.logger.debug("No reporter provided, skipping reporting.")
            return

        operation_desc = (
            f"Transformation field: {self.field_name}"
            if self.field_name
            else "Transformation on all applicable fields"
        )

        reporter.add_operation(
            operation_desc,
            details={
                "field_name": self.field_name or "N/A",
                "output_fields": (
                    output_fields[0]
                    if len(output_fields) == 1
                    else ", ".join(output_fields)
                ),
                "mode": self.mode,
                "operation_type": self.operation_name,
            },
        )

    def _process_dataframe(
        self, df: pd.DataFrame, progress_tracker: Optional[HierarchicalProgressTracker]
    ) -> 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

        processed_df = process_dataframe_with_config(
            df=df,
            process_function=self.process_batch,
            chunk_size=self.chunk_size,
            use_dask=self.use_dask,
            npartitions=self.npartitions,
            meta=None,
            use_vectorization=self.use_vectorization,
            parallel_processes=self.parallel_processes,
            progress_tracker=progress_tracker,
            task_logger=self.logger,
        )

        return processed_df

    def _calculate_all_metrics(
        self,
        original_data: Union[pd.Series, pd.DataFrame],
        transformed_data: Union[pd.Series, pd.DataFrame],
    ) -> Dict[str, Any]:
        """
        Calculate all metrics for the transformation operation.

        Supports both single-field (Series) and multi-field (DataFrame) inputs.

        Parameters
        ----------
        original_data : Union[pd.Series, pd.DataFrame]
            The original data before transformation.
        transformed_data : Union[pd.Series, pd.DataFrame]
            The transformed data after processing.

        Returns
        -------
        Dict[str, Any]
            A dictionary containing calculated metrics.
        """
        # Calculate basic metrics (e.g., accuracy, difference, etc.)
        metrics = self._collect_metrics(original_data, transformed_data)

        # Calculate execution time in seconds, safely handling missing timestamps
        self.execution_time = (
            self.end_time - self.start_time
            if self.end_time is not None and self.start_time is not None
            else 0
        )

        # Calculate processing rate (records per second), avoid division by zero
        records_per_second = (
            self.process_count / self.execution_time if self.execution_time > 0 else 0
        )

        # Update metrics dictionary with performance metrics
        metrics.update(
            {
                "execution_time_seconds": self.execution_time,
                "records_processed": self.process_count,
                "records_per_second": records_per_second,
                "processing_date": datetime.now(),
            }
        )

        return metrics

    def _handle_visualizations(
        self,
        original_data: pd.Series,
        transformed_data: pd.Series,
        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
        -----------
        original_data : pd.Series
            The original data before transformation
        transformed_data : pd.Series
            The transformed data after processing
        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 : Any
            Additional keyword arguments for visualization functions.
        """
        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_label}, 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(
                        original_data=original_data,
                        transformed_data=transformed_data,
                        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_label}",
                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_label} {viz_type} visualization",
                category=Constants.Artifact_Category_Visualization,
            )

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

        return visualization_paths

    def _generate_visualizations(
        self,
        original_data: Union[pd.Series, pd.DataFrame],
        transformed_data: Union[pd.Series, pd.DataFrame],
        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: Any,
    ) -> Dict[str, Path]:
        """
        Generate visualizations and metadata summaries using all core visualization utilities.

        Parameters
        ----------
        original_data : pd.Series or pd.DataFrame
            Original data before transformation.
        transformed_data : pd.Series or pd.DataFrame
            Transformed data after processing.
        task_dir : Path
            Task directory for saving visualizations
        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
        progress_tracker : Optional[HierarchicalProgressTracker]
            Progress tracker for monitoring visualization generation
        timestamp : Optional[str]
            Timestamp for naming visualization files
        **kwargs : Any
            Additional keyword arguments for visualization functions.

        Returns
        --------
        Dict[str, Path]
            Dictionary with visualization types and paths
        """
        operation_name = self.operation_name
        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_label} (backend=None)"
            )
            return visualization_paths

        self.logger.info(
            f"[VIZ] Starting visualization generation for {self.field_label}"
        )
        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"})

            # Sample large datasets for visualization
            if len(original_data) > 10000:
                self.logger.info(
                    f"[VIZ] Sampling large dataset: {len(original_data)} -> 10000 samples"
                )
                original_for_viz = sample_large_dataset(
                    original_data, max_samples=10000
                )
                transformed_for_viz = sample_large_dataset(
                    transformed_data, max_samples=10000
                )
            else:
                original_for_viz = original_data
                transformed_for_viz = transformed_data

            self.logger.debug(
                f"[VIZ] Data prepared for visualization: {len(original_for_viz)} samples"
            )

            # Log dtype info depending on whether data is Series or DataFrame
            if isinstance(original_for_viz, pd.Series) and isinstance(
                transformed_for_viz, pd.Series
            ):
                self.logger.debug(
                    f"[VIZ] Original data type: {original_for_viz.dtype}, "
                    f"transformed data type: {transformed_for_viz.dtype}"
                )

            elif isinstance(original_for_viz, pd.DataFrame) and isinstance(
                transformed_for_viz, pd.DataFrame
            ):
                original_cols = list(original_for_viz.dtypes.to_dict().items())
                transformed_cols = list(transformed_for_viz.dtypes.to_dict().items())
                self.logger.debug(
                    f"[VIZ] Original DataFrame columns: {original_cols}, "
                    f"transformed DataFrame columns: {transformed_cols}"
                )

            # Convert original and transformed Series to DataFrames (if Series)
            original_df: pd.DataFrame = (
                original_for_viz.to_frame()
                if isinstance(original_for_viz, pd.Series)
                else original_for_viz
            )
            transformed_df: pd.DataFrame = (
                transformed_for_viz.to_frame()
                if isinstance(transformed_for_viz, pd.Series)
                else transformed_for_viz
            )

            # Step 2: Create visualization
            if progress_tracker:
                progress_tracker.update(2, {"step": "Creating visualization"})

            # Generate visualizations for each column (supports both Series and DataFrame)
            if original_data is not None and transformed_data is not None:
                for column in original_df.columns:
                    # Handle ENRICH mode column renaming
                    transformed_column = column
                    if self.mode == "ENRICH":
                        transformed_column = (
                            self.output_field_name or f"{self.column_prefix}{column}"
                        )
                        self.logger.debug(
                            f"Mode ENRICH: output field set to ['{transformed_column}']."
                        )

                    # Generate field count comparison visualization

                    visualization_paths.update(
                        generate_field_count_comparison_vis(
                            original_df=original_df[[column]],
                            transformed_df=transformed_df[[transformed_column]],
                            field_label=column,
                            operation_name=operation_name,
                            task_dir=viz_dir,
                            timestamp=timestamp,
                            theme=vis_theme,
                            backend=vis_backend,
                            strict=vis_strict,
                            visualization_paths=visualization_paths,
                            **kwargs,
                        )
                    )

                    # Record count comparison for each column
                    transformed_dfs: Dict[str, pd.DataFrame] = {
                        column: transformed_df[[transformed_column]]
                    }
                    visualization_paths.update(
                        generate_record_count_comparison_vis(
                            original_df=original_df[[column]],
                            transformed_dfs=transformed_dfs,
                            field_label=column,
                            operation_name=operation_name,
                            task_dir=viz_dir,
                            timestamp=timestamp,
                            theme=vis_theme,
                            backend=vis_backend,
                            strict=vis_strict,
                            visualization_paths=visualization_paths,
                            **kwargs,
                        )
                    )

                    # Data distribution comparison for each column
                    visualization_paths.update(
                        generate_data_distribution_comparison_vis(
                            original_df=original_df[column],
                            transformed_data=transformed_df[transformed_column],
                            field_label=column,
                            operation_name=operation_name,
                            task_dir=viz_dir,
                            timestamp=timestamp,
                            theme=vis_theme,
                            backend=vis_backend,
                            strict=vis_strict,
                            visualization_paths=visualization_paths,
                            **kwargs,
                        )
                    )

                # Dataset overview visualizations for original
                if original_df is not None:
                    visualization_paths.update(
                        generate_dataset_overview_vis(
                            df=original_df,
                            operation_name=operation_name,
                            dataset_type="original",
                            field_label=self.field_label,
                            task_dir=viz_dir,
                            timestamp=timestamp,
                            theme=vis_theme,
                            backend=vis_backend,
                            strict=vis_strict,
                            visualization_paths=visualization_paths,
                            **kwargs,
                        )
                    )

                # Dataset overview visualizations for transformed
                if transformed_df is not None:
                    visualization_paths.update(
                        generate_dataset_overview_vis(
                            df=transformed_df,
                            operation_name=operation_name,
                            dataset_type="transformed",
                            field_label=self.field_label,
                            task_dir=viz_dir,
                            timestamp=timestamp,
                            theme=vis_theme,
                            backend=vis_backend,
                            strict=vis_strict,
                            visualization_paths=visualization_paths,
                            **kwargs,
                        )
                    )

            # 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 output file
        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"})

        custom_kwargs = self._get_custom_kwargs(result_df, **kwargs)

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

        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,
            **custom_kwargs,
        )

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

        # Report to reporter
        if reporter:
            reporter.add_artifact(
                self.output_format,
                str(output_result.path),
                f"{self.operation_name} transformed data",
            )

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

        return str(output_result.path)

    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 to check in the cache
        reporter : Any
            Reporter to log cache hits/misses

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

        try:

            cache_key = self._generate_cache_key(df)
            # Check for cached result
            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 = helpers.get_cache_result(cached_result)

            if reporter:
                reporter.add_operation(
                    f"Transformation of {self.field_label} (cached)",
                    details={"field_label": self.field_label, "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: Union[pd.Series, pd.DataFrame],
        transformed_data: Union[pd.Series, pd.DataFrame, Dict[str, pd.DataFrame]],
        result: OperationResult,
        task_dir: Path,
    ) -> bool:
        """Save operation results to cache."""
        if not self.use_cache:
            return False

        try:
            cache_key = self._generate_cache_key(original_data)
            operation_params = self._get_base_parameters()

            self.logger.debug(f"Operation parameters for cache: {operation_params}")

            cache_data = helpers.build_base_cache(
                parameters=self._get_base_parameters(), result=result
            )

            # Handle dict case
            if isinstance(transformed_data, dict):
                # Aggregate info from all DataFrames in dict
                total_length = sum(len(df) for df in transformed_data.values())
                total_null_count = sum(
                    df.isna().sum().sum() for df in transformed_data.values()
                )

                cache_data.update(
                    {
                        "data_info": {
                            "original_length": len(original_data),
                            "transformed_length": total_length,
                            "original_null_count": int(
                                original_data.isna().sum().sum()
                                if isinstance(original_data, pd.DataFrame)
                                else original_data.isna().sum()
                            ),
                            "transformed_null_count": int(total_null_count),
                            "num_dataframes": len(transformed_data),  # Extra info
                        }
                    }
                )
            else:
                # Original logic for Series/DataFrame
                cache_data.update(
                    {
                        "data_info": {
                            "original_length": len(original_data),
                            "transformed_length": len(transformed_data),
                            "original_null_count": int(
                                original_data.isna().sum().sum()
                                if isinstance(original_data, pd.DataFrame)
                                else original_data.isna().sum()
                            ),
                            "transformed_null_count": int(
                                transformed_data.isna().sum().sum()
                                if isinstance(transformed_data, pd.DataFrame)
                                else transformed_data.isna().sum()
                            ),
                        }
                    }
                )

            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_label} transformation results to cache"
                )
            else:
                self.logger.warning(
                    f"Failed to save {self.field_label} transformation results to cache"
                )

            return success

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

    def _cleanup_memory(
        self,
        result_df: Union[pd.Series, pd.DataFrame, Dict[str, pd.DataFrame]],
        original_data: Optional[Union[pd.Series, pd.DataFrame]] = None,
        transformed_data: Optional[Union[pd.Series, pd.DataFrame]] = None,
    ) -> None:
        """
        Clean up memory after operation completes.

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

        Parameters
        -----------
        result_df : pd.DataFrame, optional
            Processed DataFrame to clear from memory
        original_data: Union[pd.Series, pd.DataFrame], optional
            Original data to clear from memory
        transformed_data: Union[pd.Series, pd.DataFrame], optional
            Transformed data to clear from memory
        """
        # Delete references
        if result_df is not None:
            del result_df
        if original_data is not None:
            del original_data
        if transformed_data is not None:
            del transformed_data

        # Clear right_df if exists
        if hasattr(self, "right_df"):
            self.right_df = None

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

    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()

        # Add operation-specific parameters
        parameters.update(
            {
                "field_label": self.field_label,
                "field_name": self.field_name,
            }
        )
        return parameters

    def process_batch(self, batch: pd.DataFrame) -> pd.DataFrame:
        """
        Process a batch of data.

        Parameters
        -----------
        batch : pd.DataFrame
            DataFrame batch to process

        Returns
        --------
        pd.DataFrame
            Processed DataFrame batch
        """
        # This method should be implemented by subclasses
        raise FeatureNotImplementedError(
            "Subclasses must implement process_batch method"
        )

    def _collect_metrics(
        self,
        original_data: Union[pd.Series, pd.DataFrame],
        transformed_data: Union[pd.Series, pd.DataFrame],
    ) -> Dict[str, Any]:
        """
        Collect universal and transformation-specific metrics.

        Parameters
        ----------
        original_data : pd.Series or pd.DataFrame
            Original data before transformation.
        transformed_data : pd.Series or pd.DataFrame
            Transformed data after processing.

        Returns
        -------
        Dict[str, Any]
            Dictionary with collected metrics.
        """
        # Ensure input is DataFrame for consistent metric calculation
        original_df = (
            original_data.to_frame()
            if isinstance(original_data, pd.Series)
            else original_data
        )
        transformed_df = (
            transformed_data.to_frame()
            if isinstance(transformed_data, pd.Series)
            else transformed_data
        )

        total_input_records = len(original_df)
        total_output_records = len(transformed_df)
        total_input_fields = original_df.shape[1]
        total_output_fields = transformed_df.shape[1]

        # Universal metrics
        metrics: Dict[str, Any] = {
            "total_input_records": total_input_records,
            "total_output_records": total_output_records,
            "total_input_fields": total_input_fields,
            "total_output_fields": total_output_fields,
            "transformation_type": self.operation_name,
            "execution_time_seconds": None,
            "records_per_second": None,
            "records_processed": None,
            "processing_date": None,
            "operation_type": self.operation_name,
            "field_name": self.field_label,
            "mode": self.mode,
            "null_count": int(original_df.isna().sum().sum()),
            "unique_values_before": int(original_df.nunique().sum()),
            "unique_values_after": int(transformed_df.nunique().sum()),
        }

        # Add transformation-specific metrics
        specific_metrics = self._collect_specific_metrics(
            original_data, transformed_data
        )
        if specific_metrics:
            metrics.update(specific_metrics)

        return metrics

    def _collect_specific_metrics(
        self,
        original_data: Union[pd.Series, pd.DataFrame],
        transformed_data: Union[pd.Series, pd.DataFrame],
    ) -> Dict[str, Any]:
        """
        Collect operation-specific metrics.

        Parameters
        ----------
        original_data : pd.Series or pd.DataFrame
            Original data before transformation.
        transformed_data : pd.Series or pd.DataFrame
            Transformed data after processing.

        Returns
        --------
        Dict[str, Any]
            Dictionary with operation-specific metrics
        """
        # This method should be overridden by subclasses to add specific metrics
        return {}

    def _get_custom_kwargs(self, df, **kwargs):
        custom_kwargs = {
            k: v
            for k, v in kwargs.items()
            if k
            not in [
                "df",
                "name",
                "format",
                "subdir",
                "timestamp_in_name",
                "encryption_key",
            ]
        }
        custom_kwargs["encryption_mode"] = get_encryption_mode(df, self.use_encryption)

        return custom_kwargs

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

        Returns
        -------
        True or False
        """
        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_label} {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_label} {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.operation_name.lower()}_metrics_{operation_timestamp}"
            )

            metrics_result = writer.write_metrics(
                metrics=metrics,
                name=metrics_filename,
                timestamp_in_name=False,  # Already included in the filename
                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()} 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()} metrics",
                )

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