import datamazing.pandas as pdz
import pandas as pd

from .outage_manager import OutageManager


class AvailabilityManager:
    """
    Manager which simplifies the process of getting availability data.
    """

    def __init__(
        self,
        db: pdz.Database,
        time_interval: pdz.TimeInterval,
        resolution: pd.Timedelta,
    ) -> None:
        self.db = db
        self.time_interval = time_interval
        self.resolution = resolution
        self.outage_manager = OutageManager(db, time_interval, resolution)

    def get_availability(self) -> pd.DataFrame:
        df_availability = self.db.query("schedule_availability", self.time_interval)
        df_outage = self.outage_manager.get_plant_outage_time_series()

        # Rename resource_gsrn to plant_gsrn for consistent merge
        df_availability = df_availability.rename(
            columns={
                "resource_gsrn": "plant_gsrn",
                "schedule_resource_capacity_max_MW": "capacity_max_MW",
                "schedule_resource_capacity_min_MW": "capacity_min_MW",
            }
        )

        required_columns = [
            "time_utc",
            "plant_gsrn",
            "capacity_max_MW",
            "capacity_min_MW",
        ]

        # Create an indicator column in outage to mark where outages exist
        df_outage["has_outage"] = True

        # Merge availability with outage data
        df_available = df_availability.merge(
            df_outage, on=["plant_gsrn", "time_utc"], how="outer"
        )

        # Where there's an outage, set capacity to zero
        has_outage = df_available["has_outage"].fillna(False)
        df_available.loc[has_outage, "capacity_max_MW"] = 0
        df_available.loc[has_outage, "capacity_min_MW"] = 0

        # Return only the required columns
        df_available = df_available[required_columns]
        return df_available
