From 05a930d1be49aa31c0b8c8fe494b387e5454b1a9 Mon Sep 17 00:00:00 2001
From: Renato Augusto Gama dos Santos <renato_0603@hotmail.com>
Date: Tue, 7 Jan 2025 07:52:35 -0300
Subject: [PATCH] Add battery support

---
 src/config.ts            |  1 +
 src/editor.ts            | 35 +++++++++++++++++++++++++++++++++--
 src/translations/en.json |  1 +
 src/types.ts             |  1 +
 src/vacuum-card.ts       | 29 ++++++++++++++++++++++-------
 5 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/src/config.ts b/src/config.ts
index 3c95ed5e..ae04b5a1 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -21,6 +21,7 @@ export default function buildConfig(
     entity: config.entity,
     map: config.map ?? '',
     map_refresh: config.map_refresh ?? 5,
+    battery: config.battery ?? '',
     image: config.image ?? 'default',
     show_name: config.show_name ?? true,
     show_status: config.show_status ?? true,
diff --git a/src/editor.ts b/src/editor.ts
index ea10fad8..d06a31e7 100644
--- a/src/editor.ts
+++ b/src/editor.ts
@@ -35,11 +35,22 @@ export class VacuumCardEditor extends LitElement implements LovelaceCardEditor {
     }
   }
 
-  private getEntitiesByType(type: string): string[] {
+  private getEntitiesByType(type: string, deviceClass?: string): string[] {
     if (!this.hass) {
       return [];
     }
-    return Object.keys(this.hass.states).filter((id) => id.startsWith(type));
+
+    const entities = Object.keys(this.hass.states).filter((id) =>
+      id.startsWith(type),
+    );
+
+    if (deviceClass) {
+      return entities.filter(
+        (id) => this.hass?.states[id]?.attributes?.device_class === deviceClass,
+      );
+    }
+
+    return entities;
   }
 
   protected render(): Template {
@@ -48,6 +59,7 @@ export class VacuumCardEditor extends LitElement implements LovelaceCardEditor {
     }
 
     const vacuumEntities = this.getEntitiesByType('vacuum');
+    const batteryEntities = this.getEntitiesByType('sensor', 'battery');
     const cameraEntities = [
       ...this.getEntitiesByType('camera'),
       ...this.getEntitiesByType('image'),
@@ -76,6 +88,25 @@ export class VacuumCardEditor extends LitElement implements LovelaceCardEditor {
           </ha-select>
         </div>
 
+        <div class="option">
+          <ha-select
+            .label=${localize('editor.battery')}
+            @selected=${this.valueChanged}
+            .configValue=${'battery'}
+            .value=${this.config.battery}
+            @closed=${(e: Event) => e.stopPropagation()}
+            fixedMenuPosition
+            naturalMenuWidth
+          >
+            ${batteryEntities.map(
+              (entity) =>
+                html` <mwc-list-item .value=${entity}
+                  >${entity}</mwc-list-item
+                >`,
+            )}
+          </ha-select>
+        </div>
+
         <div class="option">
           <ha-select
             .label=${localize('editor.map')}
diff --git a/src/translations/en.json b/src/translations/en.json
index b5762cdd..cc17347b 100755
--- a/src/translations/en.json
+++ b/src/translations/en.json
@@ -61,6 +61,7 @@
   },
   "editor": {
     "entity": "Entity (Required)",
+    "battery": "Battery Sensor (Optional)",
     "map": "Map Camera (Optional)",
     "image": "Image (Optional)",
     "compact_view": "Compact View",
diff --git a/src/types.ts b/src/types.ts
index 7776b47f..e2c60acf 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -58,6 +58,7 @@ export interface VacuumCardShortcut {
 
 export interface VacuumCardConfig {
   entity: string;
+  battery: string;
   map: string;
   map_refresh: number;
   image: string;
diff --git a/src/vacuum-card.ts b/src/vacuum-card.ts
index 23841894..65311ef6 100755
--- a/src/vacuum-card.ts
+++ b/src/vacuum-card.ts
@@ -5,6 +5,7 @@ import {
   fireEvent,
   HomeAssistant,
   ServiceCallRequest,
+  stateIcon,
 } from 'custom-card-helpers';
 import registerTemplates from 'ha-template';
 import get from 'lodash/get';
@@ -67,7 +68,8 @@ export class VacuumCard extends LitElement {
   }
 
   get entity(): VacuumEntity {
-    return this.hass.states[this.config.entity] as VacuumEntity;
+    const vacuum = this.hass.states[this.config.entity] as VacuumEntity;
+    return vacuum;
   }
 
   get map(): HassEntity | null {
@@ -77,6 +79,13 @@ export class VacuumCard extends LitElement {
     return this.hass.states[this.config.map];
   }
 
+  get battery(): HassEntity | null {
+    if (!this.hass || !this.config.battery) {
+      return null;
+    }
+    return this.hass.states[this.config.battery];
+  }
+
   public setConfig(config: VacuumCardConfig): void {
     this.config = buildConfig(config);
   }
@@ -218,13 +227,19 @@ export class VacuumCard extends LitElement {
 
   private renderBattery(): Template {
     const { battery_level, battery_icon } = this.getAttributes(this.entity);
+    const batteryLevel = this.battery?.state || battery_level;
+    const batteryIcon = this?.battery ? stateIcon(this.battery) : battery_icon;
+
+    if (batteryLevel && batteryIcon) {
+      return html`
+        <div class="tip" @click="${() => this.handleMore()}">
+          <ha-icon icon="${batteryIcon}"></ha-icon>
+          <span class="icon-title">${batteryLevel}%</span>
+        </div>
+      `;
+    }
 
-    return html`
-      <div class="tip" @click="${() => this.handleMore()}">
-        <ha-icon icon="${battery_icon}"></ha-icon>
-        <span class="icon-title">${battery_level}%</span>
-      </div>
-    `;
+    return nothing;
   }
 
   private renderMapOrImage(state: VacuumEntityState): Template {