YouTube_predownloader/homeassistant_integration.py

#!/usr/bin/env python3
"""
Home Assistant Integration for YouTube Predownloader
Publishes real-time statistics via MQTT for Home Assistant monitoring
"""

import json
import time
import threading
import psutil
import os
import logging
from datetime import datetime
from typing import Dict, Any, Optional
import paho.mqtt.client as mqtt

# MQTT Configuration
MQTT_BROKER = "server.local.lan"  # Change to your Home Assistant IP
MQTT_PORT = 1883
MQTT_USERNAME = ""  # Set if authentication required
MQTT_PASSWORD = ""  # Set if authentication required
MQTT_BASE_TOPIC = "youtube_predownloader"

# Home Assistant Discovery Topics
HA_DISCOVERY_PREFIX = "homeassistant"

class HomeAssistantIntegration:
    def __init__(self, mqtt_broker=MQTT_BROKER, mqtt_port=MQTT_PORT):
        self.mqtt_broker = mqtt_broker
        self.mqtt_port = mqtt_port
        self.mqtt_client = None
        self.process = psutil.Process()
        self.stats = {
            'download_queue_length': 0,
            'conversion_queue_length': 0,
            'download_progress': 0,
            'conversion_progress': 0,
            'current_download_progress': 0,
            'current_download_file': '',
            'current_conversion_progress': 0,
            'current_conversion_file': '',
            'download_speed': 0,
            'cpu_usage': 0,
            'memory_usage': 0,
            'memory_usage_mb': 0,
            'active_downloads': 0,
            'active_conversions': 0,
            'total_downloads_today': 0,
            'total_conversions_today': 0,
            'service_status': 'waiting',  # waiting, running, completed, error
            'next_run_time': None,
            'current_run_start': None,
            'last_run_completed': None,
            'last_updated': None
        }
        self.is_connected = False
        self.update_thread = None
        self.running = False
        
    def connect_mqtt(self):
        """Connect to MQTT broker"""
        logger = logging.getLogger(__name__)
        
        try:
            self.mqtt_client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
            
            if MQTT_USERNAME and MQTT_PASSWORD:
                self.mqtt_client.username_pw_set(MQTT_USERNAME, MQTT_PASSWORD)
            
            self.mqtt_client.on_connect = self._on_mqtt_connect
            self.mqtt_client.on_disconnect = self._on_mqtt_disconnect
            
            self.mqtt_client.connect(self.mqtt_broker, self.mqtt_port, 60)
            self.mqtt_client.loop_start()
            
            logger.info(f"Connecting to Home Assistant MQTT broker at {self.mqtt_broker}:{self.mqtt_port}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to connect to MQTT broker: {e}")
            return False
    
    def _on_mqtt_connect(self, client, userdata, flags, reason_code, properties):
        """MQTT connection callback"""
        logger = logging.getLogger(__name__)
        
        if reason_code == 0:
            self.is_connected = True
            logger.info("Connected to Home Assistant MQTT broker")
            self._publish_discovery_configs()
        else:
            logger.error(f"Failed to connect to MQTT broker, return code {reason_code}")
    
    def _on_mqtt_disconnect(self, client, userdata, flags, reason_code, properties):
        """MQTT disconnection callback"""
        logger = logging.getLogger(__name__)
        self.is_connected = False
        logger.warning("Disconnected from MQTT broker")
    
    def _publish_discovery_configs(self):
        """Publish Home Assistant MQTT Discovery configurations"""
        logger = logging.getLogger(__name__)
        
        sensors = [
            {
                'name': 'YouTube Download Queue',
                'key': 'download_queue_length',
                'unit': 'videos',
                'icon': 'mdi:download-multiple',
                'device_class': None
            },
            {
                'name': 'YouTube Conversion Queue',
                'key': 'conversion_queue_length', 
                'unit': 'videos',
                'icon': 'mdi:video-vintage',
                'device_class': None
            },
            {
                'name': 'YouTube Download Progress',
                'key': 'download_progress',
                'unit': '%',
                'icon': 'mdi:progress-download',
                'device_class': None
            },
            {
                'name': 'YouTube Current Download Progress',
                'key': 'current_download_progress',
                'unit': '%',
                'icon': 'mdi:download',
                'device_class': None
            },
            {
                'name': 'YouTube Current Download File',
                'key': 'current_download_file',
                'unit': None,
                'icon': 'mdi:file-video',
                'device_class': None
            },
            {
                'name': 'YouTube Conversion Progress',
                'key': 'conversion_progress',
                'unit': '%', 
                'icon': 'mdi:progress-wrench',
                'device_class': None
            },
            {
                'name': 'YouTube Current Conversion Progress',
                'key': 'current_conversion_progress',
                'unit': '%',
                'icon': 'mdi:cog',
                'device_class': None
            },
            {
                'name': 'YouTube Current Conversion File',
                'key': 'current_conversion_file',
                'unit': None,
                'icon': 'mdi:file-cog',
                'device_class': None
            },
            {
                'name': 'YouTube Download Speed',
                'key': 'download_speed',
                'unit': 'MB/s',
                'icon': 'mdi:speedometer',
                'device_class': None
            },
            {
                'name': 'YouTube CPU Usage',
                'key': 'cpu_usage',
                'unit': '%',
                'icon': 'mdi:cpu-64-bit',
                'device_class': None
            },
            {
                'name': 'YouTube Memory Usage',
                'key': 'memory_usage',
                'unit': '%',
                'icon': 'mdi:memory',
                'device_class': None
            },
            {
                'name': 'YouTube Memory Usage MB',
                'key': 'memory_usage_mb',
                'unit': 'MB',
                'icon': 'mdi:memory',
                'device_class': None
            },
            {
                'name': 'YouTube Active Downloads',
                'key': 'active_downloads',
                'unit': 'downloads',
                'icon': 'mdi:download',
                'device_class': None
            },
            {
                'name': 'YouTube Active Conversions',
                'key': 'active_conversions',
                'unit': 'conversions',
                'icon': 'mdi:cog',
                'device_class': None
            },
            {
                'name': 'YouTube Service Status',
                'key': 'service_status',
                'unit': None,
                'icon': 'mdi:play-circle-outline',
                'device_class': None
            },
            {
                'name': 'YouTube Next Run Time',
                'key': 'next_run_time',
                'unit': None,
                'icon': 'mdi:clock-outline',
                'device_class': 'timestamp'
            },
            {
                'name': 'YouTube Current Run Start',
                'key': 'current_run_start',
                'unit': None,
                'icon': 'mdi:clock-start',
                'device_class': 'timestamp'
            },
            {
                'name': 'YouTube Last Run Completed',
                'key': 'last_run_completed',
                'unit': None,
                'icon': 'mdi:clock-check-outline',
                'device_class': 'timestamp'
            }
        ]
        
        discovery_count = 0
        for sensor in sensors:
            try:
                config = {
                    'name': sensor['name'],
                    'unique_id': f"{MQTT_BASE_TOPIC}_{sensor['key']}",
                    'state_topic': f"{MQTT_BASE_TOPIC}/sensor/{sensor['key']}/state",
                    'icon': sensor['icon'],
                    'device': {
                        'identifiers': [MQTT_BASE_TOPIC],
                        'name': 'YouTube Predownloader',
                        'model': 'YouTube Predownloader v1.0',
                        'manufacturer': 'Custom',
                        'sw_version': '1.0.0'
                    }
                }
                
                if sensor['unit']:
                    config['unit_of_measurement'] = sensor['unit']
                
                if sensor['device_class']:
                    config['device_class'] = sensor['device_class']
                
                discovery_topic = f"{HA_DISCOVERY_PREFIX}/sensor/{MQTT_BASE_TOPIC}_{sensor['key']}/config"
                self.mqtt_client.publish(discovery_topic, json.dumps(config), retain=True)
                discovery_count += 1
                
            except Exception as e:
                logger.error(f"Failed to publish discovery config for {sensor['name']}: {e}")
        
        logger.info(f"Published {discovery_count} Home Assistant discovery configurations")
    
    def update_stats(self, download_queue=0, conversion_queue=0, download_progress=0, 
                    conversion_progress=0, download_speed=0, active_downloads=0, active_conversions=0,
                    current_download_progress=None, current_download_file=None,
                    current_conversion_progress=None, current_conversion_file=None,
                    service_status=None, next_run_time=None, current_run_start=None, last_run_completed=None):
        """Update statistics"""
        # Update provided stats
        update_dict = {
            'download_queue_length': download_queue,
            'conversion_queue_length': conversion_queue,
            'download_progress': round(download_progress, 1),
            'conversion_progress': round(conversion_progress, 1),
            'download_speed': round(download_speed, 2),
            'active_downloads': active_downloads,
            'active_conversions': active_conversions,
            'last_updated': datetime.now().isoformat()
        }
        
        # Add current operation progress if provided
        if current_download_progress is not None:
            update_dict['current_download_progress'] = round(current_download_progress, 1)
        if current_download_file is not None:
            update_dict['current_download_file'] = str(current_download_file)[:100]  # Limit length
        if current_conversion_progress is not None:
            update_dict['current_conversion_progress'] = round(current_conversion_progress, 1)
        if current_conversion_file is not None:
            update_dict['current_conversion_file'] = str(current_conversion_file)[:100]  # Limit length
        
        # Add optional scheduling stats
        if service_status is not None:
            update_dict['service_status'] = service_status
        if next_run_time is not None:
            update_dict['next_run_time'] = next_run_time
        if current_run_start is not None:
            update_dict['current_run_start'] = current_run_start
        if last_run_completed is not None:
            update_dict['last_run_completed'] = last_run_completed
        
        self.stats.update(update_dict)
        
        # Update system stats
        try:
            self.stats['cpu_usage'] = round(self.process.cpu_percent(), 1)
            memory_info = self.process.memory_info()
            memory_percent = self.process.memory_percent()
            self.stats['memory_usage'] = round(memory_percent, 1)
            self.stats['memory_usage_mb'] = round(memory_info.rss / 1024 / 1024, 1)
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
        
        # Publish to MQTT
        if self.is_connected:
            self._publish_stats()
    
    def _publish_stats(self):
        """Publish statistics to MQTT"""
        for key, value in self.stats.items():
            if key != 'last_updated':
                topic = f"{MQTT_BASE_TOPIC}/sensor/{key}/state"
                self.mqtt_client.publish(topic, str(value))
        
        # Publish combined status
        status_topic = f"{MQTT_BASE_TOPIC}/status"
        self.mqtt_client.publish(status_topic, json.dumps(self.stats))
    
    def start_monitoring(self, update_interval=10):
        """Start monitoring and publishing stats"""
        logger = logging.getLogger(__name__)
        
        self.running = True
        self.update_thread = threading.Thread(target=self._monitoring_loop, args=(update_interval,), daemon=True)
        self.update_thread.start()
        logger.info(f"Started Home Assistant monitoring (update interval: {update_interval}s)")
    
    def stop_monitoring(self):
        """Stop monitoring"""
        logger = logging.getLogger(__name__)
        
        self.running = False
        if self.mqtt_client:
            self.mqtt_client.loop_stop()
            self.mqtt_client.disconnect()
        logger.info("Stopped Home Assistant monitoring")
    
    def _monitoring_loop(self, interval):
        """Background monitoring loop"""
        logger = logging.getLogger(__name__)
        
        while self.running:
            try:
                # Update system stats only (other stats updated via update_stats method)
                self.update_stats(
                    self.stats['download_queue_length'],
                    self.stats['conversion_queue_length'],
                    self.stats['download_progress'],
                    self.stats['conversion_progress'],
                    self.stats['download_speed'],
                    self.stats['active_downloads'],
                    self.stats['active_conversions']
                )
                time.sleep(interval)
            except Exception as e:
                logger.error(f"Monitoring error: {e}")
                time.sleep(interval)

