wled-screen-controller-mixed/custom_components/wled_screen_controller/__init__.py

"""The LED Screen Controller integration."""
from __future__ import annotations

import logging
from datetime import timedelta

import voluptuous as vol

from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.aiohttp_client import async_get_clientsession

from .const import (
    DOMAIN,
    CONF_SERVER_NAME,
    CONF_SERVER_URL,
    CONF_API_KEY,
    DEFAULT_SCAN_INTERVAL,
    TARGET_TYPE_KEY_COLORS,
    DATA_COORDINATOR,
    DATA_WS_MANAGER,
    DATA_EVENT_LISTENER,
)
from .coordinator import WLEDScreenControllerCoordinator
from .event_listener import EventStreamListener
from .ws_manager import KeyColorsWebSocketManager

_LOGGER = logging.getLogger(__name__)

PLATFORMS: list[Platform] = [
    Platform.BUTTON,
    Platform.LIGHT,
    Platform.SWITCH,
    Platform.SENSOR,
    Platform.NUMBER,
    Platform.SELECT,
]


async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Set up LED Screen Controller from a config entry."""
    server_name = entry.data.get(CONF_SERVER_NAME, "LED Screen Controller")
    server_url = entry.data[CONF_SERVER_URL]
    api_key = entry.data[CONF_API_KEY]

    session = async_get_clientsession(hass)
    coordinator = WLEDScreenControllerCoordinator(
        hass,
        session,
        server_url,
        api_key,
        update_interval=timedelta(seconds=DEFAULT_SCAN_INTERVAL),
    )

    await coordinator.async_config_entry_first_refresh()

    ws_manager = KeyColorsWebSocketManager(hass, server_url, api_key)

    event_listener = EventStreamListener(hass, server_url, api_key, coordinator)
    await event_listener.start()

    # Create device entries for each target and remove stale ones
    device_registry = dr.async_get(hass)
    current_identifiers: set[tuple[str, str]] = set()
    if coordinator.data and "targets" in coordinator.data:
        for target_id, target_data in coordinator.data["targets"].items():
            info = target_data["info"]
            target_type = info.get("target_type", "led")
            model = (
                "Key Colors Target"
                if target_type == TARGET_TYPE_KEY_COLORS
                else "LED Target"
            )
            device_registry.async_get_or_create(
                config_entry_id=entry.entry_id,
                identifiers={(DOMAIN, target_id)},
                name=info.get("name", target_id),
                manufacturer=server_name,
                model=model,
                configuration_url=server_url,
            )
            current_identifiers.add((DOMAIN, target_id))

    # Create a single "Scenes" device for scene preset buttons
    scenes_identifier = (DOMAIN, f"{entry.entry_id}_scenes")
    scene_presets = coordinator.data.get("scene_presets", []) if coordinator.data else []
    if scene_presets:
        device_registry.async_get_or_create(
            config_entry_id=entry.entry_id,
            identifiers={scenes_identifier},
            name="Scenes",
            manufacturer=server_name,
            model="Scene Presets",
            configuration_url=server_url,
        )
        current_identifiers.add(scenes_identifier)

    # Remove devices for targets that no longer exist
    for device_entry in dr.async_entries_for_config_entry(
        device_registry, entry.entry_id
    ):
        if not device_entry.identifiers & current_identifiers:
            _LOGGER.info("Removing stale device: %s", device_entry.name)
            device_registry.async_remove_device(device_entry.id)

    # Store data
    hass.data.setdefault(DOMAIN, {})
    hass.data[DOMAIN][entry.entry_id] = {
        DATA_COORDINATOR: coordinator,
        DATA_WS_MANAGER: ws_manager,
        DATA_EVENT_LISTENER: event_listener,
    }

    # Track target and scene IDs to detect changes
    known_target_ids = set(
        coordinator.data.get("targets", {}).keys() if coordinator.data else []
    )
    known_scene_ids = set(
        p["id"] for p in (coordinator.data.get("scene_presets", []) if coordinator.data else [])
    )

    def _on_coordinator_update() -> None:
        """Manage WS connections and detect target list changes."""
        nonlocal known_target_ids, known_scene_ids

        if not coordinator.data:
            return

        targets = coordinator.data.get("targets", {})

        # Start/stop WS connections for KC targets based on processing state
        for target_id, target_data in targets.items():
            info = target_data.get("info", {})
            state = target_data.get("state") or {}
            if info.get("target_type") == TARGET_TYPE_KEY_COLORS:
                if state.get("processing"):
                    if target_id not in ws_manager._connections:
                        hass.async_create_task(ws_manager.start_listening(target_id))
                else:
                    if target_id in ws_manager._connections:
                        hass.async_create_task(ws_manager.stop_listening(target_id))

        # Reload if target or scene list changed
        current_ids = set(targets.keys())
        current_scene_ids = set(
            p["id"] for p in coordinator.data.get("scene_presets", [])
        )
        if current_ids != known_target_ids or current_scene_ids != known_scene_ids:
            known_target_ids = current_ids
            known_scene_ids = current_scene_ids
            _LOGGER.info("Target or scene list changed, reloading integration")
            hass.async_create_task(
                hass.config_entries.async_reload(entry.entry_id)
            )

    coordinator.async_add_listener(_on_coordinator_update)

    # Register set_leds service (once across all entries)
    async def handle_set_leds(call) -> None:
        """Handle the set_leds service call."""
        source_id = call.data["source_id"]
        segments = call.data["segments"]
        # Route to the coordinator that owns this source
        for entry_data in hass.data[DOMAIN].values():
            coord = entry_data.get(DATA_COORDINATOR)
            if not coord or not coord.data:
                continue
            source_ids = {
                s["id"] for s in coord.data.get("css_sources", [])
            }
            if source_id in source_ids:
                await coord.push_segments(source_id, segments)
                return
        _LOGGER.error("No server found with source_id %s", source_id)

    if not hass.services.has_service(DOMAIN, "set_leds"):
        hass.services.async_register(
            DOMAIN,
            "set_leds",
            handle_set_leds,
            schema=vol.Schema({
                vol.Required("source_id"): str,
                vol.Required("segments"): list,
            }),
        )

    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)

    return True


async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Unload a config entry."""
    entry_data = hass.data[DOMAIN][entry.entry_id]
    await entry_data[DATA_WS_MANAGER].shutdown()
    await entry_data[DATA_EVENT_LISTENER].shutdown()

    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

    if unload_ok:
        hass.data[DOMAIN].pop(entry.entry_id)
        # Unregister service if no entries remain
        if not hass.data[DOMAIN]:
            hass.services.async_remove(DOMAIN, "set_leds")

    return unload_ok