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

const docker = new Docker();

/**
 * DC-072: WebSocket scope authorization — admin-only by default.
 *
 * Container exec is full root-equivalent access inside the target
 * container. Granting it to a key whose scope is `['read']` violates
 * least privilege. The validScopes list (`['read','write','admin']`)
 * is defined in routes/auth/keys.js; exec requires `admin`.
 *
 * Defensive: the scope field is coerced via `Array.isArray(...) ? ... : []`
 * so a malformed payload (string, object, null, undefined) cannot reach
 * `.includes('admin')` and accidentally grant access. Every malformed
 * shape falls into the rejection branch with the same 403 envelope.
 *
 * Tests should call `__test.assertExecScope(auth)` directly rather
 * than spinning up a WebSocket server.
 */
function assertExecScope(auth) {
  const scope = Array.isArray(auth && auth.scope) ? auth.scope : [];
  if (!scope.includes('admin')) {
    const err = new Error('Container exec requires admin scope');
    err.code = 'DC-072_INSUFFICIENT_SCOPE';
    err.statusCode = 403;
    err.requiredScope = 'admin';
    err.actualScope = scope;
    throw err;
  }
}

/**
 * DC-072: Tighten containerId validation.
 *
 * Docker container IDs are exactly 64 lowercase hex chars (or 12-char
 * short form). The pre-fix regex accepted `_`, `-`, `.`, mixed case,
 * and up to 128 chars — Docker would then 404 the inspect call and
 * the rejection would surface as a generic 500 in the WS error
 * envelope. Pre-validate at the upgrade layer so the rejection is
 * fast and the log line discriminates "malformed" from "unknown".
 */
function isValidContainerId(id) {
  if (typeof id !== 'string') return false;
  // Full 64-char hex, or 12-char short hex
  return /^[0-9a-f]{64}$/.test(id) || /^[0-9a-f]{12}$/.test(id);
}

/**
 * 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, authManager) {
  const wss = new WebSocketServer({ noServer: true });

  // Authenticate WebSocket upgrade request before accepting it
  server.on('upgrade', async (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]);

    // DC-072: Tighten containerId charset (64-char / 12-char lowercase hex)
    if (!isValidContainerId(containerId)) {
      log.warn('exec', 'Invalid container ID in WebSocket path', { containerId });
      socket.write('HTTP/1.1 400 Bad Request\r\n\r\n');
      socket.destroy();
      return;
    }

    // Check auth — require valid JWT or API key from the upgrade request
    let auth = null;
    const authHeader = req.headers.authorization;
    if (authHeader && authHeader.startsWith('Bearer ')) {
      const token = authHeader.substring(7);
      const payload = await authManager.verifyJWT(token);
      if (payload) {
        auth = { type: 'jwt', userId: payload.userId, scope: payload.scope || [] };
      }
    } else {
      const apiKey = req.headers['x-api-key'];
      if (apiKey) {
        const keyData = await authManager.verifyAPIKey(apiKey);
        if (keyData) {
          auth = { type: 'apikey', keyId: keyData.keyId, scope: keyData.scopes || [] };
        }
      }
    }

    if (!auth) {
      log.warn('exec', 'Unauthenticated WebSocket exec attempt', { containerId, ip: req.socket.remoteAddress });
      socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
      socket.destroy();
      return;
    }

    // DC-072: Container exec is root-equivalent — require admin scope.
    // Pre-fix, a key issued with scope `['read']` (e.g., for monitoring)
    // would get a full PTY shell inside any running container. The
    // `auth.scope` was captured at lines 39/46 but never checked.
    try {
      assertExecScope(auth);
    } catch (err) {
      log.warn('exec', 'Insufficient scope for exec attempt', {
        containerId,
        authType: auth.type,
        authId: auth.type === 'jwt' ? auth.userId : auth.keyId,
        actualScope: err.actualScope,
        requiredScope: err.requiredScope,
        ip: req.socket.remoteAddress,
      });
      // 403 with a JSON error envelope over the upgrade socket so the
      // dashboard can display "admin required" instead of guessing.
      socket.write('HTTP/1.1 403 Forbidden\r\n');
      socket.write('Content-Type: application/json\r\n');
      socket.write('\r\n');
      socket.end(JSON.stringify({
        error: err.message,
        code: err.code,
        requiredScope: err.requiredScope,
        actualScope: err.actualScope,
      }));
      return;
    }

    // Auth passed — proceed with WebSocket upgrade
    wss.handleUpgrade(req, socket, head, (ws) => {
      handleExec(ws, containerId, log, auth);
    });
  });

  return wss;
};

async function handleExec(ws, containerId, log, auth) {
  let execStream = null;
  let execInstance = null;
  const sessionStart = Date.now();

  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;
    }

    // DC-072: Audit-log the exec session start. Pairs with the end-log
    // below so the operator can correlate who opened which shell.
    log.info('exec', 'Authenticated exec session started', {
      containerId,
      authType: auth.type,
      authId: auth.type === 'jwt' ? auth.userId : auth.keyId,
      containerName: info.Name,
    });

    // 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 (_) {
      // Fall back to /bin/sh when bash detection fails
    }

    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);
      }
    });

    // DC-072: Track whether the end-log has fired so we don't double-log
    // when both execStream 'end' and ws 'close' fire (Docker stream end
    // closes the WS, which then fires 'close' too — without the flag
    // we'd emit the same audit line twice with the same durationMs).
    let ended = false;
    const logSessionEnd = (reason) => {
      if (ended) return;
      ended = true;
      log.info('exec', 'Exec session ended', {
        containerId,
        authType: auth.type,
        authId: auth.type === 'jwt' ? auth.userId : auth.keyId,
        durationMs: Date.now() - sessionStart,
        reason,
      });
    };

    execStream.on('end', () => {
      // DC-072: Audit-log the session end (duration + container) so a
      // long-running session is observable in the error log. Normal
      // shutdown path: Docker exec stream closes → log + tell client.
      logSessionEnd('exec-stream-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 (_) {
        // Treat message as raw terminal input if JSON parsing fails
      }
      // Regular terminal input
      execStream.write(data);
    });

    ws.on('close', () => {
      // DC-072: Fallback audit-log for abnormal close (browser tab
      // closed, network drop, container killed mid-session) where the
      // execStream 'end' event never fires. The ended-flag guard makes
      // this idempotent with the normal path above.
      logSessionEnd('ws-close');
      if (execStream) {
        try { execStream.destroy(); } catch (_) {
          // Ignore stream teardown errors on socket close
        }
      }
    });

    ws.on('error', (err) => {
      log.warn('exec', 'WebSocket error', { containerId, error: err.message });
      if (execStream) {
        try { execStream.destroy(); } catch (_) {
          // Ignore stream teardown errors after websocket errors
        }
      }
    });

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

// Internal-only export for unit tests. Stripped from the public
// surface; tests import this via the destructure form
// `const { __test } = require('./routes/exec')`.
module.exports.__test = {
  assertExecScope,
  isValidContainerId,
};