From 7f68c8cad20762be70ea232f101c5a04bc2351ac Mon Sep 17 00:00:00 2001
From: John Lancaster <32917998+jsl12@users.noreply.github.com>
Date: Thu, 25 Jul 2024 22:57:47 -0500
Subject: [PATCH] WIP

---
 src/room_control/model.py        |  6 ++---
 src/room_control/room_control.py | 42 ++++++++++++--------------------
 2 files changed, 18 insertions(+), 30 deletions(-)

diff --git a/src/room_control/model.py b/src/room_control/model.py
index c70c348..3da9e0a 100644
--- a/src/room_control/model.py
+++ b/src/room_control/model.py
@@ -54,9 +54,9 @@ class ControllerStateConfig(BaseModel):
     @model_validator(mode='before')
     def check_args(cls, values):
         time, elevation = values.get('time'), values.get('elevation')
-        if time is not None and elevation is not None:
-            raise PydanticCustomError('bad_time_spec', 'Only one of time or elevation can be set.')
-        elif elevation is not None and 'direction' not in values:
+        # if time is not None and elevation is not None:
+        #     raise PydanticCustomError('bad_time_spec', 'Only one of time or elevation can be set.')
+        if elevation is not None and 'direction' not in values:
             raise PydanticCustomError('no_sun_dir', 'Needs sun direction with elevation')
         return values
 
diff --git a/src/room_control/room_control.py b/src/room_control/room_control.py
index ce27a39..376baf2 100755
--- a/src/room_control/room_control.py
+++ b/src/room_control/room_control.py
@@ -54,39 +54,27 @@ class RoomController(Hass):
         self.register_service(f'{self.name}/activate', self.service_activate)
         self.register_service(f'{self.name}/deactivate', self.service_deactivate)
 
-        self.app_entities = self.gather_app_entities()
-        # self.log(f'entities: {self.app_entities}')
-
         self.refresh_state_times()
         self.run_daily(callback=self.refresh_state_times, start='00:00:00')
 
+        self.app_entities = set(self.gather_app_entities())
+        self.log(f'entities: {self.app_entities}', level='DEBUG')
+
         self.log(f'Initialized [bold green]{type(self).__name__}[/]')
 
     def terminate(self):
         self.log('[bold red]Terminating[/]', level='DEBUG')
 
-    def gather_app_entities(self) -> List[str]:
-        """Returns a list of all the entities involved in any of the states"""
+    def gather_app_entities(self):
+        """Yields all the entities involved in any of the states"""
 
-        def generator():
-            for settings in deepcopy(self.args['states']):
-                if scene := settings.get('scene'):
-                    if isinstance(scene, str):
-                        assert scene.startswith(
-                            'scene.'
-                        ), f"Scene definition must start with 'scene.' for app {self.name}"
-                        entity: Entity = self.get_entity(scene)
-                        entity_state = entity.get_state('all')
-                        attributes = entity_state['attributes']
-                        for entity in attributes['entity_id']:
-                            yield entity
-                    else:
-                        for key in scene.keys():
-                            yield key
-                else:
-                    yield self.args['entity']
-
-        return set(list(generator()))
+        for state in self._room_config.states:
+            if isinstance(state.scene, str):
+                assert state.scene.startswith('scene.'), "Scene definition must start with 'scene.'"
+                entities = self.get_state(state.scene, namespace='default', attribute='entity_id')
+                yield from entities
+            else:
+                yield from state.scene.keys()
 
     def refresh_state_times(self, *args, **kwargs):
         """Resets the `self.states` attribute to a newly parsed version of the states.
@@ -102,18 +90,18 @@ class RoomController(Hass):
 
         for state in self._room_config.states:
             if state.time is None and state.elevation is not None:
-                transition_time = self.time_at_elevation(
+                state.time = self.time_at_elevation(
                     elevation=state.elevation, direction=state.direction
                 ).time()
             elif isinstance(state.time, str):
-                transition_time = self.parse_time(state.time)
+                state.time = self.parse_time(state.time)
 
             assert isinstance(state.time, datetime.time), f'Invalid time: {state.time}'
 
             try:
                 self.run_at(
                     callback=lambda cb_args: self.set_controller_scene(cb_args['state']),
-                    start=transition_time.strftime('%H:%M:%S'),
+                    start=state.time.strftime('%H:%M:%S'),
                     state=state,
                 )
             except Exception as e: