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 wss = new WebSocketServer({ server });

// --- START: منطق جدید برای چرخش API Key ---

// ۱. خواندن تمام کلیدهای API از Secrets (متغیرهای محیطی در هاگینگ فیس)
const apiKeys = [];
let i = 1;
while (process.env[`GEMINI_API_KEY_${i}`]) {
  apiKeys.push(process.env[`GEMINI_API_KEY_${i}`]);
  i++;
}

const numKeys = apiKeys.length;

if (numKeys === 0) {
  // اگر هیچ کلیدی پیدا نشد، خطا داده و خارج شو
  console.error(
    'خطای حیاتی: هیچ Secret با نام GEMINI_API_KEY_n یافت نشد!' +
      ' لطفاً Secret ها را مانند GEMINI_API_KEY_1, GEMINI_API_KEY_2, ... ' +
      'در تنظیمات Space خود اضافه کنید.'
  );
  process.exit(1); // برنامه را متوقف کن
} else {
  console.log(`تعداد ${numKeys} کلید API جیمینای بارگذاری شد.`);
}

// ۲. شمارنده سراسری برای انتخاب کلید بعدی (شروع از اندیس ۰)
let currentKeyIndex = 0;

// ۳. تابع برای گرفتن کلید بعدی به صورت چرخشی
function getNextApiKey() {
  if (numKeys === 0) {
    console.error('تلاش برای گرفتن کلید API در حالی که هیچ کلیدی بارگذاری نشده است.');
    return null; // یا یک خطا پرتاب کنید
  }
  // محاسبه اندیس کلید با استفاده از باقیمانده تقسیم
  const keyIndexToUse = currentKeyIndex % numKeys;
  const selectedKey = apiKeys[keyIndexToUse];
  console.log(`اختصاص کلید API با اندیس: ${keyIndexToUse}`); // لاگ برای اشکال‌زدایی

  // افزایش شمارنده برای درخواست بعدی
  currentKeyIndex++;

  // اختیاری: جلوگیری از خیلی بزرگ شدن شمارنده (گرچه جاوااسکریپت اعداد بزرگ را مدیریت می‌کند)
  // if (currentKeyIndex >= Number.MAX_SAFE_INTEGER - 10) { // نزدیک به حداکثر عدد امن
  //     currentKeyIndex = currentKeyIndex % numKeys;
  // }

  return selectedKey;
}

// --- END: منطق جدید برای چرخش API Key ---

// سرو کردن فایل‌های استاتیک از پوشه بیلد React
app.use(express.static(path.join(__dirname, '../build')));

