From 5fd3bcd260d909aa3d051404e2cfee7c7c2708b6 Mon Sep 17 00:00:00 2001
From: lunDreame <87955512+lunDreame@users.noreply.github.com>
Date: Sat, 17 Aug 2024 03:01:31 +0900
Subject: [PATCH] v1.1.1

---
 custom_components/bestin/api.py               | 70 +++++++++++--------
 custom_components/bestin/climate.py           |  2 +-
 custom_components/bestin/config_flow.py       |  2 +
 custom_components/bestin/const.py             | 19 +++--
 custom_components/bestin/controller.py        |  8 ++-
 custom_components/bestin/device.py            | 58 ++++++++-------
 custom_components/bestin/fan.py               |  2 +-
 custom_components/bestin/hub.py               | 23 ++----
 custom_components/bestin/light.py             |  2 +-
 custom_components/bestin/manifest.json        |  2 +-
 custom_components/bestin/sensor.py            |  2 +-
 custom_components/bestin/switch.py            |  4 +-
 custom_components/bestin/translations/en.json |  1 +
 custom_components/bestin/translations/ko.json |  1 +
 14 files changed, 106 insertions(+), 90 deletions(-)

diff --git a/custom_components/bestin/api.py b/custom_components/bestin/api.py
index 858b3b1..87eddad 100644
--- a/custom_components/bestin/api.py
+++ b/custom_components/bestin/api.py
@@ -43,18 +43,25 @@
 class BestinAPIv2:
     """Bestin HDC Smarthome API v2 Class."""
 
-    def __init__(self) -> None:        
+    def __init__(self, hass, entry) -> None:        
         """API initialization."""
+        self.elevator_count = entry.data.get("elevator_count", 1)
         self.elevator_arrived = False
 
         self.features_list: list = []
-        self.elev_moveinfo: dict = {}
+        self.elevator_data: dict = {}
+
+    def setup_elevators(self):
+        for count in range(1, self.elevator_count + 1):
+            self.setup_device("elevator", 1, str(count), False)
+            self.setup_device("elevator", 1, f"floor_{str(count)}", "대기 층")
+            self.setup_device("elevator", 1, f"direction_{str(count)}", "대기")
 
     async def _v2_device_status(self, args=None):
         if args is not None:
             LOGGER.debug(f"Task execution started with argument: {args}")
             self.last_update_time = args
-
+        
         if self.features_list:
             await self.process_features(self.features_list)
         else:
@@ -107,6 +114,8 @@ async def elevator_call_request(self) -> None:
 
                 if response.status == 200 and result_status == "ok":
                     LOGGER.info(f"Just a central server elevator request successful")
+                    for count in range(1, self.elevator_count + 1):
+                        self.setup_device("elevator", 1, str(count), False)
                     self.hass.create_task(self.fetch_elevator_status())
                 else:
                     LOGGER.error(f"Only central server elevator request failed: {response_data}")
@@ -139,28 +148,26 @@ async def handle_message_info(self, message) -> None:
         data = json.loads(message)
         
         if "move_info" in data:
-            if not os.path.exists('data.json'):
-                async with aiofiles.open('data.json', 'w') as file:
-                    await file.write(json.dumps(data["move_info"], indent=4))
-            
-            self.elev_moveinfo.update(data["move_info"])
-            #LOGGER.debug(f"Elevator move updated: {self.elev_moveinfo}")
+            serial = data["move_info"]["Serial"]
+            move_info = data["move_info"]
+
+            self.elevator_data = {serial: move_info}
+            self.elevator_data = dict(sorted(self.elevator_data.items()))
+            LOGGER.debug(f"Elevator data: {self.elevator_data}")
+            if len(self.elevator_data) >= 2:
+                for idx, (serial, info) in enumerate(self.elevator_data.items(), start=1):
+                    floor = info["move_info"]["Floor"]
+                    move_dir = info["move_info"]["MoveDir"]
             
-            serial = str(self.elev_moveinfo.get("Serial", "1"))
-            floor = f"{str(self.elev_moveinfo.get('Floor', '대기'))} 층"
-            movedir = self.elev_moveinfo.get("MoveDir", "대기")
-
-            self.setup_device("elevator", 1, f"floor_{serial}", floor)
-            self.setup_device("elevator", 1, f"direction_{serial}", movedir)
+                    self.setup_device("elevator", 1, f"floor_{str(idx)}", floor)
+                    self.setup_device("elevator", 1, f"direction_{str(idx)}", move_dir)
+            else:
+                self.setup_device("elevator", 1, f"floor_1", move_info["Floor"])
+                self.setup_device("elevator", 1, f"direction_1", move_info["MoveDir"])
         else:
