import WebSocket, { WebSocketServer } from 'ws';
import http from 'http';
import { DBModel } from "./models/index.js";
import { chatTask } from "./agent/task.js";

export default class WebSocketServerManager {
  constructor(port = 8080) {
    this.port = port;
    this.wss = null;
    this.server = null;
  }

  start() {
    this.server = http.createServer((req, res) => {
      res.writeHead(426, { 'Content-Type': 'text/plain' });
      res.end('Upgrade Required: WebSocket expected');
    });

    this.wss = new WebSocketServer({
      server: this.server,
      perMessageDeflate: {
        zlibDeflateOptions: {
          chunkSize: 1024,
          memLevel: 7,
          level: 3
        },
        zlibInflateOptions: {
          chunkSize: 10 * 1024
        },
        clientNoContextTakeover: true,
        serverNoContextTakeover: true,
        serverMaxWindowBits: 10,
        concurrencyLimit: 10,
        threshold: 1024
      }
    });

    this.wss.on('connection', (ws, req) => {
      const path = req.url;
      console.log(`WebSocket connected: ${req.socket.remoteAddress}, path: ${path}`);

      if (path === '/chat') {
        this._handleChatConnection(ws);
      } else {
        ws.send(JSON.stringify({ type: 'system', content: 'Unknown path' }));
        ws.close(1000, 'Unknown path');
      }
    });

    this.server.listen(this.port, () => {
      console.log(`WebSocket server started on ws://localhost:${this.port}`);
    });
  }

  _handleChatConnection(ws) {
    ws.send(JSON.stringify({ type: 'system', content: 'Connected to chat server' }));

    ws.on('message', (data) => this._handleMessage(ws, data));

    ws.on('close', () => {
      console.log('Chat client disconnected');
    });

    ws.on('error', (err) => {
      console.error('WebSocket error:', err);
    });
  }

  async _handleMessage(ws, data) {
    try {
      const raw = data.toString();
      console.log('Received:', raw);

      const msg = JSON.parse(raw);

      if (msg.type === 'ping') {
        ws.send(JSON.stringify({ type: 'pong' }));
        return;
      }

      if (msg.type === 'chat' || msg.type === 'clear') {
        const userInfo = await this.getUserInfo(msg.userId);
        chatTask.streamChat(userInfo, msg, (source, type, content, id) => {
          ws.send(JSON.stringify({ source, type, content, id }));
        });
        return;
      }

      if (msg.type === 'system') {
        console.log('[System]', msg.content);
        return;
      }

      console.log('[Unknown type]', msg);
    } catch (e) {
      console.error('Message parse error:', e);
      ws.send(JSON.stringify({ type: 'error', content: 'Invalid JSON' }));
    }
  }

  broadcast(message) {
    if (!this.wss) return;
    this.wss.clients.forEach((client) => {
      if (client.readyState === WebSocket.OPEN) {
        client.send(message);
      }
    });
  }

  stop() {
    if (this.wss) {
      this.wss.close();
      this.wss = null;
    }
    if (this.server) {
      this.server.close();
      this.server = null;
    }
    console.log('WebSocket server stopped');
  }

  getClientsCount() {
    if (!this.wss) return 0;
    return this.wss.clients.size;
  }

  async getUserInfo(userId) {
    if (!DBModel.User || userId.length === 0) return null;
    return await DBModel.User.findById(userId).lean().exec();
  }
}