import http from 'node:http';
import net from 'node:net';
import { WebSocketServer } from 'ws';
import { HTMLRewriter } from 'html-rewriter-wasm';
import worker from './worker.js';

// Polyfill HTMLRewriter globally for worker.js
globalThis.HTMLRewriter = HTMLRewriter;

// Polyfill caches.default globally in Node.js
if (!globalThis.caches) {
  const cacheStore = new Map();
  globalThis.caches = {
    default: {
      async match(req) {
        const key = typeof req === 'string' ? req : (req.url || String(req));
        const item = cacheStore.get(key);
        if (!item) return null;
        return item.clone();
      },
      async put(req, res) {
        const key = typeof req === 'string' ? req : (req.url || String(req));
        cacheStore.set(key, res.clone());
      },
      async delete(req) {
        const key = typeof req === 'string' ? req : (req.url || String(req));
        return cacheStore.delete(key);
      }
    }
  };
}

const PORT = process.env.PORT || 7860;

// Convert Node.js http.IncomingMessage to WHATWG Request
async function nodeReqToWebReq(req, origin) {
  const url = new URL(req.url, origin);
  const headers = new Headers();
  for (const [key, val] of Object.entries(req.headers)) {
    if (Array.isArray(val)) {
      val.forEach(v => headers.append(key, v));
    } else if (val !== undefined) {
      headers.set(key, val);
    }
  }

  let body = null;
  if (req.method !== 'GET' && req.method !== 'HEAD') {
    const chunks = [];
    for await (const chunk of req) {
      chunks.push(chunk);
    }
    body = Buffer.concat(chunks);
  }

  return new Request(url.toString(), {
    method: req.method,
    headers,
    body,
    duplex: body ? 'half' : undefined
  });
}

// Convert WHATWG Response to Node.js http.ServerResponse
async function webResToNodeRes(webRes, res) {
  res.statusCode = webRes.status;
  res.statusMessage = webRes.statusText || '';

  // Write headers
  if (typeof webRes.headers.getSetCookie === 'function') {
    const cookies = webRes.headers.getSetCookie();
    if (cookies && cookies.length > 0) {
      res.setHeader('Set-Cookie', cookies);
    }
  }
  const skipHeaders = new Set(['set-cookie', 'content-encoding', 'transfer-encoding']);
  webRes.headers.forEach((value, key) => {
    const lk = key.toLowerCase();
    if (!skipHeaders.has(lk)) {
      res.setHeader(key, value);
    }
  });

  if (!webRes.body) {
    res.end();
    return;
  }

  const reader = webRes.body.getReader();
  let closed = false;
  let completed = false;
  const onClose = () => {
    if (closed) return;
    closed = true;
    try { reader.cancel(); } catch (_) {}
  };
  res.on('error', () => {
    closed = true;
    try { reader.cancel(); } catch (_) {}
  });
  res.on('close', onClose);
  try {
    while (!closed) {
      const { done, value } = await reader.read();
      if (done) { completed = true; break; }
      if (closed) break;
      if (!res.write(Buffer.from(value))) {
        await new Promise((resolve) => {
          const done2 = () => resolve();
          res.once('drain', done2);
          res.once('close', done2);
        });
        if (closed) break;
      }
    }
  } catch (err) {
    completed = false;
  } finally {
    res.removeListener('close', onClose);
    if (!closed) {
      if (completed) {
        try { res.end(); } catch (_) {}
      } else {
        try { res.destroy(); } catch (_) {}
      }
    }
  }
}

const server = http.createServer(async (req, res) => {
  try {
    const protocol = req.headers['x-forwarded-proto'] || 'http';
    const host = req.headers.host || `localhost:${PORT}`;
    const origin = `${protocol}://${host}`;

    const webReq = await nodeReqToWebReq(req, origin);

    const ctx = {
      waitUntil(promise) {
        promise?.catch?.(err => console.error('[waitUntil error]', err));
      },
      passThroughOnException() {}
    };

    const webRes = await worker.fetch(webReq, process.env, ctx);
    await webResToNodeRes(webRes, res);
  } catch (err) {
    console.error('[Server Error]', err);
    if (!res.headersSent) {
      res.statusCode = 500;
      res.end(`Internal Server Error: ${err.message}`);
    } else {
      try { res.destroy(); } catch (_) {}
    }
  }
});

// WebSocket Server for VLESS on Hugging Face Spaces / Node.js
const wss = new WebSocketServer({ noServer: true });

server.on('upgrade', (req, socket, head) => {
  wss.handleUpgrade(req, socket, head, (ws) => {
    let tcpSocket = null;
    let isHeaderParsed = false;

    ws.on('message', (data) => {
      try {
        const buf = Buffer.from(data);
        if (!isHeaderParsed) {
          isHeaderParsed = true;
          if (buf.length < 24) {
            ws.close(1008, 'Header too short');
            return;
          }
          const version = buf[0];
          const addonLen = buf[17];
          let offset = 18 + addonLen;
          const command = buf[offset++];
          const port = buf.readUInt16BE(offset);
          offset += 2;
          const addrType = buf[offset++];
          let address = '';
          if (addrType === 1) {
            address = Array.from(buf.subarray(offset, offset + 4)).join('.');
            offset += 4;
          } else if (addrType === 2) {
            const domainLen = buf[offset++];
            address = buf.toString('utf8', offset, offset + domainLen);
            offset += domainLen;
          } else if (addrType === 3) {
            const parts = [];
            for (let i = 0; i < 8; i++) parts.push(buf.readUInt16BE(offset + i * 2).toString(16));
            address = parts.join(':');
            offset += 16;
          }

          const rawData = buf.subarray(offset);

          // Reply VLESS OK: [version, 0]
          ws.send(Buffer.from([version, 0]));

          tcpSocket = net.connect({ host: address, port: port }, () => {
            if (rawData && rawData.length > 0) {
              tcpSocket.write(rawData);
            }
          });

          tcpSocket.on('data', (chunk) => {
            if (ws.readyState === 1) ws.send(chunk);
          });

          tcpSocket.on('error', (err) => {
            console.error('[TCP Socket Error]', address, port, err.message);
            try { ws.close(); } catch (_) {}
          });

          tcpSocket.on('close', () => {
            try { ws.close(); } catch (_) {}
          });

        } else if (tcpSocket && tcpSocket.writable) {
          tcpSocket.write(buf);
        }
      } catch (err) {
        console.error('[WS Message Error]', err);
        try { ws.close(); } catch (_) {}
      }
    });

    ws.on('close', () => {
      if (tcpSocket) try { tcpSocket.destroy(); } catch (_) {}
    });

    ws.on('error', () => {
      if (tcpSocket) try { tcpSocket.destroy(); } catch (_) {}
    });
  });
});

server.listen(PORT, '0.0.0.0', () => {
  console.log(`=======================================================`);
  console.log(`🚀 Etrog Bypass Proxy running on Hugging Face Space`);
  console.log(`📡 Listening on http://0.0.0.0:${PORT}`);
  console.log(`=======================================================`);
});