# Global instance
ha_integration = None

def init_homeassistant_integration(mqtt_broker=MQTT_BROKER, mqtt_port=MQTT_PORT):
    """Initialize Home Assistant integration"""
    logger = logging.getLogger(__name__)
    global ha_integration
    
    if ha_integration is None:
        logger.debug(f"Initializing HA integration with broker: {mqtt_broker}:{mqtt_port}")
        ha_integration = HomeAssistantIntegration(mqtt_broker, mqtt_port)
        if ha_integration.connect_mqtt():
            ha_integration.start_monitoring()
            logger.info("Home Assistant integration initialized successfully")
            return ha_integration
        else:
            logger.error("Home Assistant integration failed to connect")
            ha_integration = None
            return None
    
    logger.debug("Home Assistant integration already initialized")
    return ha_integration

def update_homeassistant_stats(**kwargs):
    """Update Home Assistant statistics"""
    global ha_integration
    
    if ha_integration and ha_integration.is_connected:
        ha_integration.update_stats(**kwargs)

def update_service_status(status, next_run=None, current_run_start=None, last_completed=None):
    """Update service scheduling status"""
    global ha_integration
    
    if ha_integration and ha_integration.is_connected:
        ha_integration.update_stats(
            service_status=status,
            next_run_time=next_run,
            current_run_start=current_run_start,
            last_run_completed=last_completed
        )

