import type { Server } from 'node:http';
import type { Duplex } from 'node:stream';
import type { WebSocket } from 'ws';
import { WebSocketServer } from 'ws';

import type { BoxHandle } from './user-container';
import { env } from './env';
import { getDockerClient } from './runtime/docker';
import { verifyTerminalToken } from './terminal-token';
import { acquireUserBox } from './user-container';
import { getTerminalWorkspace } from './worker';

const clampDimension = (value: unknown) => {
  const n = Math.trunc(Number(value));
  if (!Number.isFinite(n)) return undefined;
  return Math.min(Math.max(n, 1), 1000);
};

export const parseResizeMessage = (
  data: Buffer,
  isBinary: boolean,
): { cols: number; rows: number } | null => {
  if (isBinary) return null;
  try {
    const message = JSON.parse(data.toString('utf8')) as {
      type?: string;
      cols?: number;
      rows?: number;
    };
    if (message.type !== 'resize') return null;
    const cols = clampDimension(message.cols);
    const rows = clampDimension(message.rows);
    if (!cols || !rows) return null;
    return { cols, rows };
  } catch {
    return null;
  }
};

const bridge = async (ws: WebSocket, jobId: string) => {
  const workspace = getTerminalWorkspace(jobId);
  if (!workspace) {
    ws.close(1011, 'Workspace is not active.');
    return;
  }

  // Register the message handler immediately and buffer input/size until the exec
  // is ready (acquiring the box can take seconds on first connect), so the initial
  // resize and early keystrokes aren't dropped.
  const execHolder: {
    current?: { exec: import('dockerode').Exec; stream: Duplex };
  } = {};
  const pendingInput: Buffer[] = [];
  let cols = 80;
  let rows = 24;

  ws.on('message', (data: Buffer, isBinary: boolean) => {
    const resize = parseResizeMessage(data, isBinary);
    if (resize) {
      cols = resize.cols;
      rows = resize.rows;
      void execHolder.current?.exec
        .resize({ h: rows, w: cols })
        .catch(() => undefined);
      return;
    }
    if (execHolder.current) execHolder.current.stream.write(data);
    else pendingInput.push(data);
  });

  const handleHolder: { current?: BoxHandle } = {};
  let released = false;
  // Read through a function so TS doesn't narrow `released` to a constant — the
  // cleanup handler flips it asynchronously when the socket closes.
  const isReleased = () => released;
  const cleanup = () => {
    if (released) return;
    released = true;
    execHolder.current?.stream.end();
    execHolder.current?.stream.destroy();
    handleHolder.current?.release();
  };
  ws.on('close', cleanup);
  ws.on('error', cleanup);

  // Hold the per-user box open while this terminal is connected; the agent and
  // the terminal share the exact same container (Phase 2).
  let boxName: string;
  try {
    const handle = await acquireUserBox({
      username: workspace.username,
      workdir: workspace.workdir,
      containerHome: workspace.containerHome,
    });
    handleHolder.current = handle;
    boxName = handle.boxName;
  } catch (error) {
    ws.close(
      1011,
      `Failed to start terminal: ${error instanceof Error ? error.message : 'unknown error'}`,
    );
    return;
  }

  // Cleanup may have run before the awaited handle existed.
  if (isReleased()) {
    handleHolder.current.release();
    return;
  }

  const docker = getDockerClient();
  const container = docker.getContainer(boxName);
  const exec = await container.exec({
    AttachStdin: true,
    AttachStdout: true,
    AttachStderr: true,
    Tty: true,
    Cmd: [
      '/bin/bash',
      '-lc',
      'if command -v tmux >/dev/null 2>&1; then exec tmux new-session -A -s spoon; else exec bash -il; fi',
    ],
    Env: [
      'TERM=xterm-256color',
      `HOME=${workspace.containerHome}`,
      ...workspace.secrets.map((secret) => `${secret.name}=${secret.value}`),
    ],
    WorkingDir: workspace.containerRepo,
  });
  const stream = await exec.start({ hijack: true, stdin: true, Tty: true });
  execHolder.current = { exec, stream };

  if (isReleased()) {
    stream.end();
    stream.destroy();
    return;
  }

  await exec.resize({ h: rows, w: cols }).catch(() => undefined);

  // Replay any keystrokes the client sent before the process was ready.
  for (const buffered of pendingInput) stream.write(buffered);
  pendingInput.length = 0;

  stream.on('data', (chunk: Buffer) => {
    if (ws.readyState === ws.OPEN) ws.send(chunk, { binary: true });
  });
  stream.on('end', () => {
    if (ws.readyState === ws.OPEN) ws.close();
  });
  stream.on('error', () => {
    if (ws.readyState === ws.OPEN) ws.close();
  });
};

/**
 * Attaches the interactive-terminal WebSocket endpoint to the worker's HTTP
 * server. Browser connects to `/jobs/:jobId/terminal?token=…` with a short-lived
 * token minted by the Next app (which has already verified job ownership).
 */
export const attachTerminalServer = (server: Server) => {
  if (env.runtime !== 'docker') return;
  const wss = new WebSocketServer({ noServer: true });

  server.on('upgrade', (request, socket, head) => {
    const url = new URL(request.url ?? '', `http://localhost:${env.httpPort}`);
    const match = /^\/jobs\/([^/]+)\/terminal$/.exec(url.pathname);
    if (!match?.[1]) {
      socket.destroy();
      return;
    }
    const jobId = decodeURIComponent(match[1]);
    const token = url.searchParams.get('token') ?? '';
    if (!verifyTerminalToken(token, jobId, env.terminalSecret)) {
      socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
      socket.destroy();
      return;
    }
    wss.handleUpgrade(request, socket, head, (ws) => {
      void bridge(ws, jobId).catch(() => {
        if (ws.readyState === ws.OPEN) ws.close();
      });
    });
  });

  console.log('Spoon agent worker terminal WebSocket endpoint enabled.');
};
