66from  __future__ import  annotations 
77
88import  logging 
9+ import  re 
910from  datetime  import  datetime 
11+ from  dataclasses  import  dataclass , field 
1012
1113import  homeassistant .helpers .config_validation  as  cv 
1214import  voluptuous  as  vol 
13- import  re 
1415
1516from  awesomeversion .awesomeversion  import  AwesomeVersion 
1617from  homeassistant .config_entries  import  ConfigEntry 
1718from  homeassistant .core  import  HomeAssistant , callback 
1819from  homeassistant .const  import  __version__  as  HA_VERSION   # noqa: N812 
1920from  homeassistant .helpers .typing  import  ConfigType 
2021from  homeassistant .helpers  import  device_registry  as  dr 
22+ from  homeassistant .helpers  import  entity_registry  as  er 
2123from  homeassistant .util  import  dt  as  dt_util 
2224
2325from  .config_flow  import  CONFIG_VERSION 
2426
27+ from  .device  import  BatteryNotesDevice 
2528from  .discovery  import  DiscoveryManager 
2629from  .library_updater  import  (
2730    LibraryUpdater ,
2831)
29- from  .coordinator  import  BatteryNotesCoordinator 
3032from  .store  import  (
3133    async_get_registry ,
3234)
3739    PLATFORMS ,
3840    CONF_ENABLE_AUTODISCOVERY ,
3941    CONF_USER_LIBRARY ,
42+     DATA ,
4043    DATA_LIBRARY_UPDATER ,
4144    CONF_SHOW_ALL_DEVICES ,
4245    CONF_ENABLE_REPLACED ,
46+     CONF_DEFAULT_BATTERY_LOW_THRESHOLD ,
47+     CONF_BATTERY_INCREASE_THRESHOLD ,
48+     CONF_HIDE_BATTERY ,
49+     DEFAULT_BATTERY_LOW_THRESHOLD ,
50+     DEFAULT_BATTERY_INCREASE_THRESHOLD ,
4351    SERVICE_BATTERY_REPLACED ,
4452    SERVICE_BATTERY_REPLACED_SCHEMA ,
45-     DATA_COORDINATOR ,
53+     SERVICE_DATA_DATE_TIME_REPLACED ,
54+     DATA_STORE ,
4655    ATTR_REMOVE ,
4756    ATTR_DEVICE_ID ,
48-     ATTR_DATE_TIME_REPLACED ,
4957    CONF_BATTERY_TYPE ,
5058    CONF_BATTERY_QUANTITY ,
5159)
6371                    vol .Optional (CONF_USER_LIBRARY , default = "" ): cv .string ,
6472                    vol .Optional (CONF_SHOW_ALL_DEVICES , default = False ): cv .boolean ,
6573                    vol .Optional (CONF_ENABLE_REPLACED , default = True ): cv .boolean ,
74+                     vol .Optional (CONF_HIDE_BATTERY , default = False ): cv .boolean ,
75+                     vol .Optional (
76+                         CONF_DEFAULT_BATTERY_LOW_THRESHOLD ,
77+                         default = DEFAULT_BATTERY_LOW_THRESHOLD ,
78+                     ): cv .positive_int ,
79+                     vol .Optional (
80+                         CONF_BATTERY_INCREASE_THRESHOLD ,
81+                         default = DEFAULT_BATTERY_INCREASE_THRESHOLD ,
82+                     ): cv .positive_int ,
6683                },
6784            ),
6885        ),
6986    },
7087    extra = vol .ALLOW_EXTRA ,
7188)
7289
90+ 
91+ @dataclass  
92+ class  BatteryNotesData :
93+     """Class for sharing data within the BatteryNotes integration.""" 
94+ 
95+     devices : dict [str , BatteryNotesDevice ] =  field (default_factory = dict )
96+     platforms : dict  =  field (default_factory = dict )
97+ 
98+ 
7399async  def  async_setup (hass : HomeAssistant , config : ConfigType ) ->  bool :
74100    """Integration setup.""" 
75101
@@ -86,60 +112,96 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
86112        CONF_ENABLE_AUTODISCOVERY : True ,
87113        CONF_SHOW_ALL_DEVICES : False ,
88114        CONF_ENABLE_REPLACED : True ,
115+         CONF_HIDE_BATTERY : False ,
116+         CONF_DEFAULT_BATTERY_LOW_THRESHOLD : DEFAULT_BATTERY_LOW_THRESHOLD ,
117+         CONF_BATTERY_INCREASE_THRESHOLD : DEFAULT_BATTERY_INCREASE_THRESHOLD ,
89118    }
90119
91120    hass .data [DOMAIN ] =  {
92121        DOMAIN_CONFIG : domain_config ,
93122    }
94123
95124    store  =  await  async_get_registry (hass )
125+     hass .data [DOMAIN ][DATA_STORE ] =  store 
96126
97-     coordinator  =  BatteryNotesCoordinator (hass , store )
98-     hass .data [DOMAIN ][DATA_COORDINATOR ] =  coordinator 
127+     hass .data [DOMAIN ][DATA ] =  BatteryNotesData ()
99128
100129    library_updater  =  LibraryUpdater (hass )
101130
102131    await  library_updater .get_library_updates (dt_util .utcnow ())
103132
104133    hass .data [DOMAIN ][DATA_LIBRARY_UPDATER ] =  library_updater 
105134
106-     await  coordinator .async_refresh ()
107- 
108135    if  domain_config .get (CONF_ENABLE_AUTODISCOVERY ):
109136        discovery_manager  =  DiscoveryManager (hass , config )
110137        await  discovery_manager .start_discovery ()
111138    else :
112139        _LOGGER .debug ("Auto discovery disabled" )
113140
141+     # Register custom services 
142+     register_services (hass )
143+ 
114144    return  True 
115145
116- async  def  async_setup_entry (hass : HomeAssistant , entry : ConfigEntry ) ->  bool :
117-     """Set up a config entry.""" 
118146
119-     await  hass .config_entries .async_forward_entry_setups (entry , PLATFORMS )
147+ async  def  async_setup_entry (hass : HomeAssistant , config_entry : ConfigEntry ) ->  bool :
148+     """Set up a config entry.""" 
120149
121-     entry . async_on_unload ( entry . add_update_listener ( async_update_options ) )
150+     device :  BatteryNotesDevice   =   BatteryNotesDevice ( hass ,  config_entry )
122151
123-     # Register custom services 
124-     register_services ( hass ) 
152+     if   not   await   device . async_setup (): 
153+          return   False 
125154
126155    return  True 
127156
128157
158+ async  def  async_unload_entry (hass : HomeAssistant , config_entry : ConfigEntry ) ->  bool :
159+     """Unload a config entry.""" 
160+     data : BatteryNotesData  =  hass .data [DOMAIN ][DATA ]
161+ 
162+     device  =  data .devices .pop (config_entry .entry_id )
163+     await  device .async_unload ()
164+ 
165+     return  await  hass .config_entries .async_unload_platforms (config_entry , PLATFORMS )
166+ 
167+ 
129168async  def  async_remove_entry (hass : HomeAssistant , config_entry : ConfigEntry ) ->  None :
130169    """Device removed, tidy up store.""" 
131170
132171    if  "device_id"  not  in config_entry .data :
133172        return 
134173
135-     device_id  =  config_entry .data ["device_id" ]
174+     device : BatteryNotesDevice  =  hass .data [DOMAIN ][DATA ].devices [config_entry .entry_id ]
175+     if  not  device :
176+         return 
136177
137-     coordinator : BatteryNotesCoordinator  =  hass .data [DOMAIN ][DATA_COORDINATOR ]
138178    data  =  {ATTR_REMOVE : True }
139179
140-     coordinator .async_update_device_config (device_id = device_id , data = data )
180+     device .coordinator .async_update_device_config (
181+         device_id = device .coordinator .device_id , data = data 
182+     )
183+ 
184+     _LOGGER .debug ("Removed Device %s" , device .coordinator .device_id )
185+ 
186+     # Unhide the battery 
187+     entity_registry  =  er .async_get (hass )
188+     if  not  device .wrapped_battery :
189+         return 
190+ 
191+     if  not  (
192+         wrapped_battery_entity_entry  :=  entity_registry .async_get (
193+             device .wrapped_battery .entity_id 
194+         )
195+     ):
196+         return 
141197
142-     _LOGGER .debug ("Removed Device %s" , device_id )
198+     if  wrapped_battery_entity_entry .hidden_by  ==  er .RegistryEntryHider .INTEGRATION :
199+         entity_registry .async_update_entity (
200+             device .wrapped_battery .entity_id , hidden_by = None 
201+         )
202+         _LOGGER .debug (
203+             "Unhidden Original Battery for device%s" , device .coordinator .device_id 
204+         )
143205
144206
145207async  def  async_migrate_entry (hass , config_entry : ConfigEntry ):
@@ -180,43 +242,47 @@ async def async_migrate_entry(hass, config_entry: ConfigEntry):
180242
181243    return  True 
182244
245+ 
183246@callback  
184247async  def  async_update_options (hass : HomeAssistant , entry : ConfigEntry ) ->  None :
185248    """Update options.""" 
186249    await  hass .config_entries .async_reload (entry .entry_id )
187250
188251
189- async  def  async_unload_entry (hass : HomeAssistant , entry : ConfigEntry ) ->  bool :
190-     """Unload a config entry.""" 
191-     return  await  hass .config_entries .async_unload_platforms (entry , PLATFORMS )
192- 
193- 
194252@callback  
195253def  register_services (hass ):
196254    """Register services used by battery notes component.""" 
197255
198256    async  def  handle_battery_replaced (call ):
199257        """Handle the service call.""" 
200258        device_id  =  call .data .get (ATTR_DEVICE_ID , "" )
201-         datetime_replaced_entry  =  call .data .get (ATTR_DATE_TIME_REPLACED )
259+         datetime_replaced_entry  =  call .data .get (SERVICE_DATA_DATE_TIME_REPLACED )
202260
203261        if  datetime_replaced_entry :
204-             datetime_replaced  =  dt_util .as_utc (datetime_replaced_entry ).replace (tzinfo = None )
262+             datetime_replaced  =  dt_util .as_utc (datetime_replaced_entry ).replace (
263+                 tzinfo = None 
264+             )
205265        else :
206266            datetime_replaced  =  datetime .utcnow ()
207267
208268        device_registry  =  dr .async_get (hass )
209269
210270        device_entry  =  device_registry .async_get (device_id )
211271        if  not  device_entry :
272+             _LOGGER .error (
273+                 "Device %s not found" ,
274+                 device_id ,
275+             )
212276            return 
213277
214278        for  entry_id  in  device_entry .config_entries :
215279            if  (
216280                entry  :=  hass .config_entries .async_get_entry (entry_id )
217281            ) and  entry .domain  ==  DOMAIN :
282+                 coordinator  =  (
283+                     hass .data [DOMAIN ][DATA ].devices [entry .entry_id ].coordinator 
284+                 )
218285
219-                 coordinator : BatteryNotesCoordinator  =  hass .data [DOMAIN ][DATA_COORDINATOR ]
220286                device_entry  =  {"battery_last_replaced" : datetime_replaced }
221287
222288                coordinator .async_update_device_config (
@@ -226,9 +292,19 @@ async def handle_battery_replaced(call):
226292                await  coordinator .async_request_refresh ()
227293
228294                _LOGGER .debug (
229-                     "Device %s battery replaced on %s" , device_id , str (datetime_replaced )
295+                     "Device %s battery replaced on %s" ,
296+                     device_id ,
297+                     str (datetime_replaced ),
230298                )
231299
300+                 # Found and dealt with, exit 
301+                 return 
302+ 
303+         _LOGGER .error (
304+             "Device %s not configured in Battery Notes" ,
305+             device_id ,
306+         )
307+ 
232308    hass .services .async_register (
233309        DOMAIN ,
234310        SERVICE_BATTERY_REPLACED ,
0 commit comments