frost-mqtt-client/app/src/main/mqtt-manager.js

import { app } from 'electron';
import { readFileSync, writeFileSync, existsSync, mkdirSync } from 'fs';
import { join } from 'path';

import mqtt from 'mqtt';

class MqttManager {

  constructor() {
    this.clients = new Map(); // serverId -> mqtt client
    this.servers = new Map(); // serverId -> server config
    this.messages = new Map(); // serverId -> message[]
    this.publishHistory = new Map(); // serverId -> publishHistory[]
    this.listeners = new Set(); // IPC callback functions
    this.msgIdCounter = 0; // 用于生成唯一消息 id
    this.dataPath = join(app.getPath('userData'), 'frost-mqtt-data.json');
    this.loadData();
  }

  nextMsgId() {
    this.msgIdCounter = (this.msgIdCounter + 1) % Number.MAX_SAFE_INTEGER;
    return 'msg_' + Date.now() + '_' + this.msgIdCounter;
  }

  // ==================== 数据持久化 ====================

  getDataDir() {

    let dir = app.getPath('userData');

    if (!existsSync(dir)) {
      mkdirSync(dir, { recursive: true });
    }

    return dir;

  }

  loadData() {
    try {
      if (existsSync(this.dataPath)) {

        let raw = readFileSync(this.dataPath, 'utf-8');
        let data = JSON.parse(raw);

        if (data.servers) {
          data.servers.forEach((s) => {
            s.status = 'disconnected';
            s.topics = (s.topics || []).map((t) => ({
              ...t,
              id: t.id || 'topic_' + Date.now() + '_' + Math.random().toString(36).slice(2, 6),
              subscribed: false,
            }));
            this.servers.set(s.id, s);
          });
        }

        if (data.messages) {
          Object.entries(data.messages).forEach(([id, msgs]) => {
            this.messages.set(id, msgs);
          });
        }

        if (data.publishHistory) {
          Object.entries(data.publishHistory).forEach(([id, history]) => {
            this.publishHistory.set(id, history);
          });
        }

      } else {
        this.initDefaultServers();
      }
    } catch (e) {
      console.error('加载数据失败:', e.message);
      this.initDefaultServers();
    }
  }

  saveData() {
    try {

      let data = {
        servers: Array.from(this.servers.values()).map((s) => ({
          ...s,
          status: 'disconnected',
          topics: s.topics.map((t) => ({ ...t, subscribed: false })),
        })),
        messages: Object.fromEntries(this.messages),
        publishHistory: Object.fromEntries(this.publishHistory),
      };

      writeFileSync(this.dataPath, JSON.stringify(data, null, 2), 'utf-8');

    } catch (e) {
      console.error('保存数据失败:', e.message);
    }
  }

  initDefaultServers() {

    let defaults = [
      {
        id: 'example_server',
        name: 'EMQX 公共测试',
        host: 'broker.emqx.io',
        port: 1883,
        protocol: 'mqtt',
        clientId: 'frost_client_' + Math.random().toString(36).slice(2, 8),
        username: '',
        password: '',
        keepAlive: 60,
        cleanSession: true,
        connectTimeout: 10000,
        reconnect: true,
        reconnectInterval: 5000,
        status: 'disconnected',
        topics: [
          { id: 'topic_default_001', topic: 'sensor/+/temperature', qos: 1, subscribed: false },
          { id: 'topic_default_002', topic: 'device/status', qos: 2, subscribed: false },
        ],
      },
    ];

    defaults.forEach((s) => {
      this.servers.set(s.id, s);
      this.messages.set(s.id, []);
      this.publishHistory.set(s.id, []);
    });

    this.saveData();

  }

  // ==================== 监听器 ====================

  addListener(callback) {
    this.listeners.add(callback);
  }

  removeListener(callback) {
    this.listeners.delete(callback);
  }

  notify(event, data) {
    this.listeners.forEach((cb) => {
      try {
        cb(event, data);
      } catch (e) {
        console.error('通知失败:', e);
      }
    });
  }

  // ==================== 服务器管理 ====================

  getServers() {
    return Array.from(this.servers.values()).map((s) => ({
      ...s,
      password: s.password || '',
    }));
  }

  getServer(id) {
    return this.servers.get(id) || null;
  }

