from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime
from decimal import Decimal
import logging

from homeassistant.components.sensor import ATTR_LAST_RESET, SensorDeviceClass, SensorEntity, SensorStateClass
from homeassistant.const import (
    ATTR_UNIT_OF_MEASUREMENT,
    CONF_NAME,
    CONF_UNIQUE_ID,
    UnitOfEnergy,
)
from homeassistant.core import Event, EventStateChangedData, HomeAssistant, State, callback
from homeassistant.helpers.event import async_track_state_change_event
from homeassistant.helpers.restore_state import RestoreEntity
from homeassistant.helpers.typing import ConfigType
from homeassistant.util import dt as dt_util

from custom_components.powercalc.common import SourceEntity
from custom_components.powercalc.const import (
    CONF_COST,
    CONF_COST_SENSOR_PRECISION,
    CONF_ENERGY_PRICE,
    CONF_ENERGY_PRICE_MULTIPLIER,
    CONF_ENERGY_PRICE_SENSOR,
    CONF_ENERGY_PRICE_SURCHARGE,
    CONF_ENERGY_SENSOR_ID,
    DEFAULT_COST_SENSOR_PRECISION,
    DOMAIN,
    DOMAIN_CONFIG,
)
from custom_components.powercalc.unit import convert_to_decimal, parse_decimal

from .abstract import (
    BaseEntity,
    generate_cost_sensor_entity_id,
    generate_cost_sensor_name,
)
from .energy import EnergySensor, resolve_existing_energy_sensor
from .utility_meter import VirtualUtilityMeter

COST_ICON = "mdi:cash"
ATTR_LAST_ENERGY = "last_energy"

_LOGGER = logging.getLogger(__name__)


def _parse_scaled(state: State | None, unit_to_factor: Callable[[str | None], Decimal]) -> Decimal | None:
    """Parse a numeric state into a Decimal, scaled by a factor derived from its unit.

    Both energy amounts (converted to kWh) and prices (converted to per kWh) are parsed this
    way; ``unit_to_factor`` maps the state's unit of measurement to the multiplier to apply.
    """
    value = parse_decimal(state)
    if value is None:
        return None
    assert state is not None  # a parsed value implies the state is present and numeric
    return value * unit_to_factor(state.attributes.get(ATTR_UNIT_OF_MEASUREMENT))


def _to_kwh_factor(unit: str | None) -> Decimal:
    """Return the multiplier to convert an energy value expressed in `unit` to kWh.

    Falls back to 1 (assume kWh) when the unit is missing or not a recognizable energy unit.
    """
    if not unit or unit == UnitOfEnergy.KILO_WATT_HOUR:
        return Decimal(1)
    if (factor := convert_to_decimal(1, unit, UnitOfEnergy.KILO_WATT_HOUR)) is not None:
        return factor
    _LOGGER.warning("Cannot convert energy unit '%s' to kWh, assuming kWh", unit)
    return Decimal(1)


def _price_per_kwh_factor(unit: str | None) -> Decimal:
    """Return the multiplier that converts a price value to a price per kWh.

    Price sensors express their unit as ``<currency>/<energy>`` (for example ``EUR/MWh``).
    The energy denominator determines how the raw value maps to a per-kWh price, e.g.
    ``EUR/MWh`` -> value * 0.001 and ``EUR/Wh`` -> value * 1000. Falls back to 1 (assume the
    value is already per kWh) when there is no denominator or it is not a recognizable unit.
    """
    if not unit or "/" not in unit:
        return Decimal(1)
    denominator = unit.rsplit("/", 1)[-1].strip()
    if not denominator or denominator == UnitOfEnergy.KILO_WATT_HOUR:
        return Decimal(1)
    if (factor := convert_to_decimal(1, UnitOfEnergy.KILO_WATT_HOUR, denominator)) is not None:
        return factor
    _LOGGER.warning("Cannot convert energy price unit '%s' to a per-kWh price, assuming per kWh", unit)
    return Decimal(1)


