sensor.py

"""Sensor platform for Microsoft 365 OneDrive."""
from __future__ import annotations

from datetime import timedelta
import logging
from typing import Any

from homeassistant.components.sensor import SensorEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed

from .api import OneDriveAPI
from .const import DOMAIN

_LOGGER = logging.getLogger(__name__)

SCAN_INTERVAL = timedelta(minutes=15)


async def async_setup_entry(
    hass: HomeAssistant,
    config_entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up the Microsoft 365 OneDrive sensor from a config entry."""
    _LOGGER.debug("Setting up OneDrive sensor for entry: %s", config_entry.entry_id)

    # Check if the API instance exists
    if DOMAIN not in hass.data or config_entry.entry_id not in hass.data[DOMAIN]:
        _LOGGER.error("OneDrive API not found in hass.data for entry: %s", config_entry.entry_id)
        return

    # Retrieve the shared API instance from hass.data
    api: OneDriveAPI = hass.data[DOMAIN][config_entry.entry_id]

    # Check if sensor already exists to prevent duplicates
    sensor_unique_id = f"{config_entry.entry_id}_storage"
    coordinator_key = f"{DOMAIN}_{config_entry.entry_id}_coordinator"

    # Check if coordinator already exists
    if coordinator_key in hass.data.get(DOMAIN, {}):
        _LOGGER.debug("Coordinator already exists for entry %s, skipping duplicate setup", config_entry.entry_id)
        return

    # Check existing entities in the entity registry
    try:
        from homeassistant.helpers import entity_registry as er
        entity_registry = er.async_get(hass)
        existing_entity = entity_registry.async_get_entity_id("sensor", DOMAIN, sensor_unique_id)
        if existing_entity:
            _LOGGER.debug("Sensor already exists with unique_id %s (entity_id: %s), skipping duplicate setup",
                          sensor_unique_id, existing_entity)
            return
    except Exception as err:
        _LOGGER.debug("Could not check entity registry: %s", err)

    async def async_update_data():
        """Fetch data from API endpoint for the sensor."""
        try:
            return await api.get_drive_info()
        except Exception as err:
            raise UpdateFailed(f"Error communicating with API: {err}") from err

    # Set up a coordinator to manage data updates for the sensor
    coordinator = DataUpdateCoordinator(
        hass,
        _LOGGER,
        name=f"onedrive_sensor_{config_entry.entry_id}",
        update_method=async_update_data,
        update_interval=SCAN_INTERVAL,
    )

    # Store coordinator to prevent duplicates
    hass.data.setdefault(DOMAIN, {})[coordinator_key] = coordinator

    # Fetch initial data so we have it when the sensor is first added
    await coordinator.async_config_entry_first_refresh()

    async_add_entities([OneDriveStorageSensor(coordinator, config_entry)])
    _LOGGER.debug("Successfully set up OneDrive sensor for entry: %s", config_entry.entry_id)


class OneDriveStorageSensor(SensorEntity):
    """Representation of a OneDrive storage sensor."""

    _attr_has_entity_name = True
    _attr_icon = "mdi:microsoft-onedrive"

    def __init__(
        self, coordinator: DataUpdateCoordinator, entry: ConfigEntry
    ) -> None:
        """Initialize the sensor."""
        self.coordinator = coordinator
        # Link the sensor to the config entry and device
        self._attr_unique_id = f"{entry.entry_id}_storage"
        self._attr_device_info = {
            "identifiers": {(DOMAIN, entry.entry_id)},
            "name": "Microsoft 365 OneDrive",
            "manufacturer": "Microsoft",
            "entry_type": "service",
        }
        self._attr_name = "Storage Used"
        self._attr_native_unit_of_measurement = "GB"

    @property
    def available(self) -> bool:
        """Return True if entity is available."""
        return self.coordinator.last_update_success and self.coordinator.data is not None

    @property
    def state(self) -> float | None:
        """Return the state of the sensor (used storage)."""
        if self.coordinator.data and "quota" in self.coordinator.data:
            used = self.coordinator.data["quota"].get("used", 0) / (1024**3)
            return round(used, 2)
        return None

    @property
    def extra_state_attributes(self) -> dict[str, Any] | None:
        """Return the state attributes of the sensor."""
        if self.coordinator.data and "quota" in self.coordinator.data:
            quota = self.coordinator.data["quota"]
            total_gb = quota.get("total", 0) / (1024**3)
            used_gb = quota.get("used", 0) / (1024**3)
            
            # The API sometimes returns 'remaining', sometimes not. Calculate if needed.
            if "remaining" in quota:
                remaining_gb = quota.get("remaining", 0) / (1024**3)
            else:
                remaining_gb = total_gb - used_gb

            return {
                "total_gb": round(total_gb, 2),
                "used_gb": round(used_gb, 2),
                "free_gb": round(remaining_gb, 2),
                "used_percent": round((used_gb / total_gb) * 100, 2) if total_gb > 0 else 0,
                "owner": self.coordinator.data.get("owner", {}).get("user", {}).get("displayName"),
            }
        return None

    async def async_added_to_hass(self) -> None:
        """When entity is added to hass."""
        self.async_on_remove(
            self.coordinator.async_add_listener(self.async_write_ha_state)
        )

    async def async_update(self) -> None:
        """Update the sensor. This is handled by the coordinator."""
        await self.coordinator.async_request_refresh()