// این تابع حالا کلید API انتخاب شده را به عنوان ورودی می‌گیرد
const createGeminiWebSocket = (clientWs, apiKey) => {
  // بررسی اینکه آیا کلید معتبری داده شده است
  if (!apiKey) {
    console.error('امکان ایجاد WebSocket جیمینای وجود ندارد: کلید API ارائه نشده است.');
    // می‌توانید یک پیام خطا به کلاینت بفرستید و اتصالش را ببندید
    if (clientWs.readyState === WebSocket.OPEN) {
        try {
             clientWs.send(JSON.stringify({ error: "خطای داخلی سرور: عدم امکان دریافت کلید API." }));
        } catch (sendError) {
             console.error("خطا در ارسال پیام خطا به کلاینت:", sendError);
        }
        clientWs.close();
    }
    return null; // برگرداندن null نشان‌دهنده شکست است
  }

  console.log(`ایجاد اتصال به جیمینای با کلید: ...${apiKey.slice(-4)}`); // نمایش ۴ کاراکتر آخر کلید برای تایید

  // استفاده از کلید API اختصاص‌یافته به این کاربر
  const geminiWs = new WebSocket(
    `wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1alpha.GenerativeService.BidiGenerateContent?key=${apiKey}`
  );

  // --- بقیه event handler های geminiWs (open, message, error, close) بدون تغییر باقی می‌مانند ---
  geminiWs.on('open', () => {
    console.log(`اتصال به جیمینای برای کلاینت با کلید ...${apiKey.slice(-4)} برقرار شد.`);
    // ارسال پیام setup معلق اگر وجود دارد
    if (geminiWs.pendingSetup) {
      console.log('ارسال پیام setup معلق:', geminiWs.pendingSetup);
      try {
        geminiWs.send(JSON.stringify(geminiWs.pendingSetup));
      } catch (sendError) {
          console.error("خطا در ارسال پیام setup معلق به جیمینای:", sendError);
          clientWs.close(); // بستن اتصال کلاینت اگر نتوانستیم با جیمینای ارتباط اولیه برقرار کنیم
      }
      geminiWs.pendingSetup = null; // پاک کردن پیام معلق
    }
  });

  geminiWs.on('message', (data) => {
    try {
      // تبدیل پیام به رشته قبل از ارسال به کلاینت (کلاینت شما انتظار Blob داشت، مطمئن شوید هنوز همینطور است)
      const message = data.toString();
      console.log(`دریافت از جیمینای (کلید ...${apiKey.slice(-4)}):`, message); // نمایش ۴ کاراکتر آخر کلید در لاگ

      // کلاینت شما انتظار Blob داشت. اگر هنوز نیاز به Blob دارید، آن را حفظ کنید:
      const blob = Buffer.from(message); // ایجاد Buffer (شبیه Blob در Node.js)
      if (clientWs.readyState === WebSocket.OPEN) {
          clientWs.send(blob, { binary: true }); // ارسال به عنوان داده باینری
      }

      // اگر کلاینت شما حالا انتظار JSON یا رشته دارد، کد بالا را تغییر دهید:
      // if (clientWs.readyState === WebSocket.OPEN) {
      //     clientWs.send(message); // ارسال به عنوان رشته
      // }

    } catch (error) {
      console.error(`خطا در پردازش پیام جیمینای (کلید ...${apiKey.slice(-4)}):`, error);
    }
  });

  geminiWs.on('error', (error) => {
    console.error(`خطای WebSocket جیمینای (کلید ...${apiKey.slice(-4)}):`, error);
     // بستن اتصال کلاینت مرتبط در صورت خطای جیمینای
     if (clientWs.readyState === WebSocket.OPEN) {
         clientWs.close();
     }
  });

  geminiWs.on('close', (code, reason) => {
    console.log(`اتصال WebSocket جیمینای بسته شد (کلید ...${apiKey.slice(-4)}):`, code, reason.toString());
     // بستن اتصال کلاینت مرتبط وقتی اتصال جیمینای بسته می‌شود
     if (clientWs.readyState === WebSocket.OPEN) {
         clientWs.close();
     }
  });

  return geminiWs;
};