def _currency_from_price_unit(unit: str | None) -> str | None:
    """Derive the monetary unit from an energy price sensor's unit of measurement.

    Price sensors typically express their unit as ``<currency>/kWh`` (for example
    ``EUR/kWh`` or ``€/kWh``). The part before the slash is the currency the cost is
    denominated in. Returns None when no currency can be determined.
    """
    if not unit:
        return None
    currency = unit.split("/", 1)[0].strip()
    return currency or None


@dataclass(frozen=True, slots=True)
class PriceConfig:
    """The resolved energy price configuration for a single cost sensor."""

    fixed_price: Decimal | None = None
    price_entity_id: str | None = None
    surcharge: Decimal = Decimal(0)
    multiplier: Decimal = Decimal(1)

    @property
    def has_price_source(self) -> bool:
        """Return whether a price source (fixed price or price sensor) is configured."""
        return self.fixed_price is not None or bool(self.price_entity_id)

    def effective_price(self, price: Decimal | None) -> Decimal | None:
        """Apply the surcharge and multiplier to a raw price per kWh."""
        if price is None:
            return None
        return (price + self.surcharge) * self.multiplier


def resolve_price_config(hass: HomeAssistant, sensor_config: ConfigType) -> PriceConfig:
    """Resolve the energy price for a sensor, preferring its own config over the global one.

    The price source is resolved as a whole: a sensor which defines a fixed price or a price
    sensor of its own fully replaces the globally configured price source, rather than mixing
    the two. The surcharge and multiplier are independent and fall back to the global value
    individually.
    """
    global_config: ConfigType = hass.data[DOMAIN].get(DOMAIN_CONFIG, {})

    if CONF_ENERGY_PRICE in sensor_config or CONF_ENERGY_PRICE_SENSOR in sensor_config:
        fixed_price = sensor_config.get(CONF_ENERGY_PRICE)
        price_entity_id = sensor_config.get(CONF_ENERGY_PRICE_SENSOR)
    else:
        fixed_price = global_config.get(CONF_ENERGY_PRICE)
        price_entity_id = global_config.get(CONF_ENERGY_PRICE_SENSOR)

    def resolve(key: str, default: float) -> Decimal:
        value = sensor_config.get(key, global_config.get(key))
        return Decimal(str(default if value is None else value))

    return PriceConfig(
        fixed_price=Decimal(str(fixed_price)) if fixed_price is not None else None,
        price_entity_id=price_entity_id,
        surcharge=resolve(CONF_ENERGY_PRICE_SURCHARGE, 0),
        multiplier=resolve(CONF_ENERGY_PRICE_MULTIPLIER, 1),
    )


def create_cost_sensor(
    hass: HomeAssistant,
    sensor_config: ConfigType,
    energy_sensor: EnergySensor | VirtualUtilityMeter,
    source_entity: SourceEntity | None = None,
    name: str | None = None,
    reset_on_source_reset: bool = False,
    unique_id: str | None = None,
) -> CostSensor | None:
    """Create a cost sensor tracking the cost of the given energy sensor.

    The energy sensor can be a regular energy sensor or a utility meter. The energy price is
    taken from the sensor configuration when it defines one, otherwise from the global
    configuration, either as a fixed price or a price sensor. When no price is configured,
    no cost sensor is created. A ``unique_id`` can be provided to override the id derived
    from the energy sensor (used for standalone cost sensor config entries).
    """
    price_config = resolve_price_config(hass, sensor_config)
    if not price_config.has_price_source:
        _LOGGER.warning(
            "Cost sensor creation is enabled but no energy price is configured. "
            "Define `energy_price` or `energy_price_sensor` on the sensor itself "
            "or in the global powercalc configuration",
        )
        return None

    name_base = name if name is not None else sensor_config.get(CONF_NAME)
    cost_name = generate_cost_sensor_name(sensor_config, name_base, source_entity)
    if unique_id is None and energy_sensor.unique_id is not None:
        unique_id = f"{energy_sensor.unique_id}_cost"
    entity_id = generate_cost_sensor_entity_id(
        hass,
        sensor_config,
        source_entity,
        name=name_base,
        unique_id=unique_id,
    )

    _LOGGER.debug(
        "Creating cost sensor (entity_id=%s, source_entity=%s, %s)",
        entity_id,
        energy_sensor.entity_id,
        price_config,
    )

    return CostSensor(
        hass=hass,
        source_energy_entity=energy_sensor.entity_id,
        entity_id=entity_id,
        unique_id=unique_id,
        name=cost_name,
        sensor_config=sensor_config,
        price_config=price_config,
        reset_on_source_reset=reset_on_source_reset,
    )


