zlock_slots/custom_components/zlock_slots/config_flow.py

from __future__ import annotations

import voluptuous as vol

from homeassistant.config_entries import ConfigFlow, ConfigFlowResult, OptionsFlow
from homeassistant.core import callback
from homeassistant.helpers import selector

from .const import DOMAIN, CONF_LOCK_ENTITY_IDS, CONF_GROUPS, CONF_GROUP_NAME


class ZLockSlotsConfigFlow(ConfigFlow, domain=DOMAIN):
  VERSION = 1

  async def async_step_user(
    self, user_input: dict | None = None
  ) -> ConfigFlowResult:
    errors: dict[str, str] = {}

    if user_input is not None:
      lock_ids: list[str] = user_input.get(CONF_LOCK_ENTITY_IDS, [])
      if not lock_ids:
        errors[CONF_LOCK_ENTITY_IDS] = "no_locks_selected"
      else:
        title = self._make_title(lock_ids)
        return self.async_create_entry(title=title, data={CONF_LOCK_ENTITY_IDS: lock_ids})

    schema = vol.Schema({
      vol.Required(CONF_LOCK_ENTITY_IDS): selector.EntitySelector(
        selector.EntitySelectorConfig(
          domain="lock",
          integration="zwave_js",
          multiple=True,
        )
      ),
    })

    return self.async_show_form(
      step_id="user",
      data_schema=schema,
      errors=errors,
    )

  @staticmethod
  @callback
  def async_get_options_flow(config_entry) -> "ZLockSlotsOptionsFlow":
    return ZLockSlotsOptionsFlow(config_entry)

  def _make_title(self, lock_ids: list[str]) -> str:
    if len(lock_ids) == 1:
      state = self.hass.states.get(lock_ids[0])
      if state:
        return state.name
    return f"Z-Lock Slots ({len(lock_ids)} locks)"


class ZLockSlotsOptionsFlow(OptionsFlow):
  def __init__(self, config_entry) -> None:
    self._entry = config_entry
    self._groups: list[dict] = list(config_entry.options.get(CONF_GROUPS, []))

  async def async_step_init(self, user_input=None) -> ConfigFlowResult:
    return self.async_show_menu(
      step_id="init",
      menu_options=["manage_locks", "add_group", "remove_group", "finish"],
    )

  async def async_step_manage_locks(self, user_input=None) -> ConfigFlowResult:
    errors: dict[str, str] = {}
    current_lock_ids: list[str] = self._entry.data.get(CONF_LOCK_ENTITY_IDS, [])

    if user_input is not None:
      new_lock_ids: list[str] = user_input.get(CONF_LOCK_ENTITY_IDS, [])
      if not new_lock_ids:
        errors[CONF_LOCK_ENTITY_IDS] = "no_locks_selected"
      else:
        removed = set(current_lock_ids) - set(new_lock_ids)
        if removed:
          self._groups = [
            g for g in self._groups
            if not all(lid in removed for lid in g.get(CONF_LOCK_ENTITY_IDS, []))
          ]
          for g in self._groups:
            g[CONF_LOCK_ENTITY_IDS] = [
              lid for lid in g.get(CONF_LOCK_ENTITY_IDS, []) if lid not in removed
            ]
          self._groups = [g for g in self._groups if g.get(CONF_LOCK_ENTITY_IDS)]
        self.hass.config_entries.async_update_entry(
          self._entry, data={CONF_LOCK_ENTITY_IDS: new_lock_ids}
        )
        return await self.async_step_init()

    schema = vol.Schema({
      vol.Required(CONF_LOCK_ENTITY_IDS, default=current_lock_ids): selector.EntitySelector(
        selector.EntitySelectorConfig(
          domain="lock",
          integration="zwave_js",
          multiple=True,
        )
      ),
    })
    return self.async_show_form(step_id="manage_locks", data_schema=schema, errors=errors)

  async def async_step_add_group(self, user_input=None) -> ConfigFlowResult:
    errors: dict[str, str] = {}
    lock_ids: list[str] = self._entry.data[CONF_LOCK_ENTITY_IDS]

    if user_input is not None:
      name = user_input.get(CONF_GROUP_NAME, "").strip()
      selected: list[str] = user_input.get(CONF_LOCK_ENTITY_IDS, [])
      if not name:
        errors[CONF_GROUP_NAME] = "group_name_required"
      elif any(g[CONF_GROUP_NAME] == name for g in self._groups):
        errors[CONF_GROUP_NAME] = "group_name_exists"
      elif not selected:
        errors[CONF_LOCK_ENTITY_IDS] = "no_locks_selected"
      else:
        self._groups.append({CONF_GROUP_NAME: name, CONF_LOCK_ENTITY_IDS: selected})
        return await self.async_step_init()

    lock_options = [
      selector.SelectOptionDict(
        value=lid,
        label=(self.hass.states.get(lid).name if self.hass.states.get(lid) else lid),
      )
      for lid in lock_ids
    ]
    schema = vol.Schema({
      vol.Required(CONF_GROUP_NAME): selector.TextSelector(),
      vol.Required(CONF_LOCK_ENTITY_IDS): selector.SelectSelector(
        selector.SelectSelectorConfig(options=lock_options, multiple=True)
      ),
    })
    return self.async_show_form(step_id="add_group", data_schema=schema, errors=errors)

  async def async_step_remove_group(self, user_input=None) -> ConfigFlowResult:
    if not self._groups:
      return await self.async_step_init()

    if user_input is not None:
      names_to_remove = set(user_input.get("groups", []))
      self._groups = [g for g in self._groups if g[CONF_GROUP_NAME] not in names_to_remove]
      return await self.async_step_init()

    schema = vol.Schema({
      vol.Required("groups"): selector.SelectSelector(
        selector.SelectSelectorConfig(
          options=[
            selector.SelectOptionDict(value=g[CONF_GROUP_NAME], label=g[CONF_GROUP_NAME])
            for g in self._groups
          ],
          multiple=True,
        )
      ),
    })
    return self.async_show_form(step_id="remove_group", data_schema=schema)

  async def async_step_finish(self, user_input=None) -> ConfigFlowResult:
    return self.async_create_entry(data={CONF_GROUPS: self._groups})