  getServerConfigSchema() {
    return {
      name: (c, e) => c.name ?? e?.name,
      host: (c, e) => c.host ?? e?.host,
      port: (c, e) => c.port || e?.port || 1883,
      protocol: (c, e) => c.protocol || e?.protocol || 'mqtt',
      clientId: (c, e) => c.clientId || e?.clientId || 'frost_' + Math.random().toString(36).slice(2, 10),
      username: (c, e) => c.username ?? e?.username ?? '',
      password: (c, e, opts) => c.password ?? (opts.keepExistingPassword ? e?.password : '') ?? '',
      keepAlive: (c, e) => c.keepAlive || e?.keepAlive || 60,
      cleanSession: (c) => c.cleanSession !== false,
      connectTimeout: (c, e) => c.connectTimeout ?? e?.connectTimeout ?? 10000,
      reconnect: (c) => c.reconnect !== false,
      reconnectInterval: (c, e) => c.reconnectInterval ?? e?.reconnectInterval ?? 5000,
    };
  }

  buildServerFields(config, existing = null, options = {}) {
    let schema = this.getServerConfigSchema();
    let fields = {};
    for (let key of Object.keys(schema)) {
      fields[key] = schema[key](config, existing, options);
    }
    return fields;
  }

  buildTopic(t) {
    return {
      id: t.id || 'topic_' + Date.now() + '_' + Math.random().toString(36).slice(2, 6),
      topic: t.topic,
      qos: t.qos ?? 1,
      subscribed: false,
    };
  }

  addServer(config) {

    let id = 'srv_' + Date.now() + '_' + Math.random().toString(36).slice(2, 6);
    let server = {
      id,
      ...this.buildServerFields(config),
      status: 'disconnected',
      topics: [],
    };

    this.servers.set(id, server);
    this.messages.set(id, []);
    this.publishHistory.set(id, []);
    this.saveData();

    return server;

  }

  updateServer(id, config) {

    let server = this.servers.get(id);

    if (!server) {
      return null;
    }

    Object.assign(server, this.buildServerFields(config, server, { keepExistingPassword: true }));

    this.saveData();

    return server;

  }

  deleteServer(id) {
    // 先断开连接
    this.disconnect(id);
    this.servers.delete(id);
    this.messages.delete(id);
    this.publishHistory.delete(id);
    this.saveData();
  }

  exportServers() {
    let schemaKeys = Object.keys(this.getServerConfigSchema());
    return Array.from(this.servers.values()).map((s) => {
      let exported = {};
      for (let key of schemaKeys) {
        exported[key] = s[key];
      }
      exported.topics = (s.topics || []).map((t) => this.buildTopicExport(t));
      return exported;
    });
  }

  buildTopicExport(t) {
    return {
      id: t.id,
      topic: t.topic,
      qos: t.qos,
    };
  }

  importServers(serverList) {
    if (!Array.isArray(serverList)) {
      throw new Error('导入格式错误：必须是服务器配置数组');
    }

    // 断开所有现有连接
    this.clients.forEach((client, id) => {
      this.disconnect(id);
    });
    this.clients.clear();

    this.servers.clear();
    this.messages.clear();
    this.publishHistory.clear();

    serverList.forEach((config) => {
      if (!config.name || !config.host) {
        return;
      }

      let id = 'srv_' + Date.now() + '_' + Math.random().toString(36).slice(2, 6);
      let server = {
        id,
        ...this.buildServerFields(config),
        status: 'disconnected',
        topics: (config.topics || []).map((t) => this.buildTopic(t)),
      };

      this.servers.set(id, server);
      this.messages.set(id, []);
      this.publishHistory.set(id, []);
    });

    this.saveData();
    return this.getServers();
  }

  // ==================== 连接管理 ====================

