const { WebSocketServer } = require('ws');
const Docker = require('dockerode');
const url = require('url');

const docker = new Docker();

/**
 * Attach WebSocket server for container exec/shell
 * Route: ws://host/ws/exec/:containerId
 * @param {http.Server} server - The HTTP server instance
 * @param {Object} log - Logger
 */
module.exports = function attachExecWS(server, log) {
  const wss = new WebSocketServer({ noServer: true });

  server.on('upgrade', (req, socket, head) => {
    const parsed = url.parse(req.url, true);
    const match = parsed.pathname.match(/^\/ws\/exec\/([a-zA-Z0-9_.-]+)$/);
    if (!match) return; // Not our route — let other handlers deal with it

    const containerId = decodeURIComponent(match[1]);

    wss.handleUpgrade(req, socket, head, (ws) => {
      handleExec(ws, containerId, log);
    });
  });

  return wss;
};

async function handleExec(ws, containerId, log) {
  let execStream = null;
  let execInstance = null;

  try {
    const container = docker.getContainer(containerId);
    // Verify container exists and is running
    const info = await container.inspect();
    if (!info.State.Running) {
      ws.send(JSON.stringify({ type: 'error', message: 'Container is not running' }));
      ws.close();
      return;
    }

    // Detect available shell
    let shell = '/bin/sh';
    try {
      const bashCheck = await container.exec({ Cmd: ['which', 'bash'], AttachStdout: true });
      const bashStream = await bashCheck.start();
      const chunks = [];
      await new Promise((resolve) => {
        bashStream.on('data', (chunk) => chunks.push(chunk));
        bashStream.on('end', resolve);
      });
      if (chunks.length > 0 && Buffer.concat(chunks).toString().includes('/bash')) {
        shell = '/bin/bash';
      }
    } catch (_) {}

    execInstance = await container.exec({
      Cmd: [shell],
      AttachStdin: true,
      AttachStdout: true,
      AttachStderr: true,
      Tty: true,
    });

    execStream = await execInstance.start({ hijack: true, stdin: true, Tty: true });

    ws.send(JSON.stringify({ type: 'connected', shell, containerId }));

    // Docker → WebSocket
    execStream.on('data', (chunk) => {
      if (ws.readyState === ws.OPEN) {
        ws.send(chunk);
      }
    });

    execStream.on('end', () => {
      if (ws.readyState === ws.OPEN) {
        ws.send(JSON.stringify({ type: 'exit' }));
        ws.close();
      }
    });

    // WebSocket → Docker
    ws.on('message', (data) => {
      if (!execStream.writable) return;
      try {
        // Check for control messages (JSON)
        const str = data.toString();
        if (str.startsWith('{"type":')) {
          const msg = JSON.parse(str);
          if (msg.type === 'resize' && execInstance && msg.cols && msg.rows) {
            execInstance.resize({ h: msg.rows, w: msg.cols }).catch(() => {});
            return;
          }
        }
      } catch (_) {}
      // Regular terminal input
      execStream.write(data);
    });

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

    ws.on('error', (err) => {
      log.warn('exec', 'WebSocket error', { containerId, error: err.message });
      if (execStream) {
        try { execStream.destroy(); } catch (_) {}
      }
    });

  } catch (err) {
    log.error('exec', 'Failed to start exec session', { containerId, error: err.message });
    if (ws.readyState === ws.OPEN) {
      ws.send(JSON.stringify({ type: 'error', message: err.message }));
      ws.close();
    }
  }
}