wss.on('connection', (ws) => {
  console.log('یک کلاینت جدید متصل شد.');

  // --- اختصاص کلید API به محض اتصال کلاینت ---
  const assignedApiKey = getNextApiKey(); // گرفتن کلید بعدی برای *این* کلاینت

  // اگر به دلایلی کلید دریافت نشد (نباید اتفاق بیفتد اگر چک اولیه انجام شده)
  if (!assignedApiKey) {
       console.error("خطا: عدم موفقیت در اختصاص کلید API به کلاینت جدید. بستن اتصال.");
       ws.close();
       return; // ادامه نده برای این کلاینت
  }
  // لاگ کردن کلید اختصاص یافته (فقط ۴ کاراکتر آخر برای امنیت)
  console.log(`کلید API اختصاص یافته به کلاینت: ...${assignedApiKey.slice(-4)}`);

  let geminiWs = null; // اتصال جیمینای مخصوص *این* کلاینت

  ws.on('message', async (message) => {
    try {
        // بررسی کنیم پیام باینری است یا متنی
        let data;
        if (message instanceof Buffer) {
            // اگر کلاینت پیام باینری می‌فرستد، آن را به رشته تبدیل کنید (فرض بر JSON بودن)
            data = JSON.parse(message.toString());
        } else {
            data = JSON.parse(message); // اگر پیام متنی است
        }

        console.log('دریافت از کلاینت:', data);

        // مقداردهی اولیه اتصال جیمینای برای این کلاینت با استفاده از کلید اختصاص یافته
        if (data.setup) {
          // جلوگیری از مقداردهی مجدد اگر کلاینت دوباره پیام setup فرستاد
          if (geminiWs) {
              console.warn("کلاینت دوباره پیام setup ارسال کرد. نادیده گرفته شد.");
              return;
          }
          console.log('مقداردهی اولیه اتصال جیمینای با تنظیمات:', data.setup);
          // ارسال کلید اختصاص یافته این کلاینت به تابع ایجاد اتصال
          geminiWs = createGeminiWebSocket(ws, assignedApiKey);

          // ذخیره پیام setup برای ارسال پس از برقراری اتصال (اگر اتصال هنوز آماده نیست)
          if (geminiWs && geminiWs.readyState !== WebSocket.OPEN) {
             // مهم: پیام معلق را به نمونه WebSocket *این* کلاینت متصل کنید
             geminiWs.pendingSetup = data;
          } else if (geminiWs) {
              // اگر اتصال بلافاصله برقرار شد، پیام setup را بفرست
              try {
                  geminiWs.send(JSON.stringify(data));
              } catch (sendError) {
                   console.error("خطا در ارسال پیام setup اولیه به جیمینای:", sendError);
                   ws.close(); // بستن اتصال کلاینت
              }
          } else {
              // اگر createGeminiWebSocket ناموفق بود (مثلا به خاطر خطای کلید در بالا)
              console.error("ایجاد اتصال WebSocket جیمینای پس از پیام setup ناموفق بود.");
              // ws قبلا باید بسته شده باشد توسط createGeminiWebSocket
          }
          return; // پردازش پیام setup تمام شد
        }

        // ارسال پیام به جیمینای اگر اتصال برای این کلاینت وجود دارد و باز است
        if (geminiWs && geminiWs.readyState === WebSocket.OPEN) {
          console.log('ارسال به جیمینای:', data);
          try {
            geminiWs.send(JSON.stringify(data));
          } catch (sendError) {
              console.error("خطا در ارسال پیام به جیمینای:", sendError);
              ws.close(); // بستن اتصال کلاینت
          }
        } else if (geminiWs) {
          // اتصال جیمینای هنوز برقرار نشده یا بسته شده
          console.log('اتصال جیمینای آماده نیست، امکان ارسال پیام وجود ندارد.');
          // TODO: می‌توانید پیام‌ها را در صف قرار دهید یا خطا به کلاینت برگردانید
        } else {
          // هنوز پیام setup دریافت نشده یا اتصال ناموفق بوده
          console.error('امکان ارسال پیام وجود ندارد: اتصال جیمینای برای این کلاینت برقرار نشده است.');
          // می‌توانید خطا به کلاینت برگردانید
           if (ws.readyState === WebSocket.OPEN) {
               try {
                   ws.send(JSON.stringify({ error: "ارتباط با سرویس برقرار نیست. لطفاً دوباره تلاش کنید." }));
               } catch (sendError) {
                   console.error("خطا در ارسال پیام خطا به کلاینت:", sendError);
               }
           }
        }
    } catch (error) {
        // خطای کلی در پردازش پیام (مثلا JSON نامعتبر)
        console.error('خطا در پردازش پیام کلاینت:', error);
        // از کرش کردن سرور جلوگیری کنید
    }
  });

  ws.on('close', () => {
    console.log(`کلاینت با کلید ...${assignedApiKey ? assignedApiKey.slice(-4) : 'N/A'} قطع شد.`);
    // بستن اتصال جیمینای مربوط به این کلاینت، اگر باز است
    if (geminiWs && geminiWs.readyState !== WebSocket.CLOSED && geminiWs.readyState !== WebSocket.CLOSING) {
      geminiWs.close();
    }
  });

  ws.on('error', (error) => {
      console.error(`خطای WebSocket کلاینت (کلید ...${assignedApiKey ? assignedApiKey.slice(-4) : 'N/A'}):`, error);
      // بستن اتصال جیمینای مربوطه در صورت خطای کلاینت
      if (geminiWs && geminiWs.readyState !== WebSocket.CLOSED && geminiWs.readyState !== WebSocket.CLOSING) {
          geminiWs.close();
      }
      // ws احتمالا به طور خودکار بسته می‌شود یا در شرف بسته شدن است
  });

});

// رسیدگی به درخواست‌های باقیمانده با برگرداندن اپ React
app.get('*', (req, res) => {
  // اطمینان حاصل کنید که مسیر درست است
  const indexPath = path.join(__dirname, '../build', 'index.html');
  res.sendFile(indexPath, (err) => {
      if (err) {
          console.error("خطا در ارسال فایل index.html:", err);
          // ارسال یک پاسخ خطای ساده اگر فایل پیدا نشد
          if (!res.headersSent) {
              res.status(err.status || 500).send('خطا در بارگذاری برنامه');
          }
      }
  });
});


const PORT = process.env.PORT || 3001; // در هاگینگ فیس معمولا پورت 7860 استفاده می‌شود، اما اینجا 3001 است
server.listen(PORT, () => {
  console.log(`سرور در حال اجرا روی پورت ${PORT}`);
  console.log(`لطفاً ${numKeys} کلید API با نام‌های GEMINI_API_KEY_1 تا GEMINI_API_KEY_${numKeys} را در Secrets تنظیم کرده باشید.`);
});