  connect(id) {

    let server = this.servers.get(id);

    if (!server) {
      return false;
    }

    // 如果已连接，先断开
    if (this.clients.has(id)) {
      this.disconnect(id);
    }

    server.status = 'connecting';

    this.notify('status-change', { id, status: 'connecting' });

    let protocol = server.protocol || 'mqtt';
    let url = `${protocol}://${server.host}:${server.port}`;

    let options = {
      clientId: server.clientId,
      username: server.username || undefined,
      password: server.password || undefined,
      keepalive: server.keepAlive,
      clean: server.cleanSession,
      reconnectPeriod: 0,
      connectTimeout: server.connectTimeout,
    };

    try {

      let client = mqtt.connect(url, options);

      client.on('connect', () => {
        server.status = 'connected';
        this.notify('status-change', { id, status: 'connected' });
        // 只在首次连接成功后，再根据配置启用自动重连
        client.options.reconnectPeriod = server.reconnect ? server.reconnectInterval : 0;
        // 自动订阅已有主题
        server.topics.forEach((t) => {
          this._doSubscribe(id, t.topic, { qos: t.qos });
        });
      });

      client.on('error', (err) => {
        console.error(`[${server.name}] 连接错误:`, err.message);
        server.status = 'error';
        this.notify('status-change', { id, status: 'error', error: err.message });
      });

      client.on('close', () => {
        if (server.status === 'connected') {
          server.status = 'disconnected';
          server.topics.forEach((t) => {
            t.subscribed = false;
          });
          this.notify('status-change', { id, status: 'disconnected' });
        }
      });

      client.on('message', (topic, payload, packet) => {

        let msg = {
          time: new Date().toLocaleTimeString('zh-CN', { hour12: false }) + '.' + String(Date.now() % 1000).padStart(3, '0'),
          topic: topic,
          payload: payload.toString(),
          qos: packet.qos,
          retain: packet.retain,
          direction: 'sub',
          id: this.nextMsgId(),
        };

        let msgs = this.messages.get(id) || [];

        msgs.unshift(msg);

        if (msgs.length > 500) {
          msgs.pop();
        }

        this.messages.set(id, msgs);
        this.notify('message', { serverId: id, message: msg });

      });

      this.clients.set(id, client);

      return true;

    } catch (err) {
      server.status = 'error';
      this.notify('status-change', { id, status: 'error', error: err.message });
      return false;
    }
  }

  disconnect(id) {

    let client = this.clients.get(id);

    if (client) {
      client.end(true);
      this.clients.delete(id);
    }

    let server = this.servers.get(id);

    if (server) {
      server.status = 'disconnected';
      server.topics.forEach((t) => {
        t.subscribed = false;
      });
      this.notify('status-change', { id, status: 'disconnected' });
    }

  }

  // ==================== 订阅管理 ====================
  _doSubscribe(serverId, topic, opts) {

    let client = this.clients.get(serverId);
    let server = this.servers.get(serverId);

    if (!client || !server || server.status !== 'connected') {
      return false;
    }

    client.subscribe(topic, opts, (err) => {
      if (err) {
        console.error(`订阅失败 [${topic}]:`, err.message);
        this.notify('subscription-error', { serverId, topic, error: err.message });
      } else {
        let t = server.topics.find((t) => t.topic === topic);
        if (t) {
          t.subscribed = true;
        }
        this.notify('subscription-change', { serverId, topic, subscribed: true });
      }
    });

    return true;

  }

  addTopic(serverId, topic, qos = 1) {
    let server = this.servers.get(serverId);
    if (!server) {
      return null;
    }
    if (server.topics.find((t) => t.topic === topic)) {
      return null;
    }

    let topicId = 'topic_' + Date.now() + '_' + Math.random().toString(36).slice(2, 6);

    let newTopic = { id: topicId, topic, qos, subscribed: false };
    server.topics.push(newTopic);

    this.saveData();

    // 如果已连接，自动订阅
    if (server.status === 'connected') {
      this._doSubscribe(serverId, topic, { qos });
    }

    return newTopic;

  }

  removeTopic(serverId, topic) {

    let server = this.servers.get(serverId);

    if (!server) {
      return false;
    }

    // 先取消订阅
    this.unsubscribe(serverId, topic);

    server.topics = server.topics.filter((t) => t.topic !== topic);

    this.saveData();

    return true;

  }

  subscribe(serverId, topic) {

    let server = this.servers.get(serverId);

    if (!server) {
      return false;
    }

    let t = server.topics.find((t) => t.topic === topic);

    if (!t) {
      return false;
    }

    return this._doSubscribe(serverId, topic, { qos: t.qos });

  }

  unsubscribe(serverId, topic) {

    let client = this.clients.get(serverId);
    let server = this.servers.get(serverId);

    if (!client || !server || server.status !== 'connected') {
      return false;
    }

    client.unsubscribe(topic, {}, (err) => {
      if (!err) {
        let t = server.topics.find((t) => t.topic === topic);
        if (t) {
          t.subscribed = false;
        }
        this.notify('subscription-change', { serverId, topic, subscribed: false });
      }
    });

    return true;

  }

