from sklearn.base import BaseEstimator, TransformerMixin


class ValueUpdateConditionalTransformer(BaseEstimator, TransformerMixin):
    """
    Sets or updates a column based on a condition applied to another column.
    Existing values are preserved if the condition is not met.

    Parameters
    ----------
    condition_column : str
        Name of the column to check the condition against.
    condition_values : list or any
        The value(s) in `condition_column` that trigger the assignment.
    target_column : str
        Name of the column to create or update.
    new_value : any
        The value to set in `target_column` when the condition is met.
    """
    def __init__(self, condition_column, condition_values, target_column, new_value):
        self.condition_column = condition_column
        self.condition_values = condition_values
        self.target_column = target_column
        self.new_value = new_value

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        X = X.copy()

        cond_vals = self.condition_values if isinstance(self.condition_values, list) else [self.condition_values]

        mask = X[self.condition_column].isin(cond_vals)

        X.loc[mask, self.target_column] = self.new_value

        return X
