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

"""
MQTT客户端服务模块
提供MQTT连接、发布、订阅等功能
支持自动重连和消息处理
"""

import asyncio
import json
import time
import threading
from typing import Dict, Any, Optional, Callable, List
from utils.logger import get_logger
import paho.mqtt.client as mqtt
from config.tf_api_config import MQTT_CONFIG

logger = get_logger("services.mqtt_service")


class MQTTService:
    """内部MQTT客户端服务"""
    
    def __init__(self, config: Dict[str, Any] = None):
        self.config = config or MQTT_CONFIG
        self.client = None
        self.connected = False
        self.reconnect_attempts = 0
        self.message_handlers: Dict[str, List[Callable]] = {}
        self._lock = threading.RLock()
        self._running = False
        
    def on_connect(self, client, userdata, flags, rc):
        """连接回调"""
        if rc == 0:
            self.connected = True
            self.reconnect_attempts = 0
            logger.info(f"MQTT连接成功: {self.config['host']}:{self.config['port']}")
        else:
            self.connected = False
            logger.error(f"MQTT连接失败，错误代码: {rc}")
    
    def on_disconnect(self, client, userdata, rc):
        """断开连接回调"""
        self.connected = False
        logger.warning(f"MQTT连接断开，代码: {rc}")
        
        if self._running and rc != 0:
            self._schedule_reconnect()
    
    def on_message(self, client, userdata, msg):
        """消息接收回调"""
        try:
            topic = msg.topic
            payload = msg.payload.decode('utf-8')
            
            with self._lock:
                handlers = self.message_handlers.get(topic, [])
                for handler in handlers:
                    try:
                        if asyncio.iscoroutinefunction(handler):
                            asyncio.create_task(handler(topic, payload))
                        else:
                            handler(topic, payload)
                    except Exception as e:
                        logger.error(f"MQTT消息处理器执行失败: {e}", exc_info=True)
                        
        except Exception as e:
            logger.error(f"处理MQTT消息失败: {e}", exc_info=True)
    
    def _schedule_reconnect(self):
        """安排重连"""
        if self.reconnect_attempts >= self.config.get('max_retries', 3):
            logger.error("MQTT重连次数已达上限，停止重连")
            return
            
        delay = self.config.get('reconnect_delay', 5) * (self.reconnect_attempts + 1)
        self.reconnect_attempts += 1
        logger.info(f"将在 {delay} 秒后重连MQTT（第{self.reconnect_attempts}次尝试）")
        
        def reconnect():
            time.sleep(delay)
            if self._running:
                self.connect()
        
        threading.Thread(target=reconnect, daemon=True).start()
    
    def connect(self):
        """连接MQTT服务器"""
        try:
            if self.client:
                self.client.disconnect()
            
            self.client = mqtt.Client(callback_api_version=mqtt.CallbackAPIVersion.VERSION2)
            self.client.on_connect = self.on_connect
            self.client.on_message = self.on_message
            self.client.on_disconnect = self.on_disconnect
            
            # 设置认证信息
            if self.config.get('username'):
                self.client.username_pw_set(
                    self.config['username'], 
                    self.config.get('password', '')
                )
            
            logger.info(f"正在连接MQTT服务器 {self.config['host']}:{self.config['port']}...")
            self.client.connect(
                self.config['host'], 
                self.config['port'], 
                self.config.get('keepalive', 60)
            )
            
            self.client.loop_start()
            self._running = True
            
            # 等待连接建立
            timeout = 10
            start_time = time.time()
            while not self.connected and (time.time() - start_time) < timeout:
                time.sleep(0.1)
            
            if self.connected:
                logger.info("MQTT连接建立成功")
                return True
            else:
                logger.error("MQTT连接建立超时")
                return False
                
        except Exception as e:
            logger.error(f"连接MQTT服务器失败: {e}")
            return False
    
    def disconnect(self):
        """断开连接"""
        self._running = False
        if self.client:
            self.client.loop_stop()
            self.client.disconnect()
        self.connected = False
        logger.info("MQTT连接已断开")
    
    def subscribe(self, topic: str, handler: Callable = None):
        """订阅主题"""
        try:
            if not self.client or not self.connected:
                logger.warning(f"MQTT未连接，无法订阅主题: {topic}")
                return False
            result = self.client.subscribe(topic)
            if result[0] == mqtt.MQTT_ERR_SUCCESS:
                logger.info(f"成功订阅MQTT主题: {topic}")
                
                # 添加消息处理器
                if handler:
                    with self._lock:
                        if topic not in self.message_handlers:
                            self.message_handlers[topic] = []
                        if handler not in self.message_handlers[topic]:
                            self.message_handlers[topic].append(handler)
                
                return True
            else:
                logger.error(f"订阅MQTT主题失败: {topic}, 错误码: {result[0]}")
                return False
        except Exception as e:
            logger.error(f"订阅MQTT主题异常: {e}")
            return False
    
    def unsubscribe(self, topic: str):
        """取消订阅主题"""
        try:
            if self.client and self.connected:
                result = self.client.unsubscribe(topic)
                if result[0] == mqtt.MQTT_ERR_SUCCESS:
                    logger.info(f"取消订阅MQTT主题: {topic}")
                    
                    # 清理消息处理器
                    with self._lock:
                        if topic in self.message_handlers:
                            del self.message_handlers[topic]
                    
                    return True
                else:
                    logger.error(f"取消订阅MQTT主题失败: {topic}, 错误码: {result[0]}")
                    return False
        except Exception as e:
            logger.error(f"取消订阅MQTT主题异常: {e}")
            return False
    
    def publish(self, topic: str, payload: Any, qos: int = 0, retain: bool = False):
        """发布消息"""
        try:
            if not self.client or not self.connected:
                logger.warning(f"MQTT未连接，无法发布消息到: {topic}")
                return False
            
            if isinstance(payload, dict):
                payload_str = json.dumps(payload, ensure_ascii=False)
            else:
                payload_str = str(payload)
            
            result = self.client.publish(topic, payload_str, qos, retain)
            if result.rc == mqtt.MQTT_ERR_SUCCESS:
                logger.debug(f"MQTT消息发布成功: {topic}")
                return True
            else:
                logger.error(f"MQTT消息发布失败: {topic}, 错误码: {result.rc}")
                return False
        except Exception as e:
            logger.error(f"发布MQTT消息异常: {e}")
            return False
    
    def add_message_handler(self, topic: str, handler: Callable):
        """为指定主题添加消息处理器"""
        with self._lock:
            if topic not in self.message_handlers:
                self.message_handlers[topic] = []
            if handler not in self.message_handlers[topic]:
                self.message_handlers[topic].append(handler)
    
    def remove_message_handler(self, topic: str, handler: Callable):
        """移除指定主题的消息处理器"""
        with self._lock:
            if topic in self.message_handlers:
                if handler in self.message_handlers[topic]:
                    self.message_handlers[topic].remove(handler)
                if not self.message_handlers[topic]:
                    del self.message_handlers[topic]
    
    def is_connected(self) -> bool:
        """检查连接状态"""
        return self.connected
    
    def get_connection_info(self) -> Dict[str, Any]:
        """获取连接信息"""
        return {
            "host": self.config['host'],
            "port": self.config['port'],
            "connected": self.connected,
            "reconnect_attempts": self.reconnect_attempts,
            "running": self._running,
            "subscribed_topics": list(self.message_handlers.keys())
        }


def create_mqtt_service(config: Dict[str, Any] = None) -> MQTTService:
    """创建MQTT服务实例"""
    return MQTTService(config)