def set_service_waiting(next_run_time):
    """Set service status to waiting with next run time"""
    update_service_status('waiting', next_run=next_run_time)

def set_service_running():
    """Set service status to running"""
    update_service_status('running', current_run_start=datetime.now().isoformat())

def set_service_completed():
    """Set service status to completed"""
    update_service_status('completed', last_completed=datetime.now().isoformat())

def set_service_error():
    """Set service status to error"""
    update_service_status('error')

def stop_homeassistant_integration():
    """Stop Home Assistant integration"""
    global ha_integration
    
    if ha_integration:
        ha_integration.stop_monitoring()
        ha_integration = None

if __name__ == "__main__":
    # Test the integration
    print("🧪 Testing Home Assistant integration...")
    
    integration = init_homeassistant_integration()
    
    if integration:
        print("✅ Integration started successfully")
        
        # Simulate some data
        for i in range(5):
            update_homeassistant_stats(
                download_queue=i,
                conversion_queue=i*2,
                download_progress=i*20,
                conversion_progress=i*15,
                download_speed=i*2.5,
                active_downloads=min(i, 2),
                active_conversions=min(i, 2)
            )
            time.sleep(2)
        
        stop_homeassistant_integration()
        print("✅ Test completed")
    else:
        print("❌ Failed to start integration")