const express = require('express');
const path = require('node:path');
const { WebSocketServer, WebSocket } = require('ws');
const http = require('node:http');
require('dotenv').config();

const app = express();
const server = http.createServer(app);

const ALLOWED_WEBSOCKET_ORIGIN = "https://www.aisada.ir"; // یا "http://www.aisada.ir"; // یا "https://aisada.ir"; // یا "http://aisada.ir"

const wss = new WebSocketServer({
  server,
  verifyClient: (info, cb) => {
    const origin = info.origin;
    // --- 👇 لاگ Verify را نگه می‌داریم چون مهم است 👇 ---
    console.log(`[WebSocket Verify] Verifying origin: ${origin}`);
    if (origin === ALLOWED_WEBSOCKET_ORIGIN) {
      console.log(`[WebSocket Verify] Origin allowed: ${origin}`);
      cb(true);
    } else {
      console.warn(`[WebSocket Verify] Forbidden origin: ${origin}. Rejecting connection.`);
      cb(false, 403, 'Forbidden Origin');
    }
  }
});

app.use(express.static(path.join(__dirname, '../build')));

const GEMINI_API_KEY = process.env.GEMINI_API_KEY;
if (!GEMINI_API_KEY) {
  console.error('FATAL ERROR: GEMINI_API_KEY environment variable is not set!');
  process.exit(1);
}

const createGeminiWebSocket = (clientWs) => {
   const geminiWsUrl = `wss://generativelanguage.googleapis.com/v1beta/models/gemini-pro:streamGenerateContent?key=${GEMINI_API_KEY}&alt=sse`;
   console.log(`[Gemini Connect] Attempting connection to Gemini...`);
   const geminiWs = new WebSocket(geminiWsUrl);

   geminiWs.on('open', () => {
     console.log('[Gemini Connect] Connected to Gemini API');
     if (geminiWs.pendingSetup) {
       // لاگ pendingSetup ممکن است حجیم باشد، خلاصه‌تر می‌کنیم
       console.log('[Gemini Connect] Sending pending setup...');
       geminiWs.send(JSON.stringify(geminiWs.pendingSetup));
       geminiWs.pendingSetup = null;
     }
   });

   geminiWs.on('message', (data) => {
     try {
       // --- 👇 لاگ دریافت از Gemini حذف شد 👇 ---
       // console.log('[Gemini Receive] Received from Gemini');
       if (clientWs.readyState === WebSocket.OPEN) {
            if (Buffer.isBuffer(data)) {
                 clientWs.send(data, { binary: true });
            } else {
                  const blob = Buffer.from(data.toString());
                  clientWs.send(blob, { binary: true });
            }
       }
     } catch (error) {
       console.error('[Gemini Receive] Error handling Gemini message:', error);
     }
   });

    geminiWs.on('error', (error) => console.error('[Gemini Connect] Gemini WebSocket error:', error));
    geminiWs.on('close', (code, reason) => console.log('[Gemini Connect] Gemini WebSocket closed:', code, reason?.toString()));

   return geminiWs;
};

wss.on('connection', (ws, req) => {
  const origin = req.headers.origin || 'N/A';
  console.log(`[WebSocket Event] Client connected (Origin: ${origin})`);
  let geminiWs = null;

  ws.on('message', async (message) => {
    try {
       const messageString = message.toString();
       const data = JSON.parse(messageString);
       // --- 👇 لاگ دریافت از Client حذف شد 👇 ---
       // console.log('[Client Receive] Received from client');

      if (data.setup) {
        // لاگ setup را نگه می‌داریم چون مهم است
        console.log('[Client Receive] Initializing Gemini connection...');
        geminiWs = createGeminiWebSocket(ws);
        if (geminiWs.readyState !== WebSocket.OPEN) {
          geminiWs.pendingSetup = data;
        } else {
          geminiWs.send(JSON.stringify(data));
        }
        return;
      }

      if (geminiWs && geminiWs.readyState === WebSocket.OPEN) {
        geminiWs.send(JSON.stringify(data));
      } else if (geminiWs) {
        // لاگ انتظار را نگه می‌داریم
        console.log('[Client Receive] Gemini connection not ready, waiting...');
      } else {
        console.error('[Client Receive] No Gemini connection established. Closing client connection.');
        ws.close(1011, 'Backend connection not ready');
      }
    } catch (error) {
      console.error('[Client Receive] Error processing message:', error, message.toString().substring(0, 100) + '...');
      ws.close(1008, 'Message processing error');
    }
  });

  ws.on('close', () => {
    console.log('[WebSocket Event] Client disconnected');
    if (geminiWs && geminiWs.readyState !== WebSocket.CLOSED && geminiWs.readyState !== WebSocket.CLOSING) {
      console.log('[WebSocket Event] Closing associated Gemini connection.');
      geminiWs.close();
    }
    geminiWs = null;
  });

  ws.on('error', (error) => {
     console.error('[WebSocket Event] Client WebSocket error:', error);
     if (geminiWs && geminiWs.readyState !== WebSocket.CLOSED && geminiWs.readyState !== WebSocket.CLOSING) {
        console.log('[WebSocket Event] Closing associated Gemini connection due to client error.');
        geminiWs.close();
     }
     geminiWs = null;
     if (ws.readyState !== WebSocket.CLOSED && ws.readyState !== WebSocket.CLOSING) {
         ws.close(1011, 'Client error');
     }
  });

});

app.get('*', (req, res) => {
  res.sendFile(path.resolve(__dirname, '../build', 'index.html'));
});

const PORT = process.env.PORT || 3001;
server.listen(PORT, () => {
  console.log(`Server listening on port ${PORT}`);
  console.log(`Allowed WebSocket Origin: ${ALLOWED_WEBSOCKET_ORIGIN}`);
});