from sklearn.base import BaseEstimator, TransformerMixin


class CalculateDatetimeDifferenceTransformer(BaseEstimator, TransformerMixin):
    """
    Compute the time difference between two datetime columns.

    Parameters
    ----------
    start_column : str
        Name of the start datetime column.
    end_column : str
        Name of the end datetime column.
    new_column : str
        Name of the output feature.
    unit : str, default="hours"
        Unit of the time difference.
        One of: "seconds", "minutes", "hours", "days".
    """

    def __init__(self, start_column, end_column, new_column, unit="hours"):
        self.start_column = start_column
        self.end_column = end_column
        self.new_column = new_column
        self.unit = unit

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

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

        if self.start_column not in X.columns or self.end_column not in X.columns:
            raise ValueError("Start or end datetime column not found in DataFrame")

        delta = X[self.end_column] - X[self.start_column]

        if self.unit == "seconds":
            X[self.new_column] = delta.dt.total_seconds()
        elif self.unit == "minutes":
            X[self.new_column] = delta.dt.total_seconds() / 60
        elif self.unit == "hours":
            X[self.new_column] = delta.dt.total_seconds() / 3600
        elif self.unit == "days":
            X[self.new_column] = delta.dt.total_seconds() / 86400
        else:
            raise ValueError(f"Unsupported unit: {self.unit}")

        return X
