#!/usr/bin/env python3
"""
MQTT/Home Assistant Discovery Integration
Erweitert das Tahoma Script um MQTT-Geräte via Home Assistant Discovery
"""

import paho.mqtt.client as mqtt
import json
import time
from typing import Dict, List, Optional
import logging

logger = logging.getLogger(__name__)


class HomeAssistantDiscovery:
    """Klasse für Home Assistant MQTT Discovery"""
    
    # Bekannte Discovery-Komponenten
    COMPONENTS = [
        'binary_sensor', 'sensor', 'switch', 'light', 'cover', 
        'climate', 'fan', 'lock', 'camera', 'vacuum', 'alarm_control_panel',
        'device_tracker', 'number', 'select', 'button', 'text'
    ]
    
    def __init__(self, broker: str, port: int = 1883, username: str = None, 
                 password: str = None, discovery_prefix: str = 'homeassistant'):
        """
        Initialisiert Home Assistant Discovery
        
        Args:
            broker: MQTT Broker IP/Hostname
            port: MQTT Port (Standard: 1883)
            username: MQTT Benutzername (optional)
            password: MQTT Passwort (optional)
            discovery_prefix: Discovery Prefix (Standard: 'homeassistant')
        """
        self.broker = broker
        self.port = port
        self.username = username
        self.password = password
        self.discovery_prefix = discovery_prefix
        self.client = None
        self.discovered_devices = {}
        
    def connect(self) -> bool:
        """
        Verbindet mit dem MQTT Broker
        
        Returns:
            True bei Erfolg, False bei Fehler
        """
        try:
            self.client = mqtt.Client()
            
            if self.username and self.password:
                self.client.username_pw_set(self.username, self.password)
            
            self.client.on_connect = self._on_connect
            self.client.on_message = self._on_message
            
            self.client.connect(self.broker, self.port, 60)
            logger.info(f"Verbunden mit MQTT Broker {self.broker}:{self.port}")
            return True
            
        except Exception as e:
            logger.error(f"MQTT Verbindungsfehler: {e}")
            return False
    
    def _on_connect(self, client, userdata, flags, rc):
        """Callback wenn Verbindung hergestellt wurde"""
        if rc == 0:
            logger.info("MQTT Verbindung erfolgreich")
            # Alle Discovery Topics abonnieren mit Wildcard für object_id
            for component in self.COMPONENTS:
                # Unterstützt beide Topic-Formate:
                # homeassistant/component/node_id/config (4 Teile)
                # homeassistant/component/node_id/object_id/config (5 Teile)
                topic = f"{self.discovery_prefix}/{component}/+/+/config"
                client.subscribe(topic)
                logger.debug(f"Abonniert: {topic}")
                # Zusätzlich auch das kürzere Format abonnieren
                topic_short = f"{self.discovery_prefix}/{component}/+/config"
                client.subscribe(topic_short)
                logger.debug(f"Abonniert: {topic_short}")
        else:
            logger.error(f"MQTT Verbindung fehlgeschlagen, Code: {rc}")
    
    def _on_message(self, client, userdata, msg):
        """Callback wenn Nachricht empfangen wurde"""
        try:
            # Topic analysieren - unterstützt beide Formate:
            # homeassistant/component/node_id/config (4 Teile)
            # homeassistant/component/node_id/object_id/config (5 Teile)
            topic_parts = msg.topic.split('/')
            
            if topic_parts[-1] != 'config':
                return  # Kein Config-Topic
            
            if len(topic_parts) == 4:
                # Format: homeassistant/component/node_id/config
                component = topic_parts[1]
                node_id = topic_parts[2]
                object_id = None
            elif len(topic_parts) == 5:
                # Format: homeassistant/component/node_id/object_id/config
                component = topic_parts[1]
                node_id = topic_parts[2]
                object_id = topic_parts[3]
            else:
                logger.debug(f"Unbekanntes Topic-Format: {msg.topic}")
                return
            
            # Payload parsen
            if msg.payload:
                config = json.loads(msg.payload.decode('utf-8'))
                
                # Eindeutigen Key erstellen
                if object_id:
                    device_key = f"{component}_{node_id}_{object_id}"
                else:
                    device_key = f"{component}_{node_id}"
                
                # Gerät speichern
                self.discovered_devices[device_key] = {
                    'component': component,
                    'node_id': node_id,
                    'object_id': object_id,
                    'config': config,
                    'topic': msg.topic
                }
                
                device_name = config.get('name', config.get('unique_id', object_id or node_id))
                logger.debug(f"Gerät gefunden: {device_name} ({component}) - {msg.topic}")
                
        except Exception as e:
            logger.error(f"Fehler beim Verarbeiten der MQTT-Nachricht von {msg.topic}: {e}")
    
    def discover_devices(self, timeout: int = 10) -> Dict:
        """
        Sucht nach Home Assistant Discovery Geräten
        
        Args:
            timeout: Timeout in Sekunden
            
        Returns:
            Dictionary mit gefundenen Geräten
        """
        logger.info(f"Starte Home Assistant Discovery (Timeout: {timeout}s)...")
        logger.info(f"Lausche auf {self.discovery_prefix}/+/+/+/config und {self.discovery_prefix}/+/+/config")
        
        self.discovered_devices = {}
        
        # MQTT Loop starten
        self.client.loop_start()
        
        # Warten auf Nachrichten - mit Fortschrittsanzeige
        for i in range(timeout):
            time.sleep(1)
            if (i + 1) % 5 == 0 or i == timeout - 1:
                logger.info(f"Discovery läuft... {len(self.discovered_devices)} Geräte gefunden ({i+1}/{timeout}s)")
        
        # Loop stoppen
        self.client.loop_stop()
        
        logger.info(f"✓ {len(self.discovered_devices)} MQTT-Geräte gefunden")
        
        # Debug: Zeige einige gefundene Topics
        if self.discovered_devices:
            logger.debug("Gefundene Geräte (Auswahl):")
            for i, (key, device) in enumerate(list(self.discovered_devices.items())[:5]):
                logger.debug(f"  - {device['config'].get('name', key)} ({device['component']}) via {device['topic']}")
            if len(self.discovered_devices) > 5:
                logger.debug(f"  ... und {len(self.discovered_devices) - 5} weitere")
        
        return self.discovered_devices
    
    def disconnect(self):
        """Trennt die MQTT-Verbindung"""
        if self.client:
            self.client.disconnect()
            logger.info("MQTT-Verbindung getrennt")