-            self.elev_moveinfo = data
-            
-            serial = str(self.elev_moveinfo.get("Serial", "1"))
-            floor = f"{str(self.elev_moveinfo.get('Floor', '도착'))} 층"
-
-            self.setup_device("elevator", 1, serial, False)
-            self.setup_device("elevator", 1, f"floor_{serial}", floor)
-            self.setup_device("elevator", 1, f"direction_{serial}", "도착")
+            for count in range(1, self.elevator_count + 1):
+                self.setup_device("elevator", 1, f"floor_{str(count)}", "도착 층")
+                self.setup_device("elevator", 1, f"direction_{str(count)}", "도착")
             self.elevator_arrived = True
 
     async def request_feature_command(self, device_type: str, room_id: int, unit: str, value: str) -> None:
@@ -273,16 +280,18 @@ def __init__(
         hub_id: str,
         version: str,
         version_identifier: str,
+        elevator_registration: bool,
         async_add_device: Callable
     ) -> None:
         """API initialization."""
-        super().__init__()
+        super().__init__(hass, entry)
         self.hass = hass
         self.entry = entry
         self.entities = entities
         self.hub_id = hub_id
         self.version = version
         self.version_identifier = version_identifier
+        self.elevator_registration = elevator_registration
         self.async_add_device = async_add_device
 
         ssl_context = ssl.create_default_context()
@@ -302,7 +311,7 @@ def __init__(
     def get_short_hash(self, id: str) -> str:
         """Generate a short hash for a given id."""
         hash_object = hashlib.sha256(id.encode()).digest()
-        return base64.urlsafe_b64encode(hash_object)[:8].decode("utf-8")
+        return base64.urlsafe_b64encode(hash_object)[:8].decode("utf-8").upper()
 
     async def start(self) -> None:
         """Start main loop with asyncio."""
@@ -310,6 +319,9 @@ async def start(self) -> None:
         self.tasks.append(asyncio.create_task(self.schedule_session_refresh()))
         await asyncio.sleep(1)
 
+        if self.elevator_registration:
+            LOGGER.debug("Setting up elevator")
+            self.setup_elevators()
         v_key = getattr(self, f"_v{self.version[7:8]}_device_status")
         scan_interval = self.entry.options.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL)
         interval = timedelta(minutes=scan_interval)
@@ -445,15 +457,17 @@ def initialize_device(self, device_id: str, unit_id: Optional[str], state: Any)
 
         if full_unique_id not in self.devices:
             device_info = DeviceInfo(
-                id=full_unique_id,
-                type=device_type,
+                unique_id=full_unique_id,
+                device_type=device_type,
                 name=unique_id,
                 room=device_room,
                 state=state,
+                sub_type=unit_id or "",
+                colon_id=device_type if ":" in device_type else ""
             )
             device = Device(
                 info=device_info,
-                platform=platform,
+                domain=platform,
                 on_command=self.on_command,
                 callbacks=set()
             )
diff --git a/custom_components/bestin/climate.py b/custom_components/bestin/climate.py
index cd259fa..e117ffb 100644
--- a/custom_components/bestin/climate.py
+++ b/custom_components/bestin/climate.py
@@ -35,7 +35,7 @@ def async_add_climate(devices=None):
         entities = [
             BestinClimate(device, hub) 
             for device in devices 
-            if device.info.id not in hub.entities[DOMAIN]
+            if device.info.unique_id not in hub.entities[DOMAIN]
         ]
 
         if entities:
diff --git a/custom_components/bestin/config_flow.py b/custom_components/bestin/config_flow.py
index c9fd47b..21dff65 100644
--- a/custom_components/bestin/config_flow.py
+++ b/custom_components/bestin/config_flow.py
@@ -31,6 +31,7 @@
     DOMAIN,
     LOGGER,
     DEFAULT_PORT,
+    DEFAULT_ELEVATOR_COUNT,
     DEFAULT_SCAN_INTERVAL,
     DEFAULT_MAX_TRANSMISSION,
     DEFAULT_PACKET_VIEWER,
