"""Convert simple Arcadia flat demand charges into URDB flat demand fields."""

from datetime import datetime
from math import inf

from tariff_fetch.arcadia.schema.tariffrate import TariffRateBand, TariffRateExtended
from tariff_fetch.urdb.schema import FlatDemandTier, URDBRate

from . import rateutils as ru
from .exception import RateConversionError
from .library import Library
from .scenario import Scenario

_SUPPORTED_DEMAND_UNIT = "kW"


def build_flat_demand(scenario: Scenario, library: Library) -> URDBRate:
    """Build URDB flat-demand fields for a scenario year.

    This first implementation intentionally supports only simple monthly flat
    demand charges:

    - Arcadia `DEMAND_BASED` rates
    - monthly charge period
    - no TOU structure
    - no variable factors, variable limits, or property-driven formulas
    - BUY transaction type only
    - rate bands expressed as demand tiers
    """

    monthly_band_sets = [get_month_flat_demand_bands(scenario, library, month) for month in range(1, 13)]

    if all(not bands for bands in monthly_band_sets):
        return {}

    band_sets: list[tuple[FlatDemandTier, ...]] = []
    band_index: dict[tuple[FlatDemandTier, ...], int] = {}
    month_indexes: list[int] = []

    for bands in monthly_band_sets:
        key = tuple(bands)
        if key not in band_index:
            band_index[key] = len(band_sets)
            band_sets.append(key)
        month_indexes.append(band_index[key])

    return {
        "flatdemandunit": _SUPPORTED_DEMAND_UNIT,
        "flatdemandstructure": [list(bands) for bands in band_sets],
        "flatdemandmonths": tuple(month_indexes),
    }


def get_month_flat_demand_bands(scenario: Scenario, library: Library, month: int) -> list[FlatDemandTier]:
    """Return the combined flat-demand bands applicable for one month."""

    sample_dt = datetime(scenario.year, month, 15, 12)
    tariff = library.tariffs.get_tariff_at_date(scenario.master_tariff_id, sample_dt.date())
    rates = [
        rate
        for rate in tariff.get("rates", [])
        if ru.rate_is_applied_to_scenario(rate, scenario, library) and ru.rate_is_applied_to_datetime(rate, sample_dt)
    ]

    demand_bands: list[list[FlatDemandTier]] = []
    for rate in rates:
        bands = get_rate_flat_demand_bands(rate, scenario, library)
        if bands is None:
            continue
        demand_bands.append(bands)

    if not demand_bands:
        return []
    if len(demand_bands) == 1:
        return demand_bands[0]
    return _sum_flat_demand_band_sets(demand_bands)


def get_rate_flat_demand_bands(
    rate: TariffRateExtended, scenario: Scenario, library: Library
) -> list[FlatDemandTier] | None:
    """Convert one Arcadia demand-based rate into URDB flat-demand tiers."""

    if rate.get("charge_type") != "DEMAND_BASED":
        return None
    if rate.get("time_of_use") is not None:
        raise RateConversionError(rate, "TOU demand charges are not supported by flat demand conversion")
    if rate.get("quantity_key") is not None:
        raise RateConversionError(rate, "Rates with quantity_key are not supported for flat demand conversion")
    if rate["charge_period"] != "MONTHLY":
        raise RateConversionError(rate, "Flat demand conversion only supports MONTHLY charge period")
    if rate["transaction_type"] != "BUY":
        raise RateConversionError(rate, "Flat demand conversion only supports BUY transaction type")

    bands = [
        _validate_flat_demand_band(rate, band)
        for band in ru.rate_filter_bands(rate, scenario, library)
    ]
    return sorted(bands, key=lambda band: band.get("max", inf))


def _validate_flat_demand_band(rate: TariffRateExtended, band: TariffRateBand) -> FlatDemandTier:
    """Validate and convert one Arcadia demand band into a URDB flat-demand tier."""

    if band["rate_unit"] != "COST_PER_UNIT":
        raise RateConversionError(rate, "Flat demand bands must have rate unit = COST_PER_UNIT")
    if band["is_credit"]:
        raise RateConversionError(rate, "Credit demand rates are not supported")
    if band["has_consumption_limit"]:
        raise RateConversionError(rate, "Flat demand bands cannot have consumption limits")
    if band.get("consumption_upper_limit") is not None:
        raise RateConversionError(rate, "Flat demand bands cannot have consumption_upper_limit")
    if band["has_property_limit"]:
        raise RateConversionError(rate, "Flat demand bands cannot have property limits")
    if band.get("property_upper_limit") is not None:
        raise RateConversionError(rate, "Flat demand bands cannot have property_upper_limit")
    if band.get("prev_upper_limit") not in {None, 0.0}:
        raise RateConversionError(rate, "Flat demand bands cannot have prev_upper_limit")

    tier: FlatDemandTier = {"rate": band["rate_amount"]}
    if band["has_demand_limit"] or band.get("demand_upper_limit") is not None:
        if (limit := band.get("demand_upper_limit")) is None:
            raise RateConversionError(rate, "Flat demand band has_demand_limit=true but no demand_upper_limit")
        tier["max"] = limit
    return tier


def _sum_flat_demand_band_sets(inputs: list[list[FlatDemandTier]]) -> list[FlatDemandTier]:
    """Sum multiple flat-demand tier sets over their combined demand breakpoints."""

    limits = sorted({tier["max"] for tiers in inputs for tier in tiers if "max" in tier})
    if any("max" not in tiers[-1] for tiers in inputs):
        limits.append(inf)

    def value_at(tiers: list[FlatDemandTier], limit: float) -> float:
        for tier in tiers:
            if limit <= tier.get("max", inf):
                return tier["rate"]
        return 0.0

    result: list[FlatDemandTier] = []
    for limit in limits:
        rate = sum(value_at(tiers, limit) for tiers in inputs)
        tier: FlatDemandTier = {"rate": rate}
        if limit != inf:
            tier["max"] = limit
        result.append(tier)

    collapsed = [result[0]]
    for tier in result[1:]:
        if tier["rate"] != collapsed[-1]["rate"]:
            collapsed.append(tier)
    return collapsed
