import pandas as pd


def add_battery_features(
    df: pd.DataFrame,
    price_df: pd.DataFrame,
    battery_col: str = "charge_level_mwh",
    capacity_col: str = "capacity_mwh",
    timestamp_col: str = "timestamp",
    battery_id_col: str = "battery_id",
) -> pd.DataFrame:
    result = df.copy()

    if battery_col in result.columns and capacity_col in result.columns:
        result["charge_level_pct"] = result[battery_col] / result[capacity_col]
        result["discharge_potential_mwh"] = result[battery_col] * result.get("efficiency", 0.9)
        result["charge_capacity_mwh"] = result[capacity_col] - result[battery_col]

    if price_df is not None and "real_time_price" in price_df.columns and timestamp_col in result.columns:
        merged = result.merge(
            price_df[[timestamp_col, "real_time_price"]],
            on=timestamp_col,
            how="left",
            suffixes=("", "_market")
        )

        if "real_time_price_market" in merged.columns:
            result["market_price"] = merged["real_time_price_market"]
            result["charge_cost_potential"] = result["charge_capacity_mwh"] * result["market_price"]
            result["discharge_revenue_potential"] = result["discharge_potential_mwh"] * result["market_price"]

    return result


__all__ = ["add_battery_features"]