import cluster from "cluster";
import os from "os";
import net from "net";
import fs from "fs";
import { spawnSync } from "child_process";
import path from "path";
import express from "express";
import { createServer } from "http";
import compression from "compression";
import WebSocket from "ws";
import { baremuxPath } from "@mercuryworkshop/bare-mux/node";
import { epoxyPath } from "@mercuryworkshop/epoxy-transport";
import { libcurlPath } from "@mercuryworkshop/libcurl-transport";
import { uvPath } from "@titaniumnetwork-dev/ultraviolet";
import wisp from "wisp-server-node";
import NodeCache from "node-cache";

const surgeConfigPath = path.resolve("surge.config.json");
const isSurgedRun = process.argv.includes("--surged");

function applySurgeAndRestartIfNeeded() {
  if (isSurgedRun) {
    try {
      const config = JSON.parse(fs.readFileSync(surgeConfigPath, "utf-8"));
      process.env.UV_THREADPOOL_SIZE = String(config.uvThreadpoolSize);
    } catch {
    }
    return;
  }

  console.log('[~] Running Surge...');
  const result = spawnSync("node", ["./others/surge.mjs"], { stdio: "inherit" });
  if (result.error) {
    console.error('[!] Surger failed:', result.error);
    process.exit(1);
  }

  const config = JSON.parse(fs.readFileSync(surgeConfigPath, "utf-8"));
  const nodeArgs = config.nodeFlags.concat([path.resolve("index.mjs"), "--surged"]);
  const env = { ...process.env, UV_THREADPOOL_SIZE: String(config.uvThreadpoolSize), ALREADY_SURGED: "true" };

  console.log(`[~] Relaunching with Node flags: ${config.nodeFlags.join(' ')}`);
  const relaunch = spawnSync(process.execPath, nodeArgs, { stdio: 'inherit', env });
  process.exit(relaunch.status || 0);
}

applySurgeAndRestartIfNeeded();

if (global.gc) {
  setInterval(() => global.gc(), 15000);
}

import './others/scaler.mjs';
import './others/warmup.mjs';

const cache = new NodeCache({ stdTTL: 345600, checkperiod: 3600, useClones: false });
const port = parseInt(process.env.PORT || "3000", 10);
cluster.schedulingPolicy = cluster.SCHED_RR;

function logInfo(message) { console.info(`[~] ${message}`); }
function logSuccess(message) { console.info(`[+] ${message}`); }
function logError(error) {
  const msg = error instanceof Error ? error.message : error;
  console.error(`[!] ${msg}`);
}

process.on("uncaughtException", err => logError(`Unhandled Exception: ${err}`));
process.on("unhandledRejection", reason => logError(`Unhandled Rejection: ${reason}`));

if (cluster.isPrimary) {
  const numCPUs = os.cpus().length;
  logInfo(`Master started. Forking ${numCPUs} workers...`);
  for (let i = 0; i < numCPUs; i++) cluster.fork();

  cluster.on("exit", (worker, code, signal) => {
    logError(`Worker ${worker.process.pid} exited (code=${code}, signal=${signal}). Restarting...`);
    cluster.fork();
  });

  let current = 0;
  const server = net.createServer({ pauseOnConnect: true }, conn => {
    const workers = Object.values(cluster.workers);
    if (!workers.length) return conn.destroy();
    const worker = workers[current++ % workers.length];
    worker.send("sticky-session:connection", conn);
  });

  server.on("error", err => logError(`Server error: ${err}`));
  server.listen(port, () => logSuccess(`Server listening on port ${port}`));

} else {
  const __dirname = process.cwd();
  const publicPath = path.join(__dirname, "public");
  const app = express();

  app.use(compression({ level: 9, threshold: 128, memLevel: 9 }));
  app.use((req, res, next) => {
    const key = req.originalUrl;
    if (cache.has(key)) {
      res.setHeader("X-Cache", "HIT");
      return res.send(cache.get(key));
    }
    res.sendResponse = res.send;
    res.send = body => {
      cache.set(key, body);
      res.setHeader("X-Cache", "MISS");
      res.sendResponse(body);
    };
    next();
  });

  const staticOpts = { maxAge: "7d", immutable: true };
  app.use("/baremux/", express.static(baremuxPath, staticOpts));
  app.use("/epoxy/", express.static(epoxyPath, staticOpts));
  app.use("/libcurl/", express.static(libcurlPath, staticOpts));
  app.use(express.static(publicPath, staticOpts));
  app.use("/wah/", express.static(uvPath, staticOpts));

  app.use(express.json());
  const sendHtml = file => (req, res) => res.sendFile(path.join(publicPath, file));
  app.get('/', sendHtml('$.html'));
  app.get('/g', sendHtml('!.html'));
  app.get('/a', sendHtml('!!.html'));
  app.get('/resent', (req, res) => res.sendFile(path.join(publicPath, 'resent', 'index.html')));
  app.use((req, res) => res.status(404).sendFile(path.join(publicPath, '404.html')));

  const server = createServer(app);
  server.keepAliveTimeout = 0;
  server.headersTimeout = 0;

  const pingWSS = new WebSocket.Server({ noServer: true, maxPayload: 4 * 1024 * 1024, perMessageDeflate: false });
  pingWSS.on("connection", (ws, req) => {
    const remote = req.socket.remoteAddress || 'unknown';
    let lat = [];
    const pingInterval = setInterval(() => {
      if (ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify({ type: 'ping', timestamp: Date.now() }));
    }, 1000);

    ws.on('message', msg => {
      try {
        const data = JSON.parse(msg);
        if (data.type === 'pong' && data.timestamp) {
          const delta = Date.now() - data.timestamp;
          lat.push(delta);
          if (lat.length > 5) lat.shift();
          ws.send(JSON.stringify({ type: 'latency', latency: delta }));
        }
      } catch (e) { logError(`Ping error: ${e}`); }
    });

    ws.on('close', () => {
      clearInterval(pingInterval);
      const avg = lat.length ? (lat.reduce((a,b)=>a+b)/lat.length).toFixed(2) : 0;
      logInfo(`WS ${remote} closed. Avg latency ${avg}ms`);
    });
  });

  server.on('upgrade', (req, sock, head) => {
    if (req.url === '/w/ping') pingWSS.handleUpgrade(req, sock, head, ws => pingWSS.emit('connection', ws, req));
    else if (req.url.startsWith('/w/')) wisp.routeRequest(req, sock, head);
    else sock.end();
  });

  server.on('error', err => logError(`Worker server error: ${err}`));
  server.listen(0, () => logSuccess(`Worker ${process.pid} ready`));

  process.on('message', (msg, conn) => {
    if (msg === 'sticky-session:connection' && conn) {
      server.emit('connection', conn);
      conn.resume();
    }
  });
}