#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：mqtt_manager.py
模块描述：MQTT 连接管理器，管理到 Broker 的连接池

作者：czj
创建日期：2026-07-29
最后修改：2026-07-29
"""

from __future__ import annotations

import json
import logging
import threading
import time
from datetime import datetime
from typing import Callable, Optional, Dict, Any

# paho-mqtt 为可选依赖：设备模拟（device_sim）模块未启用 / 未联调时，
# Broker 不配置，paho 可能未安装。此处降级为可选导入，避免阻塞整个应用启动。
# 仅在真正调用 MQTT 功能时（连接/发布）才会因 mqtt 为 None 抛错。
try:
    import paho.mqtt.client as mqtt
    PAHO_AVAILABLE: bool = True
except ImportError:  # pragma: no cover - 依赖未安装时的优雅降级
    mqtt = None  # type: ignore[assignment]
    PAHO_AVAILABLE = False

logger = logging.getLogger(__name__)


class MqttConnection:
    """
    单个 MQTT 连接

    封装 paho-mqtt 客户端，管理连接生命周期。

    Attributes:
        client (mqtt.Client): paho-mqtt 客户端实例
        env_config_id (str): 关联环境配置 ID
        connected (bool): 是否已连接
        reconnect_attempts (int): 重连尝试次数
        last_connected_at (datetime): 最后连接成功时间
        _lock (threading.Lock): 线程安全锁
    """

    def __init__(self, env_config_id: str):
        self.env_config_id = env_config_id
        self.client: Optional[mqtt.Client] = None
        self.connected = False
        self.reconnect_attempts = 0
        self.last_connected_at: Optional[datetime] = None
        self._lock = threading.Lock()
        self._subscribers: Dict[str, Callable] = {}  # topic -> callback

    def on_connect(self, client, userdata, flags, rc):
        """连接成功/失败回调"""
        with self._lock:
            if rc == 0:
                self.connected = True
                self.reconnect_attempts = 0
                self.last_connected_at = datetime.now()
                logger.info(f"MQTT 连接成功: env_config_id={self.env_config_id}")
                # 重新订阅已有主题
                for topic, callback in self._subscribers.items():
                    client.subscribe(topic)
                    logger.info(f"重连后重新订阅主题: {topic}")
            else:
                self.connected = False
                logger.error(f"MQTT 连接失败: env_config_id={self.env_config_id}, rc={rc}")

    def on_disconnect(self, client, userdata, rc):
        """断开连接回调"""
        with self._lock:
            self.connected = False
            if rc != 0:
                # 非主动断开，尝试重连
                self.reconnect_attempts += 1
                delay = min(2 ** self.reconnect_attempts, 60)  # 指数退避，最大 60 秒
                logger.warning(
                    f"MQTT 意外断开: env_config_id={self.env_config_id}, "
                    f"重连尝试={self.reconnect_attempts}, 等待={delay}s"
                )
                time.sleep(delay)
                try:
                    client.reconnect()
                except Exception as e:
                    logger.error(f"MQTT 重连失败: {e}")
            else:
                logger.info(f"MQTT 主动断开: env_config_id={self.env_config_id}")


class MqttManager:
    """
    MQTT 连接管理器（单例）

    管理多个环境配置的 MQTT 连接，提供连接/断开/发布/订阅接口。
    同一环境下的所有模拟设备共享一个 MQTT 连接。

    Attributes:
        _connections (dict): env_config_id -> MqttConnection
        _simulators (dict): device_id -> BaseSimulator 引用
        _lock (threading.Lock): 全局线程安全锁
    """

    _instance = None
    _lock = threading.Lock()

    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._connections = {}
                    cls._instance._simulators = {}
                    cls._instance._manager_lock = threading.Lock()
        return cls._instance

    def connect(self, env_config_id: str, host: str, port: int,
                username: Optional[str] = None, password: Optional[str] = None,
                client_id: Optional[str] = None, use_tls: bool = False) -> bool:
        """
        连接 MQTT Broker

        Args:
            env_config_id: 环境配置 ID
            host: Broker 地址
            port: 端口号
            username: 用户名
            password: 密码
            client_id: 客户端 ID
            use_tls: 是否启用 TLS

        Returns:
            bool: 是否连接成功
        """
        with self._manager_lock:
            # 如果已有连接，先断开
            if env_config_id in self._connections:
                self.disconnect(env_config_id)

            try:
                conn = MqttConnection(env_config_id)
                client = mqtt.Client(
                    client_id=client_id or f"sim_{env_config_id[:8]}",
                    protocol=mqtt.MQTTv311
                )

                # 设置回调
                client.on_connect = conn.on_connect
                client.on_disconnect = conn.on_disconnect

                # 设置认证
                if username:
                    client.username_pw_set(username, password)

                # 设置 TLS
                if use_tls:
                    client.tls_set()

                # 设置自动重连（paho-mqtt 内置机制）
                client.reconnect_delay_set(min_delay=1, max_delay=60)

                # 连接
                client.connect(host, port, keepalive=60)
                client.loop_start()  # 启动网络循环

                conn.client = client
                # 等待连接完成（最多 5 秒）
                for _ in range(50):
                    if conn.connected:
                        break
                    time.sleep(0.1)

                self._connections[env_config_id] = conn
                if conn.connected:
                    logger.info(f"MQTT 连接成功: {host}:{port}")
                    return True
                else:
                    logger.warning(f"MQTT 连接超时: {host}:{port}")
                    return False

            except Exception as e:
                logger.error(f"MQTT 连接异常: {host}:{port}, error={e}")
                return False

    def disconnect(self, env_config_id: str) -> None:
        """
        断开 MQTT 连接

        Args:
            env_config_id: 环境配置 ID
        """
        with self._manager_lock:
            conn = self._connections.pop(env_config_id, None)
            if conn and conn.client:
                try:
                    conn.client.loop_stop()
                    conn.client.disconnect()
                    logger.info(f"MQTT 断开连接: env_config_id={env_config_id}")
                except Exception as e:
                    logger.warning(f"MQTT 断开异常: {e}")

    def publish(self, env_config_id: str, topic: str, payload: dict) -> bool:
        """
        发布消息到 MQTT 主题

        Args:
            env_config_id: 环境配置 ID
            topic: MQTT 主题
            payload: 消息内容

        Returns:
            bool: 是否发布成功
        """
        conn = self._connections.get(env_config_id)
        if not conn or not conn.connected or not conn.client:
            logger.warning(f"MQTT 未连接，无法发布: env_config_id={env_config_id}")
            return False

        try:
            payload_str = json.dumps(payload, ensure_ascii=False)
            result = conn.client.publish(topic, payload_str, qos=1)
            if result.rc == mqtt.MQTT_ERR_SUCCESS:
                return True
            else:
                logger.warning(f"MQTT 发布失败: topic={topic}, rc={result.rc}")
                return False
        except Exception as e:
            logger.error(f"MQTT 发布异常: topic={topic}, error={e}")
            return False

    def subscribe(self, env_config_id: str, topic: str,
                  callback: Callable[[str, dict], None]) -> bool:
        """
        订阅 MQTT 主题

        Args:
            env_config_id: 环境配置 ID
            topic: MQTT 主题
            callback: 消息回调函数 (topic, payload) -> None

        Returns:
            bool: 是否订阅成功
        """
        conn = self._connections.get(env_config_id)
        if not conn or not conn.client:
            logger.warning(f"MQTT 未连接，无法订阅: env_config_id={env_config_id}")
            return False

        try:
            # 注册回调
            conn._subscribers[topic] = callback

            # 设置消息回调（使用包装函数）
            def on_message(client, userdata, msg):
                try:
                    payload = json.loads(msg.payload.decode())
                    cb = conn._subscribers.get(msg.topic)
                    if cb:
                        cb(msg.topic, payload)
                except Exception as e:
                    logger.error(f"MQTT 消息处理异常: topic={msg.topic}, error={e}")

            conn.client.on_message = on_message
            conn.client.subscribe(topic, qos=1)
            logger.info(f"MQTT 订阅主题: {topic}")
            return True
        except Exception as e:
            logger.error(f"MQTT 订阅异常: topic={topic}, error={e}")
            return False

    def unsubscribe(self, env_config_id: str, topic: str) -> bool:
        """
        取消订阅 MQTT 主题

        Args:
            env_config_id: 环境配置 ID
            topic: MQTT 主题

        Returns:
            bool: 是否取消成功
        """
        conn = self._connections.get(env_config_id)
        if not conn or not conn.client:
            return False

        try:
            conn.client.unsubscribe(topic)
            conn._subscribers.pop(topic, None)
            return True
        except Exception as e:
            logger.error(f"MQTT 取消订阅异常: topic={topic}, error={e}")
            return False

    def get_connection_status(self, env_config_id: str) -> str:
        """
        获取连接状态

        Args:
            env_config_id: 环境配置 ID

        Returns:
            str: connected/disconnected/error
        """
        conn = self._connections.get(env_config_id)
        if not conn:
            return "disconnected"
        if conn.connected:
            return "connected"
        return "error"

    def is_connected(self, env_config_id: str) -> bool:
        """
        检查是否已连接

        Args:
            env_config_id: 环境配置 ID

        Returns:
            bool: 是否已连接
        """
        conn = self._connections.get(env_config_id)
        return conn is not None and conn.connected

    def register_simulator(self, device_id: str, simulator) -> None:
        """
        注册模拟设备到管理器

        Args:
            device_id: 设备 ID
            simulator: 模拟器实例
        """
        self._simulators[device_id] = simulator

    def unregister_simulator(self, device_id: str) -> None:
        """
        注销模拟设备

        Args:
            device_id: 设备 ID
        """
        self._simulators.pop(device_id, None)

    def test_connection(self, host: str, port: int,
                        username: Optional[str] = None,
                        password: Optional[str] = None,
                        use_tls: bool = False,
                        timeout: int = 5) -> tuple[bool, str, Optional[int]]:
        """
        测试 MQTT Broker 连接

        Args:
            host: Broker 地址
            port: 端口号
            username: 用户名
            password: 密码
            use_tls: 是否启用 TLS
            timeout: 超时秒数

        Returns:
            tuple: (success, message, latency_ms)
        """
        test_client_id = f"test_{int(time.time())}"
        try:
            client = mqtt.Client(client_id=test_client_id, protocol=mqtt.MQTTv311)
            if username:
                client.username_pw_set(username, password)
            if use_tls:
                client.tls_set()

            start = time.time()
            client.connect(host, port, keepalive=timeout)
            client.loop_start()

            # 等待连接
            for _ in range(timeout * 10):
                if client.is_connected():
                    break
                time.sleep(0.1)

            elapsed = int((time.time() - start) * 1000)

            if client.is_connected():
                client.disconnect()
                client.loop_stop()
                return True, "连接成功", elapsed
            else:
                client.loop_stop()
                return False, "连接超时", None

        except Exception as e:
            return False, f"连接失败: {str(e)}", None


# 全局单例
mqtt_manager = MqttManager()