class MQTTDeviceConverter:
    """Konvertiert MQTT Discovery Entities in Datenbank-Format, gruppiert nach Gerät"""
    
    # Mapping von HA Komponenten zu Actor/Sensor
    ACTOR_COMPONENTS = ['switch', 'light', 'cover', 'fan', 'lock', 'climate', 
                        'vacuum', 'alarm_control_panel', ' ', 'number', 'select']
    SENSOR_COMPONENTS = ['binary_sensor', 'sensor', 'device_tracker']
    
    @staticmethod
    def group_entities_by_device(discovered_devices: Dict) -> Dict[str, List]:
        """
        Gruppiert Discovery-Entities nach Gerät (node_id)
        
        Args:
            discovered_devices: Dictionary mit allen gefundenen Entities
            
        Returns:
            Dictionary: {device_id: [entity1, entity2, ...]}
        """
        devices = {}
        
        for entity_key, entity in discovered_devices.items():
            # Device Identifier aus Config extrahieren
            config = entity.get('config', {})
            device_info = config.get('device') or config.get('dev') or {}
            # Node ID als Geräte-Identifier verwenden
            node_id = entity.get('node_id', 'unknown')
            
            # Zusätzlich Device Identifiers prüfen falls vorhanden
            if device_info and 'identifiers' in device_info:
                identifiers = device_info['identifiers']
                if isinstance(identifiers, list) and identifiers:
                    node_id = identifiers[0]
            
            if node_id not in devices:
                devices[node_id] = {
                    'entities': [],
                    'device_info': device_info,
                    'node_id': node_id
                }
            
            devices[node_id]['entities'].append(entity)
        
        return devices
    
    @staticmethod
    def is_actor_entity(component: str) -> bool:
        """Prüft ob Entity-Komponente ein Aktor ist"""
        if(component == "climate"):
           logger.info("Climate device gefunden"); 
        return component in MQTTDeviceConverter.ACTOR_COMPONENTS
    
    @staticmethod
    def is_sensor_entity(component: str) -> bool:
        """Prüft ob Entity-Komponente ein Sensor ist"""
        return component in MQTTDeviceConverter.SENSOR_COMPONENTS
    
    @staticmethod
    def convert_device_to_actors_and_sensors(device_id: str, device_data: Dict) -> tuple:
        """
        Konvertiert ein Gerät mit allen seinen Entities in Actor/Sensor-Format
        
        Args:
            device_id: Geräte-ID (node_id)
            device_data: Device-Daten mit Entities-Liste
            
        Returns:
            Tuple (actor_dict or None, sensor_dict or None)
        """
        entities = device_data['entities']
        device_info = device_data.get('device_info')
        # Gerätename vom ersten Entity oder aus device_info
        device_name = device_info.get('name', device_id)
        if not device_name or device_name == device_id:
            # Fallback: Name vom ersten Entity
            if entities:
                device_name = entities[0]['config'].get('name', device_id)
        
        # Device URL
        device_url = f"mqtt://{device_id}"
        
        # Entities nach Actor/Sensor trennen
        actor_entities = [e for e in entities if MQTTDeviceConverter.is_actor_entity(e['component'])]
        sensor_entities = [e for e in entities if MQTTDeviceConverter.is_sensor_entity(e['component'])]
        
        actor_result = None
        sensor_result = None
        
        # Actor erstellen falls Actor-Entities vorhanden
        if actor_entities:
            commands = []
            states = []
            
            for entity in actor_entities:
                component = entity['component']
                object_id = entity.get('object_id', 'unknown')
                config = entity['config']
                
                # Command aus der Entity erstellen
                command_entry = MQTTDeviceConverter._entity_to_command(component, object_id, config)
                if command_entry:
                    commands.append(command_entry)
                
                # States aus der Entity extrahieren
                entity_states = MQTTDeviceConverter._entity_to_states(component, object_id, config)
                states.extend(entity_states)
            
            if commands:  # Nur Actor erstellen wenn Commands vorhanden
                actor_result = {
                    'name': device_name,
                    'type': f"mqtt_device",  # Allgemeiner Typ für Multi-Entity-Geräte
                    'url': device_url,
                    'commands': commands,
                    'states': states
                }
        
        # Sensor erstellen falls Sensor-Entities vorhanden
        if sensor_entities:
            states = []
            
            for entity in sensor_entities:
                component = entity['component']
                object_id = entity.get('object_id', 'unknown')
                config = entity['config']
                
                # States aus der Entity extrahieren
                entity_states = MQTTDeviceConverter._entity_to_states(component, object_id, config)
                states.extend(entity_states)
            
            if states:  # Nur Sensor erstellen wenn States vorhanden
                sensor_result = {
                    'name': device_name,
                    'type': f"mqtt_device",
                    'url': device_url,
                    'states': states
                }
        
        return actor_result, sensor_result
    
    @staticmethod
    def _entity_to_command(component: str, object_id: str, config: Dict) -> Optional[Dict]:
        """
        Konvertiert eine MQTT Entity in ein Command
        
        Args:
            component: Entity-Typ (number, button, switch, etc.)
            object_id: Object ID (z.B. set_max_ampere_limit)
            config: Entity-Konfiguration
            
        Returns:
            Command-Dictionary oder None
        """
        # Command Topic - verschiedene mögliche Feldnamen
        command_topic = config.get('command_topic') or config.get('cmd_t') or config.get('temperature_command_topic')
        
        if not command_topic:
            return None
        
        # Command-Name aus object_id ableiten
        command_name = object_id.replace('_', ' ').title().replace(' ', '')
        # Oder aus dem Namen
        entity_name = config.get('name', object_id)
        
        command_entry = {
            'command': command_name,
            'parameters': []
        }
        
        # Parameter basierend auf Component-Typ
        if component == 'number':
            # Number hat einen Wert-Parameter
            param = {
                'name': 'value',
                'type': 'number',
                'url': command_topic
            }
            
            # Min/Max aus Config
            if 'min' in config:
                param['min'] = config['min']
            if 'max' in config:
                param['max'] = config['max']
            
            # Unit hinzufügen - verschiedene mögliche Feldnamen
            unit = (
                config.get('unit_of_measurement') or
                config.get('unit_of_meas') or
                config.get('unit') or
                config.get('u')
            )
            if unit:
                param['unit'] = unit
            
            command_entry['parameters'].append(param)
        
        elif component == 'select':
            # Select hat Optionen
            param = {
                'name': 'option',
                'type': 'string',
                'url': command_topic
            }
            
            options = config.get('options') or config.get('ops')
            if options:
                param['values'] = options
            
            command_entry['parameters'].append(param)
        
        elif component in ['switch', 'light']:
            # Switch/Light haben on/off
            param = {
                'name': 'state',
                'type': 'string',
                'url': command_topic,
                'values': [
                    config.get('payload_on', config.get('pl_on', 'ON')),
                    config.get('payload_off', config.get('pl_off', 'OFF'))
                ]
            }
            command_entry['parameters'].append(param)
            
            # Brightness für Light
            brightness_cmd_topic = (
                config.get('brightness_command_topic') or
                config.get('bri_cmd_t')
            )
            if component == 'light' and brightness_cmd_topic:
                command_entry['parameters'].append({
                    'name': 'brightness',
                    'type': 'integer',
                    'min': 0,
                    'max': 255,
                    'url': brightness_cmd_topic
                })
        
        elif component == 'cover':
            # Cover hat position
            set_pos_topic = (
                config.get('set_position_topic') or
                config.get('pos_cmd_t')
            )
            if set_pos_topic:
                param = {
                    'name': 'position',
                    'type': 'integer',
                    'min': 0,
                    'max': 100,
                    'url': set_pos_topic
                }
                command_entry['parameters'].append(param)
            else:
                # Nur open/close/stop
                param = {
                    'name': 'action',
                    'type': 'string',
                    'url': command_topic,
                    'values': ['OPEN', 'CLOSE', 'STOP']
                }
                command_entry['parameters'].append(param)
        
        elif component == 'button':
            # Button hat normalerweise keinen Parameter, nur das Topic
            param = {
                'name': 'press',
                'type': 'trigger',
                'url': command_topic
            }
            command_entry['parameters'].append(param)
        
        elif component == 'climate':
            # Climate hat Temperatur-Setpoint
            temp_cmd_topic = (
                config.get('temperature_command_topic') or
                config.get('temp_cmd_t')
            )
            if temp_cmd_topic:
                param = {
                    'name': 'temperature',
                    'type': 'number',
                    'url': temp_cmd_topic
                }
                
                if 'min_temp' in config:
                    param['min'] = config['min_temp']
                if 'max_temp' in config:
                    param['max'] = config['max_temp']
                
                command_entry['parameters'].append(param)
            mode_cmd_topic = (
                config.get('mode_command_topic') or
                config.get('mode_cmd_t')
            )
            if mode_cmd_topic:
                param = {
                    'name': 'mode',
                    'type': 'string',
                    'url': temp_cmd_topic,
                    'values': config.get('modes', [])
                }
        else:
            # Generischer Command mit dem Topic
            param = {
                'name': 'value',
                'type': 'string',
                'url': command_topic
            }
            command_entry['parameters'].append(param)
        
        return command_entry
    
    @staticmethod
    def _entity_to_states(component: str, object_id: str, config: Dict) -> List[Dict]:
        """
        Extrahiert States aus einer MQTT Entity
        
        Args:
            component: Entity-Typ
            object_id: Object ID
            config: Entity-Konfiguration
            
        Returns:
            Liste von State-Dictionaries
        """
        states = []
        
        # State Topic - verschiedene mögliche Feldnamen prüfen
        state_topic = (
            config.get('state_topic') or 
            config.get('stat_t') or  # Abkürzung
            config.get('~') and config.get('stat_t')  # Mit Base Topic
        )
        
        # Bei number/select: oft kein separates state_topic, dann command_topic verwenden
        if not state_topic and component in ['number', 'select', 'button']:
            # Bei diesen Komponenten kann der State über command_topic abgefragt werden
            # oder es gibt ein explizites state_topic
            state_topic = config.get('command_topic') or config.get('cmd_t')
        
        if state_topic:
            state_entry = {
                'name': object_id,
                'type': 'string',
                'url': state_topic
            }
            
            # Unit hinzufügen - verschiedene mögliche Feldnamen
            unit = (
                config.get('unit_of_measurement') or
                config.get('unit_of_meas') or
                config.get('unit') or
                config.get('u')  # Weitere Abkürzung
            )
            if unit:
                state_entry['unit'] = unit
            
            # Device Class als zusätzliche Info
            if 'device_class' in config:
                state_entry['device_class'] = config['device_class']
            elif 'dev_cla' in config:
                state_entry['device_class'] = config['dev_cla']
            
            # Typ anpassen basierend auf Component
            if component == 'number':
                state_entry['type'] = 'number'
            elif component == 'binary_sensor':
                state_entry['type'] = 'boolean'
            elif component == 'sensor':
                # Bei Sensor den Typ aus value_template ableiten oder number annehmen
                state_entry['type'] = 'number'  # Default für Sensoren
            
            states.append(state_entry)
        
        # Zusätzliche State Topics (z.B. brightness bei Light)
        if component == 'light':
            brightness_topic = (
                config.get('brightness_state_topic') or
                config.get('bri_stat_t')
            )
            if brightness_topic:
                states.append({
                    'name': f"{object_id}_brightness",
                    'type': 'integer',
                    'url': brightness_topic
                })
        
        if component == 'cover':
            position_topic = (
                config.get('position_topic') or
                config.get('pos_t')
            )
            if position_topic:
                states.append({
                    'name': f"{object_id}_position",
                    'type': 'integer',
                    'url': position_topic
                })
        
        if component == 'climate':
            current_temp_topic = (
                config.get('current_temperature_topic') or
                config.get('curr_temp_t')
            )
            if current_temp_topic:
                states.append({
                    'name': f"{object_id}_current_temp",
                    'type': 'number',
                    'unit': '°C',
                    'url': current_temp_topic
                })
        
        return states