  updateTopic(serverId, topicId, newTopic, newQos) {

    let server = this.servers.get(serverId);

    if (!server) {
      return false;
    }

    let t = server.topics.find((item) => item.id === topicId);

    if (!t) {
      return false;
    }

    let topicChanged = t.topic !== newTopic;

    if (
      topicChanged &&
      server.topics.find((item) => item.topic === newTopic && item.id !== topicId)
    ) {
      return false;
    }

    let oldTopic = t.topic;
    let wasSubscribed = t.subscribed;

    // 如果主题名变更且已订阅，先取消旧主题订阅
    if (topicChanged && wasSubscribed) {
      this.unsubscribe(serverId, oldTopic);
    }

    t.topic = newTopic;
    t.qos = newQos;

    // 如果已连接且此前已订阅，重新订阅新主题
    if (server.status === 'connected' && wasSubscribed) {
      this._doSubscribe(serverId, newTopic, { qos: newQos });
    }

    this.saveData();
    this.notify('subscription-change', {
      serverId,
      topic: newTopic,
      subscribed: t.subscribed,
      topicId,
    });

    return true;

  }

  subscribeAll(serverId) {

    let client = this.clients.get(serverId);
    let server = this.servers.get(serverId);

    if (!client || !server || server.status !== 'connected') {
      return false;
    }

    let topicsToSubscribe = server.topics.filter((t) => !t.subscribed);

    if (topicsToSubscribe.length === 0) {
      return true;
    }

    let subMap = Object.fromEntries(topicsToSubscribe.map((t) => [t.topic, { qos: t.qos }]));

    client.subscribe(subMap, (err) => {
      if (err) {
        console.error('批量订阅失败:', err.message);
        topicsToSubscribe.forEach((t) => {
          this.notify('subscription-error', { serverId, topic: t.topic, error: err.message });
        });
      } else {
        topicsToSubscribe.forEach((t) => {
          t.subscribed = true;
          this.notify('subscription-change', {
            serverId,
            topic: t.topic,
            subscribed: true,
            topicId: t.id,
          });
        });
      }
    });

    return true;

  }

  // ==================== 消息发布 ====================

  publish(serverId, topic, payload, opts = {}) {

    let client = this.clients.get(serverId);
    let server = this.servers.get(serverId);

    if (!client || !server || server.status !== 'connected') {
      return false;
    }

    let pubOpts = {
      qos: opts.qos ?? 1,
      retain: opts.retain ?? false,
    };

    client.publish(topic, payload, pubOpts, (err) => {
      if (err) {
        console.error(`发布失败 [${topic}]:`, err.message);
        this.notify('publish-error', { serverId, topic, error: err.message });
      }
    });

    // 记录发布历史
    let pubEntry = {
      time:
        new Date().toLocaleTimeString('zh-CN', { hour12: false }) +
        '.' +
        String(Date.now() % 1000).padStart(3, '0'),
      topic,
      payload,
      qos: pubOpts.qos,
      retain: pubOpts.retain,
      direction: 'pub',
    };

    let history = this.publishHistory.get(serverId) || [];

    history.unshift(pubEntry);

    if (history.length > 50) {
      history.pop();
    }

    this.publishHistory.set(serverId, history);

    // 同时添加到消息日志
    let pubMsgId = this.nextMsgId();
    let msgs = this.messages.get(serverId) || [];

    msgs.unshift({
      ...pubEntry,
      id: pubMsgId,
    });

    if (msgs.length > 500) {
      msgs.pop();
    }

    this.messages.set(serverId, msgs);

    this.notify('message', {
      serverId,
      message: {
        ...pubEntry,
        id: pubMsgId,
      },
    });

    this.saveData();

    return true;

  }

  // ==================== 消息管理 ====================

  getMessages(serverId) {
    return this.messages.get(serverId) || [];
  }

  getPublishHistory(serverId) {
    return this.publishHistory.get(serverId) || [];
  }

  clearMessages(serverId) {
    this.messages.set(serverId, []);
    this.saveData();
  }

  clearPublishHistory(serverId) {
    this.publishHistory.set(serverId, []);
    this.saveData();
  }

}

// 单例
const mqttManager = new MqttManager();
export default mqttManager;