#----------------------------------------------------------------------------------------------
# This is a custom LQP rewriter that filters extra columns from CSV export. It is used in the
# LQP executor, when the format="csv", to ensure only intended columns are being exported.
#----------------------------------------------------------------------------------------------

from dataclasses import replace
from relationalai.semantics.lqp import lqp_bundle

class PartitioningDisabler:
    """Disables output file partitioning for CSV exports by setting partition_size=-1."""

    def rewrite_read(self, read: lqp_bundle.ir.Read) -> lqp_bundle.ir.Read:
        if not isinstance(read.read_type, lqp_bundle.ir.Export):
            return read
        config = read.read_type.config
        new_config = replace(config, partition_size=-1)
        new_read_type = replace(read.read_type, config=new_config)
        return replace(read, read_type=new_read_type)

class ExtraColumnsFilter:

    def __init__(self, original_cols: list[str]):
        self.original_cols = set(original_cols)

    def filter_read(self, read: lqp_bundle.ir.Read) -> lqp_bundle.ir.Read:
        """Filter data_columns to only include columns in original_cols."""
        if not isinstance(read.read_type, lqp_bundle.ir.Export):
            return read

        config = read.read_type.config
        assert isinstance(config, lqp_bundle.ir.ExportCSVConfig) and config.data_columns is not None, \
            "Expected ExportCSVConfig with data_columns in the read type"

        data_columns = config.data_columns

        new_data_columns = [col for col in data_columns if col.column_name in self.original_cols]

        # Reconstruct the nested structure with filtered data_columns
        new_config = replace(config, data_columns=new_data_columns)
        new_read_type = replace(read.read_type, config=new_config)
        new_read = replace(read, read_type=new_read_type)

        return new_read
