from enum import StrEnum
import logging

from homeassistant.components.sensor import SensorDeviceClass
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_UNIQUE_ID, EVENT_HOMEASSISTANT_STARTED
from homeassistant.core import Event, HomeAssistant, callback
from homeassistant.helpers.entity import Entity
import homeassistant.helpers.entity_registry as er
from homeassistant.helpers.entity_registry import EVENT_ENTITY_REGISTRY_UPDATED, EventEntityRegistryUpdatedData
from homeassistant.helpers.typing import ConfigType

from custom_components.powercalc.const import (
    CONF_CREATE_ENERGY_SENSOR,
    CONF_DISABLE_EXTENDED_ATTRIBUTES,
    CONF_ENERGY_SENSOR_UNIT_PREFIX,
    CONF_EXCLUDE_ENTITIES,
    CONF_GROUP_TRACKED_AUTO,
    CONF_GROUP_TRACKED_POWER_ENTITIES,
    CONF_MAIN_POWER_SENSOR,
    CONF_UTILITY_METER_NET_CONSUMPTION,
    GroupType,
    UnitPrefix,
)
from custom_components.powercalc.group_include.filter import LambdaFilter
from custom_components.powercalc.group_include.include import find_entities
from custom_components.powercalc.sensors.abstract import (
    generate_energy_sensor_entity_id,
    generate_energy_sensor_name,
    generate_power_sensor_entity_id,
    generate_power_sensor_name,
)
from custom_components.powercalc.sensors.energy import VirtualEnergySensor
from custom_components.powercalc.sensors.energy_related import create_energy_related_sensors
from custom_components.powercalc.sensors.group.custom import GroupedPowerSensor, GroupedSensor
from custom_components.powercalc.sensors.group.subtract import SubtractGroupSensor
from custom_components.powercalc.sensors.power import PowerSensor

_LOGGER = logging.getLogger(__name__)


class SensorType(StrEnum):
    TRACKED = "tracked"
    UNTRACKED = "untracked"

    @property
    def label(self) -> str:
        """Capitalized variant, used in the friendly names of the created sensors."""
        return self.value.capitalize()


async def find_auto_tracked_power_entities(hass: HomeAssistant, exclude_entities: set[str] | None = None) -> set[str]:
    """Find tracked power entities."""
    entity_filter = None
    if exclude_entities:
        entity_filter = LambdaFilter(lambda entity: entity.entity_id not in exclude_entities)
    result = await find_entities(hass, entity_filter)
    return {
        entity.entity_id
        for entity in result.resolved
        if isinstance(entity, PowerSensor) and not isinstance(entity, GroupedSensor)
    }


