#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
MQTT相关功能模块
"""

import queue
import time
from typing import Optional, Dict, Any
from utils.logger import get_logger
import paho.mqtt.client as mqtt

logger = get_logger("services.online_script.built_in_modules.mqtt_module")



class VWEDMqttModule:
    """MQTT模块 - 提供MQTT订阅和发布功能"""
    
    def __init__(self, script_id: str):
        self.script_id = script_id
        self._clients: Dict[str, mqtt.Client] = {}
        self._message_queues: Dict[str, queue.Queue] = {}
        self._config = self._load_config()
        self._connected_clients: Dict[str, bool] = {}
    
    def _load_config(self) -> Dict[str, Any]:
        """加载MQTT配置"""
        return {
            'enable': True, 
            'pub_config': { # 发布者
                'broker': 'tcp://192.168.189.97:1883', # 服务器端点url 
                'topics': ['oagv/v2/asbm2_IRAYPLE/SIM-1/state'], # 订阅的主题，MQTT允许使用通配符订阅主题，不允许使用通配符pub发布消息
                'qos': 2, # 设置message的服务质量（0：消息最多传递一次（零次或一次）1：至少传递一次（一次或多次）。2：只传递一次）
                'client_id': 'RDS-Pub', # 客户端唯一标识
                'username': None, # 连接的用户名，不需要就填 null
                'password': None, # 连接的密码,不需要就填 null
                'clean_session': False, # 设置是否清空session,false表示服务器会保留客户端的连接记录，true每次都以新的身份连接服务器
                'connection_timeout': 30, # 超时时间(seconds)
                'keep_alive_interval': 60, # 设置会话心跳时间(seconds)
                'automatic_reconnect': True, # 设置断开后重新连接
                'retained': False,  # 表示发送的消息需要一直持久保存（不受服务器重启影响），不但要发送给当前的订阅者，并且以后新来的订阅了此Topic name的订阅者会马上得到推送。
                'will_msg': None, # 作为publish的遗嘱消息，会存到服务器，在publish端非正常断连的情况下，发送给所有订阅的客户端 不需要就填 null
                'will_topic': 'Examples1' # 遗嘱消息发布的topic
            },
            'sub_config': { # 订阅者
                'broker': 'tcp://192.168.189.97:1883', # 服务器端点url
                'topics': ['uagv/v2/IRAYPLE/SIM-1/state'], # MQTT允许使用通配符sub订阅主题，但不允许使用通配符pub发布消息
                'qos': 1, # 设置message的服务质量（0：消息最多传递一次（零次或一次）1：至少传递一次（一次或多次）。2：只传递一次）
                'client_id': 'RDS-Sub', # 客户端唯一标识
                'username': None, # 连接的用户名，不需要就填 null
                'password': None, # 连接的密码,不需要就填 null
                'clean_session': False, # 设置是否清空session,false表示服务器会保留客户端的连接记录，true每次都以新的身份连接服务器
                'connection_timeout': 30, # 超时时间(seconds)
                'keep_alive_interval': 60, # 设置会话心跳时间(seconds)
                'automatic_reconnect': True, # 设置断开后重新连接
                'retained': False # 表示发送的消息需要一直持久保存（不受服务器重启影响），不但要发送给当前的订阅者，并且以后新来的订阅了此Topic name的订阅者会马上得到推送。
            }
        }
    
    def _parse_broker_url(self, broker: str) -> tuple:
        """解析broker URL"""
        if broker.startswith('tcp://'):
            broker = broker[6:]
        elif broker.startswith('mqtt://'):
            broker = broker[7:]
        
        if ':' in broker:
            host, port = broker.split(':', 1)
            return host, int(port)
        else:
            return broker, 1883
    
    def _get_or_create_client(self, client_type: str, config: Dict[str, Any]) -> Optional[mqtt.Client]:
        """获取或创建MQTT客户端"""
        
        client_key = f"{client_type}_{self.script_id}"
        
        if client_key in self._clients:
            return self._clients[client_key]
        
        try:
            client = mqtt.Client(
                client_id=config['client_id'] + f"_{self.script_id}",
                callback_api_version=mqtt.CallbackAPIVersion.VERSION2
            )
            
            if config.get('username') and config.get('password'):
                client.username_pw_set(config['username'], config['password'])
            
            client.clean_session = config.get('clean_session', False)
            
            if client_type == 'sub':
                if client_key not in self._message_queues:
                    self._message_queues[client_key] = queue.Queue()
                
                def on_message(_client, _userdata, msg, _properties=None):
                    try:
                        message = msg.payload.decode('utf-8')
                        self._message_queues[client_key].put(message)
                    except Exception as e:
                        logger.error(f"Error processing MQTT message: {e}")
                
                client.on_message = on_message
            
            def on_connect(_client, _userdata, _flags, rc, _properties=None):
                if rc == 0:
                    self._connected_clients[client_key] = True
                    logger.info(f"MQTT client {client_key} connected successfully")
                else:
                    self._connected_clients[client_key] = False
                    logger.error(f"MQTT client {client_key} failed to connect with code {rc}")
            
            def on_disconnect(_client, _userdata, rc, _properties=None):
                self._connected_clients[client_key] = False
                logger.info(f"MQTT client {client_key} disconnected")
            
            client.on_connect = on_connect
            client.on_disconnect = on_disconnect
            
            host, port = self._parse_broker_url(config['broker'])
            client.connect(host, port, config.get('keep_alive_interval', 60))
            client.loop_start()
            
            self._clients[client_key] = client
            
            # 等待连接建立
            timeout = time.time() + config.get('connection_timeout', 30)
            while time.time() < timeout and not self._connected_clients.get(client_key, False):
                time.sleep(0.1)
            
            if not self._connected_clients.get(client_key, False):
                logger.error(f"MQTT client {client_key} connection timeout")
                return None
            
            return client
            
        except Exception as e:
            logger.error(f"Failed to create MQTT client: {e}")
            return None
    
    def mqtt_subscribe(self, topic: Optional[str] = None) -> Optional[str]:
        """MQTT订阅信息
        
        Args:
            topic: 订阅的topic，如果不提供则使用配置文件中的topic
            
        Returns:
            订阅到的信息，失败返回None
        """
        if not self._config.get('enable', False):
            logger.warning("MQTT is disabled in configuration")
            return None
        
        sub_config = self._config.get('sub_config', {})
        client = self._get_or_create_client('sub', sub_config)
        
        if not client:
            return None
        
        try:
            # 确定要订阅的topic
            if topic:
                subscribe_topic = topic
            else:
                topics = sub_config.get('topics', [])
                if not topics:
                    logger.error("No topics configured for subscription")
                    return None
                subscribe_topic = topics[0]  # 使用第一个配置的topic
            print("subscribe_topic::", subscribe_topic)
            # 订阅topic
            client.subscribe(subscribe_topic, sub_config.get('qos', 1))
            
            # 等待消息
            client_key = f"sub_{self.script_id}"
            message_queue = self._message_queues.get(client_key)
            
            if not message_queue:
                logger.error("Message queue not available")
                return None
            
            try:
                # 等待消息，超时时间5秒
                message = message_queue.get(timeout=5)
                return message
            except queue.Empty:
                logger.warning("No message received within timeout")
                return None
                
        except Exception as e:
            logger.error(f"MQTT subscribe failed: {e}")
            return None
    
    def mqtt_publish(self, topic_or_message: str, message: Optional[str] = None) -> bool:
        """MQTT发布信息
        
        Args:
            topic_or_message: 如果message为None，则此参数为消息内容；否则为topic
            message: 消息内容
            
        Returns:
            发布是否成功
        """
        if not self._config.get('enable', False):
            logger.warning("MQTT is disabled in configuration")
            return False
        
        pub_config = self._config.get('pub_config', {})
        client = self._get_or_create_client('pub', pub_config)
        
        if not client:
            return False
        
        try:
            # 确定topic和message
            if message is None:
                # 使用配置文件中的topic
                topics = pub_config.get('topics', [])
                if not topics:
                    logger.error("No topics configured for publishing")
                    return False
                publish_topic = topics[0]
                publish_message = topic_or_message
            else:
                # 使用传入的topic
                publish_topic = topic_or_message
                publish_message = message
            
            # 发布消息
            result = client.publish(
                publish_topic,
                publish_message,
                qos=pub_config.get('qos', 2),
                retain=pub_config.get('retained', False)
            )
            
            return result.rc == mqtt.MQTT_ERR_SUCCESS
            
        except Exception as e:
            logger.error(f"MQTT publish failed: {e}")
            return False
    
    def cleanup(self):
        """清理资源"""
        for client in self._clients.values():
            try:
                client.loop_stop()
                client.disconnect()
            except:
                pass
        
        self._clients.clear()
        self._message_queues.clear()
        self._connected_clients.clear()