def create_cost_sensor_for_energy_entity(hass: HomeAssistant, sensor_config: ConfigType) -> CostSensor | None:
    """Create a standalone cost sensor tracking an existing (non-powercalc) energy sensor.

    The tracked energy sensor is read from the ``cost`` block (YAML) or the flat
    ``energy_sensor_id`` key (GUI config entry).
    """
    cost_config = sensor_config.get(CONF_COST, {})
    energy_sensor_id = cost_config.get(CONF_ENERGY_SENSOR_ID) or sensor_config[CONF_ENERGY_SENSOR_ID]
    energy_sensor = resolve_existing_energy_sensor(hass, energy_sensor_id)
    # In YAML the name is optional, fall back to the name of the tracked energy sensor.
    name = sensor_config.get(CONF_NAME) or energy_sensor.name
    return create_cost_sensor(
        hass,
        sensor_config,
        energy_sensor,
        name=name,
        unique_id=sensor_config.get(CONF_UNIQUE_ID),
    )


class CostSensor(BaseEntity, RestoreEntity, SensorEntity):
    """Cost sensor, accumulating the cost of the energy consumed at price-at-consumption."""

    _attr_device_class = SensorDeviceClass.MONETARY
    _attr_state_class = SensorStateClass.TOTAL
    _attr_should_poll = False
    _attr_icon = COST_ICON
    _unrecorded_attributes = frozenset({ATTR_LAST_ENERGY})

    def __init__(
        self,
        hass: HomeAssistant,
        source_energy_entity: str,
        entity_id: str,
        sensor_config: ConfigType,
        name: str | None = None,
        unique_id: str | None = None,
        price_config: PriceConfig | None = None,
        reset_on_source_reset: bool = False,
    ) -> None:
        self._source_energy_entity = source_energy_entity
        self._sensor_config = sensor_config
        self._reset_on_source_reset = reset_on_source_reset
        self._attr_name = name
        self._attr_unique_id = unique_id
        self._attr_native_unit_of_measurement = hass.config.currency
        self._price_config = price_config or PriceConfig()
        self._price_entity_id = self._price_config.price_entity_id
        self._rounding_digits = int(sensor_config.get(CONF_COST_SENSOR_PRECISION, DEFAULT_COST_SENSOR_PRECISION))
        self._attr_suggested_display_precision = self._rounding_digits
        self._state: Decimal = Decimal(0)
        self._last_energy: Decimal | None = None
        self._current_price: Decimal | None = self._price_config.effective_price(self._price_config.fixed_price)
        # Only a resetting (per utility meter cycle) sensor uses last_reset; a lifetime cost
        # sensor accumulates monotonically and leaves it None.
        self._attr_last_reset: datetime | None = None
        self.entity_id = entity_id

    async def async_added_to_hass(self) -> None:
        """Restore state and start tracking the source energy and price sensors."""
        await super().async_added_to_hass()

        if (state := await self.async_get_last_state()) is not None:
            self._state = parse_decimal(state) or Decimal(0)
            self._last_energy = parse_decimal(state.attributes.get(ATTR_LAST_ENERGY))

        # A resetting (per utility meter cycle) cost sensor exposes last_reset so long-term
        # statistics treat each cycle reset as a new cycle rather than negative consumption.
        # Restore it across restarts, falling back to now for a freshly created sensor.
        if self._reset_on_source_reset:
            restored = dt_util.parse_datetime(state.attributes.get(ATTR_LAST_RESET, "")) if state else None
            self._attr_last_reset = restored or dt_util.utcnow()

        _LOGGER.debug("%s: Restoring cost sensor state: %s", self.entity_id, self._state)

        # Seed the current price and, for a price sensor, track its changes so consumption
        # is always settled at the price that was in effect when it was consumed.
        if self._price_entity_id is not None:
            price_state = self.hass.states.get(self._price_entity_id)
            # Prefer the currency of the price sensor (e.g. `€/kWh` -> `€`) over the HA currency.
            price_unit = price_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) if price_state else None
            if (currency := _currency_from_price_unit(price_unit)) is not None:
                self._attr_native_unit_of_measurement = currency
            self._current_price = self._price_config.effective_price(
                _parse_scaled(price_state, _price_per_kwh_factor),
            )
            self.async_on_remove(
                async_track_state_change_event(
                    self.hass,
                    [self._price_entity_id],
                    self._handle_price_state_change,
                ),
            )

        self.async_on_remove(
            async_track_state_change_event(
                self.hass,
                [self._source_energy_entity],
                self._handle_energy_state_change,
            ),
        )

    @callback
    def _handle_energy_state_change(self, event: Event[EventStateChangedData]) -> None:
        """Accumulate cost based on the delta of the energy sensor and the current price."""
        new_energy = _parse_scaled(event.data["new_state"], _to_kwh_factor)
        if new_energy is None:
            return

        # First reading only establishes a baseline to measure deltas from.
        if self._last_energy is None:
            self._last_energy = new_energy
            return

        self._accumulate(new_energy, self._current_price)

    @callback
    def _handle_price_state_change(self, event: Event[EventStateChangedData]) -> None:
        """Settle the energy consumed so far at the previous price, then adopt the new price."""
        # Recalculate the outstanding energy delta with the previously known price before switching.
        if self._last_energy is not None and self._current_price is not None:
            current_energy = _parse_scaled(self.hass.states.get(self._source_energy_entity), _to_kwh_factor)
            if current_energy is not None:
                self._accumulate(current_energy, self._current_price)

        self._current_price = self._price_config.effective_price(
            _parse_scaled(event.data["new_state"], _price_per_kwh_factor),
        )

    def _accumulate(self, new_energy: Decimal, price: Decimal | None) -> None:
        """Add the cost of the consumed energy at the given price and advance the baseline."""
        # Leave _last_energy untouched when no price is known, so the consumption is
        # priced once a price becomes available again.
        if price is None or self._last_energy is None:
            return

        delta = new_energy - self._last_energy
        if delta < 0:
            if self._reset_on_source_reset:
                # The source (utility meter) reset for a new cycle, start the cost cycle over.
                self._state = Decimal(0)
                self._last_energy = new_energy
                self._attr_last_reset = dt_util.utcnow()
                self.async_write_ha_state()
                return
            # The energy sensor got reset (e.g. restart or calibrate), treat the new value as the delta.
            delta = new_energy
        if delta == 0:
            return

        self._state += delta * price
        self._last_energy = new_energy
        self.async_write_ha_state()

    @callback
    def async_reset(self) -> None:
        """Reset the cost sensor to zero from the current source energy reading."""
        _LOGGER.debug("%s: Reset cost sensor", self.entity_id)
        self._state = Decimal(0)
        self._attr_last_reset = dt_util.utcnow()
        self._set_current_energy_baseline()
        self.async_write_ha_state()

    async def async_calibrate(self, value: str) -> None:
        """Set the cost sensor to the given value from the current source energy reading."""
        _LOGGER.debug("%s: Calibrate cost sensor to: %s", self.entity_id, value)
        self._state = Decimal(value)
        self._set_current_energy_baseline()
        self.async_write_ha_state()

    def _set_current_energy_baseline(self) -> None:
        current_energy = _parse_scaled(self.hass.states.get(self._source_energy_entity), _to_kwh_factor)
        if current_energy is not None:
            self._last_energy = current_energy

    @property
    def native_value(self) -> Decimal:
        """Return the accumulated cost."""
        return Decimal(round(self._state, self._rounding_digits))

    @property
    def extra_state_attributes(self) -> dict[str, str] | None:
        """Return the state attributes of the cost sensor."""
        if self._last_energy is None:
            return None
        return {ATTR_LAST_ENERGY: str(self._last_energy)}