class TrackedPowerSensorFactory:
    def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry, config: ConfigType) -> None:
        self.hass = hass
        self.tracked_entities: set[str] = set()
        self.config_entry = config_entry
        self.config = config

    async def create_tracked_untracked_group_sensors(self) -> list[Entity]:
        """Create tracked/untracked group sensors."""

        unique_id = str(self.config.get(CONF_UNIQUE_ID))
        main_power_sensor = (
            str(self.config.get(CONF_MAIN_POWER_SENSOR)) if self.config.get(CONF_MAIN_POWER_SENSOR) else None
        )
        self.config[CONF_DISABLE_EXTENDED_ATTRIBUTES] = True  # prevent adding all entities in the state attributes

        self.tracked_entities = await self.get_tracked_power_entities()
        if main_power_sensor and main_power_sensor in self.tracked_entities:
            self.tracked_entities.remove(main_power_sensor)

        should_create_energy_sensor = bool(self.config.get(CONF_CREATE_ENERGY_SENSOR, False))

        entities: list[Entity] = []
        tracked_sensor = await self.create_tracked_power_sensor(SensorType.TRACKED, unique_id, self.tracked_entities)
        await self._add_power_sensor_with_energy(
            entities,
            SensorType.TRACKED,
            tracked_sensor,
            should_create_energy_sensor,
        )

        if main_power_sensor:
            untracked_sensor = await self.create_untracked_power_sensor(
                SensorType.UNTRACKED,
                unique_id,
                main_power_sensor,
                tracked_sensor.entity_id,
            )
            await self._add_power_sensor_with_energy(
                entities,
                SensorType.UNTRACKED,
                untracked_sensor,
                should_create_energy_sensor,
            )

        return entities

    async def _add_power_sensor_with_energy(
        self,
        entities: list[Entity],
        sensor_type: SensorType,
        power_sensor: GroupedPowerSensor,
        should_create_energy_sensor: bool,
    ) -> None:
        """Append the power sensor and, when enabled, its energy and energy-related sensors."""
        entities.append(power_sensor)
        if not should_create_energy_sensor:
            return

        energy_sensor = await self.create_energy_sensor(sensor_type, power_sensor)
        entities.append(energy_sensor)
        entities.extend(
            create_energy_related_sensors(
                self.hass,
                self.config,
                energy_sensor,
                utility_meter_config={CONF_UTILITY_METER_NET_CONSUMPTION: True, **self.config},
                cost_name=sensor_type.label,
            ),
        )

    async def get_tracked_power_entities(self) -> set[str]:
        """
        Get all power entities which are part of the tracked sensor group
        """
        if not bool(self.config.get(CONF_GROUP_TRACKED_AUTO, False)):
            tracked_entities: list[str] | None = self.config.get(CONF_GROUP_TRACKED_POWER_ENTITIES)
            if not isinstance(tracked_entities, list):
                return set()
            return set(tracked_entities)

        # For auto mode, we also want to listen for any changes in the entity registry
        # Dynamically add/remove power sensors from the tracked group
        @callback
        def _start_entity_registry_listener(_: Event) -> None:
            self.hass.bus.async_listen(EVENT_ENTITY_REGISTRY_UPDATED, self._handle_entity_registry_updated)

        self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STARTED, _start_entity_registry_listener)

        exclude_entities = self.config.get(CONF_EXCLUDE_ENTITIES)
        return await find_auto_tracked_power_entities(self.hass, set(exclude_entities) if exclude_entities else None)

    async def _handle_entity_registry_updated(
        self,
        event: Event[EventEntityRegistryUpdatedData],
    ) -> None:
        """Listen to all entity registry updates and reload the config entry if a power sensor is added/removed."""
        entity_id = event.data["entity_id"]
        action = event.data["action"]

        old_entity_id = event.data.get("old_entity_id")
        if action == "update" and old_entity_id is not None:
            if old_entity_id in self.tracked_entities:
                return await self.reload()
            return None  # pragma: no cover

        if action == "remove" and entity_id in self.tracked_entities:
            return await self.reload()

        if action == "create":
            registry = er.async_get(self.hass)
            entity_entry = registry.async_get(entity_id)
            if entity_entry and entity_entry.original_device_class == SensorDeviceClass.POWER:
                return await self.reload()
        return None

    async def reload(self) -> None:
        """Reload the config entry."""
        await self.hass.config_entries.async_reload(self.config_entry.entry_id)

    async def create_tracked_power_sensor(
        self,
        sensor_type: SensorType,
        unique_id: str,
        tracked_entities: set[str],
    ) -> GroupedPowerSensor:
        _LOGGER.debug("Creating tracked grouped power sensor, entities: %s", tracked_entities)
        unique_id = f"{unique_id}_{sensor_type}_power"
        entity_id = generate_power_sensor_entity_id(self.hass, self.config, name=sensor_type, unique_id=unique_id)
        name = generate_power_sensor_name(self.config, name=sensor_type.label)
        return GroupedPowerSensor(
            self.hass,
            sensor_config=self.config,
            group_type=GroupType.TRACKED_UNTRACKED,
            entities=tracked_entities,
            entity_id=entity_id,
            name=name,
            unique_id=unique_id,
        )

    async def create_untracked_power_sensor(
        self,
        sensor_type: SensorType,
        unique_id: str,
        main_power_entity_id: str,
        tracked_entity_id: str,
    ) -> GroupedPowerSensor:
        _LOGGER.debug("Creating untracked grouped power sensor")
        unique_id = f"{unique_id}_{sensor_type}_power"
        entity_id = generate_power_sensor_entity_id(self.hass, self.config, name=sensor_type, unique_id=unique_id)
        name = generate_power_sensor_name(self.config, name=sensor_type.label)
        return SubtractGroupSensor(
            self.hass,
            entity_id=entity_id,
            name=name,
            sensor_config=self.config,
            base_entity_id=main_power_entity_id,
            subtract_entities=[tracked_entity_id],
            unique_id=unique_id,
        )

    async def create_energy_sensor(
        self,
        sensor_type: SensorType,
        power_sensor: GroupedPowerSensor,
    ) -> VirtualEnergySensor:
        """Create an energy sensor for a power sensor."""
        _LOGGER.debug("Creating %s grouped energy sensor", sensor_type)
        unique_id = f"{power_sensor.unique_id}_{sensor_type}_energy"
        name = generate_energy_sensor_name(self.config, sensor_type.label)
        entity_id = generate_energy_sensor_entity_id(self.hass, self.config, name=sensor_type, unique_id=unique_id)
        return VirtualEnergySensor(
            hass=self.hass,
            source_entity=power_sensor.entity_id,
            entity_id=entity_id,
            name=name,
            unique_id=unique_id,
            sensor_config=self.config,
            unit_prefix=self.config.get(CONF_ENERGY_SENSOR_UNIT_PREFIX, UnitPrefix.KILO),
        )
