"""MQTT client for Heiman device communication.

Handles MQTT communication with Heiman Cloud for real-time
device property updates and event handling.
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import socket
import ssl
import time
from typing import Any, Callable

try:
    import paho.mqtt.client as mqtt

    PAHO_AVAILABLE = True
except ImportError:
    PAHO_AVAILABLE = False
    mqtt = None

from .models import HeimanDevice

_LOGGER = logging.getLogger(__name__)

# MQTT 配置常量
DEFAULT_MQTT_BROKER = "spmqtt.heiman.cn"
DEFAULT_MQTT_PORT_SSL = 1884
DEFAULT_SECURE_ID = "htJXYn5TyM3zZ7ji"
DEFAULT_SECURE_KEY = "htJXYn5TyM3zZ7ji"


class HeimanMqttError(Exception):
    """Base exception for Heiman MQTT errors."""



class HeimanMqttClient:
    """Heiman MQTT client for real-time device communication.
    
    This client connects to Heiman MQTT broker and subscribes to device topics
    for receiving real-time property updates and events.
    """

    def __init__(
        self,
        hass,
        access_token: str,
        user_id: str,
        secure_id: str = DEFAULT_SECURE_ID,
        secure_key: str = DEFAULT_SECURE_KEY,
        mqtt_broker: str = DEFAULT_MQTT_BROKER,
        mqtt_port: int = DEFAULT_MQTT_PORT_SSL,
        user_display_name: str | None = None,
        cloud_client: Any = None,
        devices: dict[str, Any] | None = None,  # New parameter
    ) -> None:
        """Initialize MQTT client.
        
        Args:
            hass: Home Assistant instance
            access_token: OAuth2 access token for authentication
            user_id: User ID for topic subscription
            secure_id: Secure ID for MQTT authentication
            secure_key: Secure key for MQTT password generation
            mqtt_broker: MQTT broker hostname
            mqtt_port: MQTT broker port
            user_display_name: User display name for device control payload
            cloud_client: Cloud client instance for accessing devices data
            devices: Dictionary of devices for child device detection
        """
        if not PAHO_AVAILABLE:
            raise ImportError(
                "paho-mqtt is required. Install with: pip install paho-mqtt",
            )

        self._hass = hass
        self._access_token = access_token
        self._user_id = user_id
        self._secure_id = secure_id
        self._secure_key = secure_key
        self._broker = mqtt_broker
        self._port = mqtt_port
        self._user_display_name = user_display_name
        self._cloud_client = cloud_client
        self._devices = devices or {}  # Store devices dictionary
        
        # Generate MQTT credentials
        self._client_id = access_token
        self._username = f"appid_{secure_id}"
        self._password = self._generate_password(secure_id, secure_key)
        
        # MQTT client
        self._client: mqtt.Client | None = None
        self._connected = False
        self._connecting = False
        self._should_reconnect = True
        
        # Message tracking for requests
        self._pending_requests: dict[str, asyncio.Future] = {}
        self._request_timeout = 10  # seconds
        
        # Subscribed topics
        self._subscribed_topics: set[str] = set()
        
        # Device callbacks - called when device properties are updated
        self._device_callbacks: list[Callable[[str, dict], None]] = []
        
        # Event loop
        self._loop = asyncio.get_running_loop() if hasattr(asyncio, 'get_running_loop') else asyncio.get_event_loop()
        
        _LOGGER.info("Initialized Heiman MQTT client for user %s", user_id)
    
    def _generate_password(self, secure_id: str, secure_key: str) -> str:
        """Generate MD5 password for authentication."""
        md5_str = f"{secure_id}|{secure_key}"
        return hashlib.md5(md5_str.encode("utf-8")).hexdigest()
    
    def _generate_message_id(self) -> str:
        """Generate a unique message ID."""
        return f"{int(time.time() * 1000)}"
    
    def _is_child_device(self, device_info: dict | None) -> bool:
        """Check if device is a child device (gateway sub-device).
        
        Args:
            device_info: Device information dictionary or HeimanDevice object
            
        Returns:
            True if device is a child device, False otherwise
        """
        if not device_info:
            return False
        
        # Handle both dict and HeimanDevice object
        if isinstance(device_info, dict):
            device_type = device_info.get("deviceType")
        else:
            # HeimanDevice object - check raw_data if available
            device_type = getattr(device_info, 'device_type', None)
            if device_type is None and hasattr(device_info, 'raw_data'):
                device_type = device_info.raw_data.get("deviceType") if device_info.raw_data else None
        
        if isinstance(device_type, dict):
            return device_type.get("value") == "childrenDevice"
        return False
    
    def _build_write_topic(
        self,
        product_id: str,
        device_id: str,
        device_info: dict | HeimanDevice | None,
    ) -> str:
        """Build MQTT topic for writing device properties.
        
        For child devices (gateway sub-devices), the topic format is:
        /{parentProductId}/{parentDeviceId}/child/{childDeviceId}/properties/write
        
        For regular devices, the topic format is:
        /{product_id}/{device_id}/properties/write
        
        Args:
            product_id: Device product ID
            device_id: Device ID
            device_info: Device information dictionary or HeimanDevice object
            
        Returns:
            MQTT topic string
        """
        if self._is_child_device(device_info):
            # Get parent device info
            # Handle both dict and HeimanDevice object
            if isinstance(device_info, dict):
                parent_id = device_info.get("parentId")
            else:
                # HeimanDevice object - check raw_data
                parent_id = getattr(device_info, 'parent_id', None)
                if parent_id is None and hasattr(device_info, 'raw_data'):
                    parent_id = device_info.raw_data.get("parentId") if device_info.raw_data else None
            
            if parent_id:
                # Get parent device info from local devices dictionary
                parent_device = self._devices.get(parent_id)
                
                if parent_device:
                    # Handle both dict and HeimanDevice object
                    if isinstance(parent_device, dict):
                        parent_product_id = parent_device.get("productId")
                    else:
                        # HeimanDevice object
                        parent_product_id = getattr(parent_device, 'product_id', None)
                    
                    if parent_product_id:
                        # Child device topic format
                        topic = f"/{parent_product_id}/{parent_id}/child/{product_id}/properties/write"
                        _LOGGER.info(
                            "[CHILD-DEVICE] Using child device topic for %s: %s (parent: %s)",
                            device_id,
                            topic,
                            parent_id,
                        )
                        return topic
                    
                    _LOGGER.warning(
                        "[CHILD-DEVICE] Parent device %s found but productId missing",
                        parent_id,
                    )
                else:
                    _LOGGER.warning(
                        "[CHILD-DEVICE] Parent device %s not found in self._devices",
                        parent_id,
                    )
            else:
                _LOGGER.warning(
                    "[CHILD-DEVICE] Cannot get parent info: parent_id is missing",
                )
        
        # Regular device topic format
        topic = f"/{product_id}/{device_id}/properties/write"
        _LOGGER.debug("[REGULAR-DEVICE] Using regular topic for %s: %s", device_id, topic)
        return topic
    
    async def connect(self) -> bool:
        """Connect to MQTT broker.
        
        Returns:
            True if connection successful
            
        Raises:
            HeimanMqttError: If connection fails
        """
        if self._connected or self._connecting:
            _LOGGER.debug("MQTT already connected or connecting")
            return True
        
        self._connecting = True
        
        try:
            _LOGGER.info("Connecting to MQTT broker %s:%d", self._broker, self._port)
            
            # Create MQTT client
            self._client = mqtt.Client(
                client_id=self._client_id,
                protocol=mqtt.MQTTv311,
                transport="tcp",
            )
            
            # Enable SSL/TLS for SSL port
            if self._port == DEFAULT_MQTT_PORT_SSL:
                _LOGGER.debug("Configuring SSL/TLS connection")
                
                # Use executor for blocking TLS configuration
                def configure_tls():
                    self._client.tls_set(
                        ca_certs=None,
                        certfile=None,
                        keyfile=None,
                        cert_reqs=mqtt.ssl.CERT_REQUIRED,
                        tls_version=mqtt.ssl.PROTOCOL_TLS,
                        ciphers=None,
                    )
                    self._client.tls_insecure_set(False)
                
                await self._hass.async_add_executor_job(configure_tls)
            
            # Set callbacks
            self._client.on_connect = self._on_connect
            self._client.on_disconnect = self._on_disconnect
            self._client.on_message = self._on_message
            
            # Set username/password
            self._client.username_pw_set(self._username, self._password)
            
            # Connect
            self._client.connect_async(self._broker, self._port, 60)
            self._client.loop_start()
            
            _LOGGER.info("MQTT client started, waiting for connection...")
            
            # Wait for connection (max 30 seconds)
            max_wait_time = 30
            check_interval = 0.2
            total_checks = int(max_wait_time / check_interval)
            
            for i in range(total_checks):
                if self._connected:
                    elapsed_time = (i + 1) * check_interval
                    _LOGGER.info("MQTT connected successfully after %.1f seconds", elapsed_time)
                    break
                await asyncio.sleep(check_interval)
            
            if not self._connected:
                raise HeimanMQTTError(f"Connection timeout after {max_wait_time}s")
            
            return True
            
        except Exception as err:
            self._connecting = False
            raise HeimanMQTTError(f"Connection failed: {err}") from err

    async def disconnect(self) -> None:
        """Disconnect from MQTT broker."""
        if not self._client:
            return
        
        self._should_reconnect = False
        
        try:
            self._client.loop_stop()
            self._client.disconnect()
            _LOGGER.info("Disconnected from MQTT broker")
        except Exception as err:
            _LOGGER.warning("Error during MQTT disconnect: %s", err)
        finally:
            self._connected = False
            self._client = None
    
    def _on_connect(self, client, userdata, flags, rc) -> None:
        """Called when MQTT client connects."""
        if rc == 0:
            self._connected = True
            self._connecting = False
            _LOGGER.info("MQTT client connected")
            
            # Subscribe to user-specific app topic
            if self._user_id:
                app_topic = f"/iot/app/11/{self._user_id}/#"
                client.subscribe(app_topic)
                self._subscribed_topics.add(app_topic)
                _LOGGER.info("Subscribed to app topic: %s", app_topic)
            else:
                _LOGGER.warning("User ID not available for MQTT subscription")
        else:
            error_messages = {
                0: "Connection accepted",
                1: "Connection refused: unacceptable protocol version",
                2: "Connection refused: identifier rejected",
                3: "Connection refused: server unavailable",
                4: "Connection refused: bad user name or password",
                5: "Connection refused: not authorized",
            }
            _LOGGER.error("MQTT connection failed with code %s: %s", rc, error_messages.get(rc, "Unknown error"))
            self._connecting = False
    
    def _on_disconnect(self, client, userdata, rc) -> None:
        """Called when MQTT client disconnects."""
        self._connected = False
        _LOGGER.info("MQTT client disconnected (rc=%s)", rc)
        
        # Auto-reconnect if needed
        if self._should_reconnect and rc != 0:
            _LOGGER.info("Attempting to reconnect...")
            try:
                client.reconnect()
            except Exception as err:
                _LOGGER.error("Reconnect failed: %s", err)
    
    def _on_message(self, client, userdata, msg) -> None:
        """Called when a message is received."""
        try:
            topic = msg.topic
            raw_payload = msg.payload.decode("utf-8")
            payload = json.loads(raw_payload)
            
            _LOGGER.debug("Received MQTT message on topic %s: %s", topic, raw_payload[:200])
            
            # Parse topic to extract device info
            topic_parts = topic.split("/")
            
            # Handle property report messages
            if "/properties/report" in topic or "/properties/report" in raw_payload:
                self._handle_property_report(topic, payload)
            elif "/properties/write/reply" in topic:
                # Handle write property reply (confirmation)
                self._handle_write_reply(topic, payload)
            elif "/event/" in topic:
                self._handle_event(topic, payload)
                
        except json.JSONDecodeError as err:
            _LOGGER.error("Failed to parse MQTT message: %s", err)
        except Exception as err:
            _LOGGER.error("Error handling MQTT message: %s", err)
    
    def _handle_property_report(self, topic: str, payload: dict) -> None:
        """Handle device property report."""
        try:
            device_id = payload.get("deviceId")
            properties = payload.get("properties", {})
            
            if device_id and properties:
                _LOGGER.info("Property report for device %s: %s", device_id, properties)
                
                # Call registered callbacks
                for callback in self._device_callbacks:
                    try:
                        callback(device_id, properties)
                    except Exception as err:
                        _LOGGER.error("Error in property callback: %s", err)
        except Exception as err:
            _LOGGER.error("Error handling property report: %s", err)
    
    def _handle_write_reply(self, topic: str, payload: dict) -> None:
        """Handle write property reply (confirmation from cloud).
        
        This is called when the cloud confirms a property write operation.
        The reply contains the updated property values which should be used
        to update the coordinator cache and notify entities.
        
        Args:
            topic: MQTT topic string
            payload: Decoded message payload
        """
        try:
            device_id = payload.get("deviceId")
            properties = payload.get("properties", {})
            success = payload.get("success", 0)
            
            if device_id and properties:
                _LOGGER.info(
                    "Write reply for device %s: properties=%s, success=%s",
                    device_id,
                    properties,
                    success,
                )
                
                # Only process successful writes
                if success == 1:
                    # Call registered callbacks to update coordinator cache
                    for callback in self._device_callbacks:
                        try:
                            callback(device_id, properties)
                        except Exception as err:
                            _LOGGER.error("Error in write reply callback: %s", err)
                else:
                    _LOGGER.warning(
                        "Write operation failed for device %s: %s",
                        device_id,
                        payload,
                    )
        except Exception as err:
            _LOGGER.error("Error handling write reply: %s", err)
    
    def _handle_event(self, topic: str, payload: dict) -> None:
        """Handle device event."""
        try:
            device_id = payload.get("deviceId")
            event_type = payload.get("eventType", "unknown")
            
            if device_id:
                _LOGGER.info("Event from device %s: %s", device_id, event_type)
                
                # Call registered callbacks
                for callback in self._device_callbacks:
                    try:
                        callback(device_id, {"event": event_type, **payload})
                    except Exception as err:
                        _LOGGER.error("Error in event callback: %s", err)
        except Exception as err:
            _LOGGER.error("Error handling event: %s", err)

    def subscribe(
        self,
        topic: str,
        callback: Callable[[str, Any], None],
    ) -> None:
        """Subscribe to MQTT topic with callback.
        
        Args:
            topic: Topic pattern to subscribe to
            callback: Callback function to invoke on message
        """
        self._subscribed_topics.add(topic)
        
        if self._connected and self._client:
            try:
                self._client.subscribe(topic)
                _LOGGER.debug("Subscribed to topic: %s", topic)
            except Exception as err:
                _LOGGER.error("Failed to subscribe to topic %s: %s", topic, err)
    
    def register_device_callback(self, callback: Callable[[str, dict], None]) -> None:
        """Register callback for device property updates.
        
        Args:
            callback: Function to call when device properties are updated
        """
        self._device_callbacks.append(callback)
        _LOGGER.debug("Registered device callback")
    
    def unregister_device_callback(self, callback: Callable[[str, dict], None]) -> None:
        """Unregister device callback.
        
        Args:
            callback: Callback to remove
        """
        if callback in self._device_callbacks:
            self._device_callbacks.remove(callback)
            _LOGGER.debug("Unregistered device callback")
    
    @property
    def is_connected(self) -> bool:
        """Check if MQTT client is connected."""
        return self._connected
    
    async def async_read_properties(self, device_id: str, product_id: str, property_identifiers: list[str] | None = None) -> dict[str, Any]:
        """Read device properties via MQTT.
        
        Args:
            device_id: Device ID to read from
            product_id: Product ID for topic construction
            property_identifiers: List of property identifiers to read (reads all if None)
            
        Returns:
            Dictionary of property values
            
        Raises:
            HeimanMqttError: If read fails or times out
        """
        if not self._connected:
            raise HeimanMQTTError("Not connected to MQTT broker")
        
        # Generate unique message ID
        message_id = self._generate_message_id()
        
        # Build request payload
        request_payload = {
            "messageId": message_id,
            "timestamp": int(time.time() * 1000),
            "deviceId": device_id,
            "properties": property_identifiers or [],  # Empty means read all
        }
        
        # Build topic for reading properties
        # Format: /iot/app/11/{user_id}/{product_id}/{device_id}/properties/read
        read_topic = f"/iot/app/11/{self._user_id}/{product_id}/{device_id}/properties/read"
        
        _LOGGER.debug("Reading properties from %s via topic: %s", device_id, read_topic)
        _LOGGER.debug("Request payload: %s", request_payload)
        
        # Create future to wait for response
        loop = asyncio.get_event_loop()
        response_future = loop.create_future()
        
        # Register temporary callback for this specific message ID
        def response_handler(device_id_str: str, properties: dict):
            """Handle response for this specific read request."""
            if device_id_str == device_id:
                if not response_future.done():
                    response_future.set_result(properties)
        
        # Add to callbacks
        self.register_device_callback(response_handler)
        
        try:
            # Publish read command
            await self.publish(read_topic, request_payload)
            
            # Wait for response with timeout
            properties = await asyncio.wait_for(response_future, timeout=10.0)
            _LOGGER.info("Successfully read properties from %s: %s", device_id, properties)
            return properties
            
        except asyncio.TimeoutError:
            raise HeimanMQTTError(f"Timeout reading properties from {device_id}")
        finally:
            # Remove callback
            self.unregister_device_callback(response_handler)
    
    async def publish(self, topic: str, payload: dict[str, Any]) -> None:
        """Publish message to MQTT topic.
        
        Args:
            topic: Topic to publish to
            payload: Message payload (will be JSON encoded)
            
        Raises:
            HeimanMqttError: If publish fails
        """
        if not self._connected or not self._client:
            raise HeimanMqttError("Not connected to MQTT broker")
        
        try:
            payload_str = json.dumps(payload)
            self._client.publish(topic, payload=payload_str, qos=1, retain=False)
            _LOGGER.debug("Published to %s: %s", topic, payload_str[:200])
        except Exception as err:
            raise HeimanMqttError(f"Failed to publish message: {str(err)}") from err
    
    async def async_write_property(
        self,
        device_id: str,
        product_id: str,
        property_identifiers: list[str],
        values: dict[str, Any],
        device_info: dict | None = None,
    ) -> bool:
        """Write device properties via MQTT.
        
        Args:
            device_id: Device ID to write to
            product_id: Product ID for topic construction
            property_identifiers: List of property identifiers to write
            values: Dictionary of property values to write
            device_info: Optional device info to determine if it's a child device
            
        Returns:
            True if successful, False otherwise
            
        Raises:
            HeimanMqttError: If write fails
        """
        if not self._connected or not self._client:
            raise HeimanMqttError("Not connected to MQTT broker")
        
        try:
            # Build topic based on device type (child or regular)
            topic = self._build_write_topic(product_id, device_id, device_info)
            
            # Generate unique message ID
            message_id = self._generate_message_id()
            
            # Convert boolean values to integers (0/1) for compatibility
            converted_values = {}
            for key, value in values.items():
                if isinstance(value, bool):
                    converted_values[key] = 1 if value else 0
                    _LOGGER.debug(
                        "Converting boolean value %s to integer %s for property %s",
                        value,
                        converted_values[key],
                        key,
                    )
                elif isinstance(value, str) and value.isdigit():
                    converted_values[key] = int(value)
                    _LOGGER.debug(
                        "Converting string number '%s' to integer %s for property %s",
                        value,
                        converted_values[key],
                        key,
                    )
                else:
                    converted_values[key] = value
            
            # Build properties with UserName if available
            properties_payload = converted_values.copy()
            if self._user_display_name:
                properties_payload["UserName"] = self._user_display_name
            
            # Build payload with correct format including messageType
            payload = {
                "messageId": message_id,
                "timestamp": int(time.time() * 1000),
                "deviceId": device_id,
                "properties": properties_payload,
                "messageType": "WRITE_PROPERTY",
            }
            
            # Create future for response
            loop = asyncio.get_event_loop()
            response_future = loop.create_future()
            self._pending_requests[message_id] = response_future
            
            _LOGGER.info(
                "Write property topic: %s payload: %s",
                topic,
                payload,
            )
            
            # Register temporary callback for this specific message ID
            def response_handler(device_id_str: str, properties: dict):
                """Handle response for this specific write request."""
                if device_id_str == device_id:
                    if not response_future.done():
                        response_future.set_result(True)
            
            # Add to callbacks
            self.register_device_callback(response_handler)
            
            try:
                # Publish write command
                await self.publish(topic, payload)
                
                # Wait for response with timeout
                await asyncio.wait_for(response_future, timeout=self._request_timeout)
                return True
                
            except asyncio.TimeoutError:
                _LOGGER.warning("Write property timeout for device %s", device_id)
                return False
            finally:
                # Remove callback
                self.unregister_device_callback(response_handler)
                
        except Exception as err:
            _LOGGER.error(
                "Failed to write properties to device %s: %s",
                device_id,
                err,
            )
            return False
