import { Client as SSH } from 'ssh2-es';
import { EventEmitter } from 'node:events';

class NodeStream extends EventEmitter { // Adapts a Web Stream {readable, writable} to a Node.js-style stream.
  constructor(socket) {
    super();
    this.writer = socket.writable.getWriter();
    (async () => {
      try { for await (const chunk of socket.readable) this.emit('data', chunk); }
      catch (e) { this.emit('error', e); }
      finally { this.emit('close'); }
    })();
  }
  write(d) { return this.writer.write(d); }
  end() { return this.writer.close(); }
}

async function handleSession(ws) {
  let ssh, shell;
  const cleanup = () => { try { shell?.end(); ssh?.end(); ws.close(); } catch {} };
  ws.addEventListener('close', cleanup);
  ws.addEventListener('error', cleanup);

  try {
    const params = await new Promise((res, rej) => ws.addEventListener('message', (msg) => { try { res(JSON.parse(msg.data.toString())); } catch (e) { rej(e); } }, { once: true }));
    const tcpSocket = await connect({ hostname: params.host, port: Number(params.port ?? 22) }, { allowTls: false });
    
    ssh = new SSH(new NodeStream(tcpSocket), params);
    ssh.on('error', (e) => ws.send(`\r\n\x1b[31mSSH Error: ${e.message}\x1b[0m\r\n`));
    await new Promise(res => ssh.on('ready', res));

    shell = await new Promise((res, rej) => ssh.shell({ term: 'xterm-256color', ...params }, (e, s) => e ? rej(e) : res(s)));
    shell.on('data', (d) => ws.send(d)).on('close', () => ws.close(1000, 'Shell closed'));

    ws.addEventListener('message', (msg) => {
      try {
        const d = JSON.parse(msg.data.toString());
        if (d.resize) shell.setWindow(d.resize.rows, d.resize.cols, d.resize.height, d.resize.width);
      } catch { shell.write(msg.data); }
    });
  } catch (e) {
    ws.send(`\r\n\x1b[31mWorker Error: ${e.message}\x1b[0m\r\n`);
  }
}

export default {
  async fetch(request, env, ctx) {
    if (request.headers.get('Upgrade') === 'websocket') {
      const [client, server] = Object.values(new WebSocketPair());
      ctx.waitUntil(handleSession(server));
      return new Response(null, { status: 101, webSocket: client });
    }
    return env.ASSETS.fetch(request);
  },
};