@@ -237,6 +238,7 @@ async def async_step_center_v2(
 
         data_schema = vol.Schema({
             vol.Optional(CONF_IP_ADDRESS): cv.string,
+            vol.Required("elevator_count", default=DEFAULT_ELEVATOR_COUNT): ConfigFlow.int_between(1, 3),
             vol.Required(CONF_UUID): cv.string,
         })
 
diff --git a/custom_components/bestin/const.py b/custom_components/bestin/const.py
index 1c994cb..dd8ac51 100644
--- a/custom_components/bestin/const.py
+++ b/custom_components/bestin/const.py
@@ -1,7 +1,7 @@
 import logging
 
-from typing import Any, Callable
-from dataclasses import dataclass
+from typing import Callable, Any, Optional, Set
+from dataclasses import dataclass, field
 
 from homeassistant.components.sensor import SensorDeviceClass
 from homeassistant.const import (
@@ -14,7 +14,7 @@
 
 DOMAIN = "bestin"
 NAME = "BESTIN"
-VERSION = "1.1.0"
+VERSION = "1.1.1"
 
 PLATFORMS: list[Platform] = [
     Platform.CLIMATE,
@@ -27,6 +27,7 @@
 LOGGER: logging.Logger = logging.getLogger(__package__)
 
 DEFAULT_PORT: int = 8899
+DEFAULT_ELEVATOR_COUNT: int = 1
 DEFAULT_SCAN_INTERVAL: int = 15
 DEFAULT_MAX_TRANSMISSION: int = 10
 
@@ -175,22 +176,26 @@
 @dataclass
 class DeviceInfo:
     """Represents the basic information of a device."""
-    id: str
-    type: str
+    unique_id: str
+    device_type: str
     name: str
     room: str
     state: Any
+    colon_id: Optional[str] = None
+    sub_type: Optional[str] = None
 
 @dataclass
 class Device:
     """Represents a device with callbacks and update functionalities."""
     info: DeviceInfo
-    platform: Platform
+    domain: str
     on_command: Callable
-    callbacks: set[Callable]
+    callbacks: Set[Callable] = field(default_factory=set)
     
     def add_callback(self, callback: Callable):
+        """Add a callback to the set of callbacks."""
         self.callbacks.add(callback)
 
     def remove_callback(self, callback: Callable):
+        """Remove a callback from the set of callbacks, if it exists."""
         self.callbacks.discard(callback)
diff --git a/custom_components/bestin/controller.py b/custom_components/bestin/controller.py
index 2abc8be..d690df2 100644
--- a/custom_components/bestin/controller.py
+++ b/custom_components/bestin/controller.py
@@ -305,15 +305,17 @@ def initialize_device(self, device_id: str, sub_id: Optional[str], state: Any) -
 
         if full_unique_id not in self.devices:
             device_info = DeviceInfo(
-                id=full_unique_id,
-                type=device_type,
+                unique_id=full_unique_id,
+                device_type=device_type,
                 name=unique_id,
                 room=device_room,
                 state=state,
+                sub_type=sub_id or "",
+                colon_id=device_type if ":" in device_type else ""
             )
             device = Device(
                 info=device_info,
-                platform=platform,
+                domain=platform,
                 on_command=self.on_command,
                 callbacks=set()
             )
diff --git a/custom_components/bestin/device.py b/custom_components/bestin/device.py
index 546c26f..9ce53e5 100644
--- a/custom_components/bestin/device.py
+++ b/custom_components/bestin/device.py
@@ -4,20 +4,12 @@
 
 from typing import Any
 
-from homeassistant.helpers.restore_state import RestoreEntity
+from homeassistant.helpers.entity import Entity, DeviceInfo
 from homeassistant.core import callback
 
 from .const import DOMAIN, MAIN_DEVICES
 
 
-def split_dt(dt: str) -> str:
-    """
-    Split the first part by a colon,
-    if there is no colon, return the entire string.
-    """
-    return dt.split(":")[0].title() if ":" in dt else dt.title()
-
-
 class BestinBase:
     """Base class for BESTIN devices."""
 
@@ -29,32 +21,44 @@ def __init__(self, device, hub):
     @property
     def unique_id(self) -> str:
         """Return a unique ID."""
-        return self._device.info.id
+        return self._device.info.unique_id
 
     @property
-    def device_info(self):
-        """Return device registry information for this entity."""
-        base_info = {
-            "connections": {(self.hub.hub_id, self.unique_id)},
-            "identifiers": {(DOMAIN, f"{self.hub.wp_version}_{split_dt(self._device.info.type)}")},
-            "manufacturer": "HDC Labs Co., Ltd.",
-            "model": self.hub.wp_version,
-            "name": f"{self.hub.name} {split_dt(self._device.info.type)}",
-            "sw_version": self.hub.sw_version,
-            "via_device": (DOMAIN, self.hub.hub_id),
-        }
-        if self._device.info.type in MAIN_DEVICES:
-            base_info["identifiers"] = {(DOMAIN, f"{self.hub.wp_version}_{self.hub.model}")}
-            base_info["name"] = f"{self.hub.name}"
+    def device_type_name(self) -> str:
+        """Returns the formatted device type name."""
+        device_type = self._device.info.device_type
+        return (device_type.split(":")[0].title() 
+                if ":" in device_type else device_type.title())
 
-        return base_info
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device registry information for this entity."""
+        if self._device.info.device_type in MAIN_DEVICES:
+            return DeviceInfo(
+                connections={(self.hub.hub_id, self.unique_id)},
+                identifiers={(DOMAIN, f"{self.hub.wp_version}_{self.hub.model}")},
+                manufacturer="HDC Labs Co., Ltd.",
+                model=self.hub.wp_version,
+                name=self.hub.name,
+                sw_version=self.hub.sw_version,
+                via_device=(DOMAIN, self.hub.hub_id),
+            )
+        return DeviceInfo(
+            connections={(self.hub.hub_id, self.unique_id)},
+            identifiers={(DOMAIN, f"{self.hub.wp_version}_{self.device_type_name}")},
+            manufacturer="HDC Labs Co., Ltd.",
+            model=self.hub.wp_version,
+            name=f"{self.hub.name} {self.device_type_name}",
+            sw_version=self.hub.sw_version,
+            via_device=(DOMAIN, self.hub.hub_id),
+        )
 
     async def _on_command(self, data: Any = None, **kwargs):
         """Send commands to the device."""
         await self._device.on_command(self.unique_id, data, **kwargs)
 
 
-class BestinDevice(BestinBase, RestoreEntity):
+class BestinDevice(BestinBase, Entity):
     """Define the Bestin Device entity."""
 
     TYPE = ""
@@ -109,7 +113,7 @@ def extra_state_attributes(self) -> dict:
         attributes = {
             "unique_id": self.unique_id,
             "device_room": self._device.info.room,
-            "device_type": self._device.info.type,
+            "device_type": self._device.info.device_type,
         }
         if self.should_poll:
             attributes["last_update_time"] = self.hub.api.last_update_time
diff --git a/custom_components/bestin/fan.py b/custom_components/bestin/fan.py
index 7f2444c..ed79e74 100644
--- a/custom_components/bestin/fan.py
+++ b/custom_components/bestin/fan.py
@@ -40,7 +40,7 @@ def async_add_fan(devices=None):
         entities = [
             BestinFan(device, hub) 
             for device in devices 
-            if device.info.id not in hub.entities[DOMAIN]
+            if device.info.unique_id not in hub.entities[DOMAIN]
         ]
 
         if entities:
diff --git a/custom_components/bestin/hub.py b/custom_components/bestin/hub.py
index 469407f..f49b55a 100644
--- a/custom_components/bestin/hub.py
+++ b/custom_components/bestin/hub.py
@@ -284,7 +284,7 @@ def wp_version(self) -> str:
         if check_ip_or_serial(self.hub_id):
             return f"{self.gateway_mode[0]}-generation"
         else:
-            return f"Center-{self.version}"
+            return f"Center-v{self.version[7:8]}"
 
     @property
     def conn_str(self) -> str:
@@ -350,8 +350,8 @@ def async_add_device_callback(
         self, device_type: str, device=None, force: bool = False
     ) -> None:
         """Add device callback if not already registered."""
-        domain = device.platform.value
-        unique_id = device.info.id
+        domain = device.domain.value
+        unique_id = device.info.unique_id
         
         if (unique_id in self.entities.get(domain, [])
             or unique_id in self.hass.data[DOMAIN]):   
@@ -410,8 +410,7 @@ async def async_initialize_serial(self) -> None:
         Asynchronously initialize the Bestin Controller for serial communication.
         """
         try:
-            if self.gateway_mode is None:
-                await self.determine_gateway_mode()
+            await self.determine_gateway_mode()
 
             self.hass.config_entries.async_update_entry(
                 entry=self.entry,
@@ -447,23 +446,11 @@ async def async_initialize_center(self) -> None:
                 self.hub_id,
                 self.version,
                 self.identifier,
+                self.version == "version2.0" and self.ip_address,
                 self.async_add_device_callback,
             )
             await self.api.start()
 
-            if self.version == "version2.0" and self.ip_address:
-                if re.match(r"^10\.\d{1,3}\.\d{1,3}\.\d{1,3}$", self.ip_address):
-                    if not os.path.exists('data.json'):
-                        await self.api.elevator_call_request()
-                    else:
-                        async with aiofiles.open('data.json', 'r') as file:
-                            content = await file.read()
-                            await self.api.handle_message_info(content)  # elevator
-                else:
-                    LOGGER.warning(
-                        f"Wallpad IP address exists, but it doesn't fit the format pattern. {self.ip_address} "
-                        f"Please report it to the developer."
-                    )
         except Exception as ex:
             self.api = None
             raise RuntimeError(
diff --git a/custom_components/bestin/light.py b/custom_components/bestin/light.py
index 3dd6828..e9ae2d8 100644
--- a/custom_components/bestin/light.py
+++ b/custom_components/bestin/light.py
@@ -34,7 +34,7 @@ def async_add_light(devices=None):
         entities = [
             BestinLight(device, hub) 
             for device in devices 
-            if device.info.id not in hub.entities[DOMAIN]
+            if device.info.unique_id not in hub.entities[DOMAIN]
         ]
         
         if entities:
diff --git a/custom_components/bestin/manifest.json b/custom_components/bestin/manifest.json
index b0c370c..9dc7b1f 100644
--- a/custom_components/bestin/manifest.json
+++ b/custom_components/bestin/manifest.json
@@ -13,5 +13,5 @@
     "aiofiles",
     "xmltodict"
   ],
-  "version": "1.1.0"
+  "version": "1.1.1"
 }
\ No newline at end of file
diff --git a/custom_components/bestin/sensor.py b/custom_components/bestin/sensor.py
index f06b834..055df49 100644
--- a/custom_components/bestin/sensor.py
+++ b/custom_components/bestin/sensor.py
@@ -45,7 +45,7 @@ def async_add_sensor(devices=None):
         entities = [
             BestinSensor(device, hub) 
             for device in devices 
-            if device.info.id not in hub.entities[DOMAIN]
+            if device.info.unique_id not in hub.entities[DOMAIN]
         ]
 
         if entities:
diff --git a/custom_components/bestin/switch.py b/custom_components/bestin/switch.py
index d50e0e7..3291baa 100644
--- a/custom_components/bestin/switch.py
+++ b/custom_components/bestin/switch.py
@@ -30,7 +30,7 @@ def async_add_switch(devices=None):
         entities = [
             BestinSwitch(device, hub) 
             for device in devices 
-            if device.info.id not in hub.entities[DOMAIN]
+            if device.info.unique_id not in hub.entities[DOMAIN]
         ]
 
         if entities:
@@ -51,7 +51,7 @@ class BestinSwitch(BestinDevice, SwitchEntity):
     def __init__(self, device, hub):
         """Initialize the switch."""
         super().__init__(device, hub)
-        self._is_gas = device.info.type == "gas"
+        self._is_gas = device.info.device_type == "gas"
         self._version_exists = getattr(hub.api, "version", False)
 
     @property
diff --git a/custom_components/bestin/translations/en.json b/custom_components/bestin/translations/en.json
index c1961ed..6dbcd88 100644
--- a/custom_components/bestin/translations/en.json
+++ b/custom_components/bestin/translations/en.json
@@ -34,6 +34,7 @@
             "center_v2": {
                 "data": {
                     "ip_address": "IP address",
+                    "elevator_count": "Elevator count",
                     "uuid": "UUID"
                 },
                 "data_description": {
diff --git a/custom_components/bestin/translations/ko.json b/custom_components/bestin/translations/ko.json
index 343bf39..f1ed63b 100644
--- a/custom_components/bestin/translations/ko.json
+++ b/custom_components/bestin/translations/ko.json
@@ -34,6 +34,7 @@
             "center_v2": {
                 "data": {
                     "ip_address": "IP 주소",
+                    "elevator_count": "엘리베이터 수",
                     "uuid": "UUID"
                 },
                 "data_description": {