"""Number entities for the Xiaoxiang Smart BMS integration.

Provides user-configurable overrides for BMS values that may be stale or
incorrect in the BMS firmware itself (e.g. nominal capacity after a cell
upgrade). Overrides are stored in HA and shadow the BMS-reported value
inside coordinator.data so all other sensors stay consistent.
"""
from __future__ import annotations

from homeassistant.components.number import (
    NumberDeviceClass,
    NumberEntity,
    NumberEntityDescription,
    NumberMode,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.restore_state import RestoreEntity
from homeassistant.helpers.update_coordinator import CoordinatorEntity

from .const import DOMAIN
from .coordinator import BmsCoordinator

NUMBER_ENTITIES: tuple[NumberEntityDescription, ...] = (
    NumberEntityDescription(
        key="nominal_capacity_override",
        name="Nominal Capacity (Override)",
        device_class=NumberDeviceClass.ENERGY_STORAGE,
        native_unit_of_measurement="Ah",
        native_min_value=1,
        native_max_value=2000,
        native_step=0.1,
        mode=NumberMode.BOX,
        entity_category=EntityCategory.CONFIG,
        icon="mdi:battery-edit",
    ),
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    coordinator: BmsCoordinator = hass.data[DOMAIN][entry.entry_id]
    async_add_entities(
        BmsNumberEntity(coordinator, description)
        for description in NUMBER_ENTITIES
    )


class BmsNumberEntity(CoordinatorEntity[BmsCoordinator], NumberEntity, RestoreEntity):
    """A number entity whose value overrides a field in coordinator.data.

    On change: immediately patches coordinator.data so the corresponding
    sensor reflects the new value without waiting for the next poll.
    On HA restart: restores the last set value from state history.
    """

    _attr_has_entity_name = True

    def __init__(
        self,
        coordinator: BmsCoordinator,
        description: NumberEntityDescription,
    ) -> None:
        super().__init__(coordinator)
        self.entity_description = description
        self._attr_unique_id = f"{coordinator.address}_{description.key}"
        self._attr_device_info = coordinator.device_info
        self._override_value: float | None = None

    # ------------------------------------------------------------------
    # State restoration across HA restarts
    # ------------------------------------------------------------------

    async def async_added_to_hass(self) -> None:
        await super().async_added_to_hass()
        if (last_state := await self.async_get_last_state()) is not None:
            try:
                self._override_value = float(last_state.state)
                self._apply_override()
            except (ValueError, TypeError):
                pass

    # ------------------------------------------------------------------
    # Entity interface
    # ------------------------------------------------------------------

    @property
    def native_value(self) -> float | None:
        if self._override_value is not None:
            return self._override_value
        # Fall back to BMS-reported value while no override is set
        data_key = self.entity_description.key.replace("_override", "")
        return self.coordinator.data.get(data_key)

    async def async_set_native_value(self, value: float) -> None:
        self._override_value = round(value, 1)
        self._apply_override()
        self.async_write_ha_state()

    def _apply_override(self) -> None:
        """Patch coordinator.data so dependent sensors update immediately."""
        if self.coordinator.data and self._override_value is not None:
            data_key = self.entity_description.key.replace("_override", "")
            self.coordinator.data[data_key] = self._override_value
            # Recalculate energy_stored since nominal_capacity changed
            if data_key == "nominal_capacity" and "voltage" in self.coordinator.data:
                pass  # energy_stored uses residual_capacity, not nominal — no recalc needed