"""
Place this code in custom_components/heat_pump_monitor/sensor.py
"""
from homeassistant.components.sensor import (
    SensorEntity,
    SensorDeviceClass,
    SensorStateClass,
)
from homeassistant.const import UnitOfTemperature

from . import DOMAIN
from .__init__ import HeatPumpEntity

async def async_setup_entry(hass, entry, async_add_entities):
    """Set up sensors."""
    coordinator = hass.data[DOMAIN][entry.entry_id]

    entities = [
        HeatPumpModeSensor(coordinator, "water_mode"),
        HeatPumpTemperatureSensor(coordinator, "current_temperature", "Current Temperature"),
    ]
    async_add_entities(entities)

class HeatPumpModeSensor(HeatPumpEntity, SensorEntity):
    """Sensor for heat pump water mode."""

    def __init__(self, coordinator, key):
        """Initialize the sensor."""
        super().__init__(coordinator, key)
        self._attr_name = "Heat Pump Water Mode"
        self._attr_unique_id = f"{DOMAIN}_water_mode"
        self._attr_device_class = SensorDeviceClass.ENUM

    @property
    def native_value(self):
        """Return the state of the sensor."""
        return self.coordinator.data.get(self.key)

class HeatPumpTemperatureSensor(HeatPumpEntity, SensorEntity):
    """Sensor for heat pump temperature values."""

    def __init__(self, coordinator, key, name):
        """Initialize the sensor."""
        super().__init__(coordinator, key)
        self._attr_name = f"Heat Pump {name}"
        self._attr_unique_id = f"{DOMAIN}_{key}"
        self._attr_device_class = SensorDeviceClass.TEMPERATURE
        self._attr_native_unit_of_measurement = UnitOfTemperature.CELSIUS
        self._attr_state_class = SensorStateClass.MEASUREMENT

    @property
    def native_value(self):
        """Return the state of the sensor."""
        return self.coordinator.data.get(self.key)