platform/src/servers/api/terminal/websocket.ts

import type { ServerWebSocket } from 'bun';
import { existsSync, mkdirSync, statSync, writeFileSync } from 'node:fs';

import { dirname, join } from 'node:path';
import { fileURLToPath } from 'node:url';
import { getHomeDir, getGlobalSkillsDir, getGlobalToolsDir, getGlobalExtensionsDir, getUserSkillsDir, getUserToolsDir, DATA_PATH, PI_CONFIG_DIR, toShellUsername } from '@@/data-path';
import { getUsers } from 'officerdb';
import { generateContainerContext, generateClaudeSettings } from '@@/generate-container-context';

const ensureDir = (dir: string) => { if (!existsSync(dir)) mkdirSync(dir, { recursive: true }); return dir; };

type WSData = { userId: number; email: string; username: string; role: string; sandboxed: boolean; sessionId?: string; cwd?: string; cols?: number; rows?: number };
type ShellInfo = { command: string; args: string[]; name: string };
type BridgeSession = {
  client: ServerWebSocket<WSData>;
  sidecar: WebSocket | null;
  dockerId: string;
  port: number;
  pendingMessages: string[];
};

type ContainerInfo = {
  userId: number;
  email: string;
  dockerId: string;
  port: number;
};

const HOST_SIDECAR_PORT = 5338;

const sessions = new Map<ServerWebSocket<WSData>, BridgeSession>();
const containerMapPath = join(getHomeDir(''), '..', 'terminal-containers.json');

let dockerImageReady = false;
let containersCache: Record<string, ContainerInfo> | null = null;
let hostSidecarProcess: ReturnType<typeof import('bun').spawn> | null = null;

const sendOutput = (ws: ServerWebSocket<WSData>, data: string) => {
  try {
    ws.send(JSON.stringify({ type: 'output', data }));
  } catch {
    // ws already closed
  }
};

const connectSidecar = async (port: number): Promise<WebSocket> => {
  const delays = [200, 300, 500, 800, 1200, 1600, 2000];
  let lastError: Error | null = null;

  for (const delay of delays) {
    try {
      const ws = await new Promise<WebSocket>((resolve, reject) => {
        const socket = new WebSocket(`ws://127.0.0.1:${port}`);
        const timeout = setTimeout(() => {
          try {
            socket.close();
          } catch {
            // ignore
          }
          reject(new Error('Terminal sidecar timeout'));
        }, 2000);

        socket.addEventListener('open', () => {
          clearTimeout(timeout);
          resolve(socket);
        });
        socket.addEventListener('error', () => {
          clearTimeout(timeout);
          reject(new Error('Terminal sidecar connection failed'));
        });
      });

      return ws;
    } catch (err) {
      lastError = err instanceof Error ? err : new Error('Terminal sidecar connection failed');
      await new Promise((resolve) => setTimeout(resolve, delay));
    }
  }

  throw lastError ?? new Error('Terminal sidecar connection failed');
};

const ensureDockerImage = () => {
  if (dockerImageReady) return;
  const dockerPath = Bun.which('docker');
  if (!dockerPath) throw new Error('Docker not found');

  const tag = 'officer-terminal-sidecar:v1';
  const inspect = Bun.spawnSync({ cmd: [dockerPath, 'image', 'inspect', tag], stdout: 'ignore', stderr: 'ignore' });
  if (inspect.exitCode === 0) {
    dockerImageReady = true;
    return;
  }

  const dockerfilePath = fileURLToPath(new URL('./Dockerfile.terminal-sidecar', import.meta.url));
  const build = Bun.spawnSync({
    cmd: [dockerPath, 'build', '-f', dockerfilePath, '-t', tag, '.'],
    cwd: fileURLToPath(new URL('./', import.meta.url)),
    stdout: 'inherit',
    stderr: 'inherit',
  });

  if (build.exitCode !== 0) throw new Error('Failed to build terminal sandbox image');
  dockerImageReady = true;
};

// Check whether a container has all expected volume mounts.
// Tests for multiple mount sources — if any is missing, the container should be recreated.
const containerHasExpectedMounts = (dockerId: string): boolean => {
  const dockerPath = Bun.which('docker') ?? 'docker';
  const result = Bun.spawnSync({
    cmd: [dockerPath, 'inspect', '--format', '{{range .Mounts}}{{.Source}}\n{{end}}', dockerId],
    stdout: 'pipe',
    stderr: 'ignore',
  });
  if (result.exitCode !== 0) return false;
  const mounts = result.stdout.toString();
  return mounts.includes(getGlobalSkillsDir()) && mounts.includes('/officer/data') && mounts.includes('.claude');
};

const startDockerSidecar = async (port: number, homeDir: string, userId: number, username: string, email: string, contextFile?: string, settingsFile?: string): Promise<{ dockerId: string }> => {
  ensureDockerImage();
  const dockerPath = Bun.which('docker') ?? 'docker';
  const dockerId = `officer-terminal-${userId}`;
  const tag = 'officer-terminal-sidecar:v1';

  // Remove stale container with same name if it exists
  if (dockerContainerExists(dockerId)) {
    Bun.spawnSync({ cmd: [dockerPath, 'rm', '-f', dockerId], stdout: 'ignore', stderr: 'ignore' });
  }

  let uid = 1000;
  let gid = 1000;
  try {
    const stats = statSync(homeDir);
    uid = stats.uid;
    gid = stats.gid;
  } catch {
    // fallback to defaults
  }

  const containerHome = `/home/${username}`;

  const run = Bun.spawnSync({
    cmd: [
      dockerPath,
      'run',
      '-d',
      '--name',
      dockerId,
      '--restart',
      'unless-stopped',
      '--network', 'host',
      '-e',
      `TERMINAL_PTY_PORT=${port}`,
      '-e',
      'TERMINAL_PTY_HOST=127.0.0.1',
      '-e',
      `TERMINAL_USER=${username}`,
      '-e',
      `TERMINAL_UID=${uid}`,
      '-e',
      `TERMINAL_GID=${gid}`,
      '-e',
      `OFFICER_EMAIL=${email}`,
      '-v', `${homeDir}:${containerHome}`,
      '-v', `${getGlobalSkillsDir()}:/officer/skills:ro`,
      '-v', `${getGlobalToolsDir()}:/officer/tools:ro`,
      '-v', `${getGlobalExtensionsDir()}:/officer/extensions:ro`,
      '-v', `${getUserSkillsDir(email)}:/officer/user/skills:ro`,
      '-v', `${getUserToolsDir(email)}:/officer/user/tools:ro`,
      '-v', `${PI_CONFIG_DIR}:/officer/pi-config:ro`,
      '-v', `${PI_CONFIG_DIR}:${containerHome}/.pi/agent`,
      '-v', `${ensureDir(join(getHomeDir(email), '.pi', 'agent', 'sessions'))}:${containerHome}/.pi/agent/sessions`,
      '-v', `${join(DATA_PATH, '.generated')}:/officer/generated:ro`,
      '-v', `${join(DATA_PATH, email)}:/officer/data`,
      ...(existsSync(join(process.env.HOME ?? '', '.claude')) ? ['-v', `${join(process.env.HOME!, '.claude')}:${containerHome}/.claude`] : []),
      ...(contextFile && existsSync(contextFile) ? ['-v', `${contextFile}:${containerHome}/.claude/CLAUDE.md:ro`] : []),
      ...(settingsFile && existsSync(settingsFile) ? ['-v', `${settingsFile}:${containerHome}/.claude/settings.json:ro`] : []),
      '-w', containerHome,
      tag,
    ],
    stdout: 'inherit',
    stderr: 'inherit',
  });

  if (run.exitCode !== 0) throw new Error('Failed to start terminal sandbox container');

  // Wait for entrypoint to finish (user creation, sidecar start)
  for (let i = 0; i < 20; i++) {
    await new Promise((r) => setTimeout(r, 500));
    if (await sidecarAlive(port)) return { dockerId };
  }
  throw new Error('Terminal sidecar did not start in time');
};

const stopDockerSidecar = (dockerId: string) => {
  const dockerPath = Bun.which('docker') ?? 'docker';
  Bun.spawnSync({ cmd: [dockerPath, 'rm', '-f', dockerId], stdout: 'ignore', stderr: 'ignore' });
};

const readDockerLogs = (dockerId: string) => {
  const dockerPath = Bun.which('docker') ?? 'docker';
  const logs = Bun.spawnSync({ cmd: [dockerPath, 'logs', '--tail', '200', dockerId], stdout: 'pipe', stderr: 'pipe' });
  if (logs.exitCode !== 0) return '';
  return logs.stdout.toString().trim();
};

const loadContainerMap = async (): Promise<Record<string, ContainerInfo>> => {
  if (containersCache) return containersCache;
  const data = await Bun.file(containerMapPath)
    .json()
    .catch(() => ({}));
  containersCache = data as Record<string, ContainerInfo>;
  return containersCache;
};

const saveContainerMap = async (map: Record<string, ContainerInfo>) => {
  containersCache = map;
  await Bun.write(containerMapPath, JSON.stringify(map, null, 2));
};

const getAvailablePort = (map: Record<string, ContainerInfo>, userId: number) => {
  const base = 54000;
  const used = new Set(Object.values(map).map((item) => item.port));
  let port = base + (userId % 1000);
  while (used.has(port)) port += 1;
  return port;
};

const dockerContainerExists = (dockerId: string) => {
  const dockerPath = Bun.which('docker') ?? 'docker';
  const result = Bun.spawnSync({ cmd: [dockerPath, 'ps', '-a', '-q', '-f', `name=${dockerId}`], stdout: 'pipe' });
  return result.exitCode === 0 && result.stdout.toString().trim().length > 0;
};

const dockerContainerRunning = (dockerId: string) => {
  const dockerPath = Bun.which('docker') ?? 'docker';
  const result = Bun.spawnSync({ cmd: [dockerPath, 'ps', '-q', '-f', `name=${dockerId}`], stdout: 'pipe' });
  return result.exitCode === 0 && result.stdout.toString().trim().length > 0;
};

const dockerStart = (dockerId: string) => {
  const dockerPath = Bun.which('docker') ?? 'docker';
  const result = Bun.spawnSync({ cmd: [dockerPath, 'start', dockerId], stdout: 'ignore', stderr: 'ignore' });
  return result.exitCode === 0;
};

export const ensureDockerContainer = async (email: string, userId: number, homeDir: string, username: string, contextFile?: string, settingsFile?: string) => {
  // Check if mount sources are stale (e.g. data dir was deleted and Docker recreated them as root)
  // Must check BEFORE mkdirSync overwrites them
  const skillsDir = getUserSkillsDir(email);
  let stale = false;
  try {
    const s = statSync(skillsDir);
    if (s.uid === 0) stale = true;
  } catch {
    // doesn't exist yet — not stale, will be created below
  }

  // Ensure user-specific resource dirs exist before mounting (Docker creates them as root if missing)
  mkdirSync(skillsDir, { recursive: true });
  mkdirSync(getUserToolsDir(email), { recursive: true });
  mkdirSync(join(DATA_PATH, email, 'integrations'), { recursive: true });

  const map = await loadContainerMap();
  const existing = map[email];

  if (existing && dockerContainerRunning(existing.dockerId)) {
    // Recreate if resource mounts are missing or data dir was recreated (stale mounts)
    if (!containerHasExpectedMounts(existing.dockerId) || stale) {
      console.log(`[terminal] recreating container for ${email} — mounts stale or missing`);
      stopDockerSidecar(existing.dockerId);
    } else {
      console.log(`[terminal] reusing running container ${existing.dockerId} for ${email} on port ${existing.port}`);
      return existing;
    }
  }

  if (existing && dockerContainerExists(existing.dockerId)) {
    if (!containerHasExpectedMounts(existing.dockerId) || stale) {
      stopDockerSidecar(existing.dockerId);
    } else if (dockerStart(existing.dockerId)) {
      return existing;
    } else {
      stopDockerSidecar(existing.dockerId);
    }
  }

  const port = existing?.port ?? getAvailablePort(map, userId);
  const docker = await startDockerSidecar(port, homeDir, userId, username, email, contextFile, settingsFile);
  const next = { userId, email, dockerId: docker.dockerId, port };
  map[email] = next;
  await saveContainerMap(map);
  return next;
};

const sidecarAlive = async (port: number): Promise<boolean> => {
  try {
    const res = await fetch(`http://127.0.0.1:${port}`, { signal: AbortSignal.timeout(500) });
    return res.ok;
  } catch {
    return false;
  }
};

const killSidecarOnPort = (port: number) => {
  try {
    const result = Bun.spawnSync({ cmd: ['fuser', '-k', `${port}/tcp`], stdout: 'ignore', stderr: 'ignore' });
    if (result.exitCode === 0) console.log(`[terminal] killed stale sidecar on port ${port}`);
  } catch {
    // fuser not available or failed
  }
};

const startHostSidecar = async () => {
  if (hostSidecarProcess) {
    hostSidecarProcess.kill();
    await hostSidecarProcess.exited.catch(() => {});
    hostSidecarProcess = null;
  }

  killSidecarOnPort(HOST_SIDECAR_PORT);
  await new Promise((resolve) => setTimeout(resolve, 200));

  const sidecarPath = fileURLToPath(new URL('./pty-sidecar.mjs', import.meta.url));
  hostSidecarProcess = Bun.spawn({
    cmd: ['node', sidecarPath],
    env: { ...process.env, TERMINAL_PTY_PORT: String(HOST_SIDECAR_PORT) },
    stdout: 'inherit',
    stderr: 'inherit',
  });
  console.log(`[terminal] host sidecar started on port ${HOST_SIDECAR_PORT}`);
};

export const ensureHostSidecar = async () => {
  const alive = await sidecarAlive(HOST_SIDECAR_PORT);
  if (!alive) {
    await startHostSidecar();
    // Wait for it to come up
    for (let i = 0; i < 10; i++) {
      await new Promise((r) => setTimeout(r, 300));
      if (await sidecarAlive(HOST_SIDECAR_PORT)) return;
    }
    throw new Error('Host sidecar failed to start');
  }
};

export const initTerminalSidecars = async () => {
  await startHostSidecar();
  ensureDockerImage();
  const users = await getUsers();
  for (const user of users) {
    const homeDir = getHomeDir(user.email);
    mkdirSync(dirname(homeDir), { recursive: true });
    mkdirSync(homeDir, { recursive: true });
    const shellUsername = toShellUsername(user.username ?? '', user.email);
    const contextFile = generateContainerContext(user.email);
    const settingsFile = generateClaudeSettings(user.email, shellUsername);
    try {
      await ensureDockerContainer(user.email, user.id, homeDir, shellUsername, contextFile, settingsFile);
      console.log(`[terminal] sidecar ready for ${user.email}`);
    } catch (err) {
      console.error(`[terminal] failed to start sidecar for ${user.email}:`, err);
    }
  }
};

const containerShell: ShellInfo = { command: '/bin/zsh', args: ['-d', '-i'], name: 'zsh' };

const resolveCwd = (home: string, cwd?: string) => {
  if (!cwd || cwd === '~') return home;
  if (cwd.startsWith('~/')) return join(home, cwd.slice(2));
  if (cwd.startsWith('/')) return join(home, cwd.slice(1));
  return home;
};

export const terminalWebsocket = {
  async open(ws: ServerWebSocket<WSData>) {
    const { email, username, role, sandboxed } = ws.data;

    if (!sandboxed && role !== 'Super Admin') {
      sendOutput(ws, '\r\n[Permission denied] Host terminal requires Super Admin role.\r\n');
      return;
    }

    if (!sandboxed) {
      const session: BridgeSession = { client: ws, sidecar: null, dockerId: '', port: HOST_SIDECAR_PORT, pendingMessages: [] };
      sessions.set(ws, session);

      let sidecar: WebSocket | null = null;
      try {
        sidecar = await connectSidecar(HOST_SIDECAR_PORT);
      } catch (err) {
        const message = err instanceof Error ? err.message : 'Failed to connect host sidecar';
        sendOutput(ws, `\r\n[Terminal error] ${message}\r\n`);
        sessions.delete(ws);
        return;
      }

      session.sidecar = sidecar;

      sidecar.addEventListener('message', (ev) => {
        try {
          if (typeof ev.data === 'string') {
            ws.send(ev.data);
          } else {
            ws.send(new TextDecoder().decode(ev.data));
          }
        } catch {
          // ws already closed
        }
      });

      sidecar.send(
        JSON.stringify({
          type: 'init',
          host: true,
          sessionId: ws.data.sessionId ?? `host-${ws.data.userId}`,
          shell: { command: process.env.SHELL ?? '/bin/zsh', args: ['-i'] },
          cwd: resolveCwd(process.env.HOME!, ws.data.cwd),
          homeDir: process.env.HOME,
          userLabel: email,
          cols: ws.data.cols,
          rows: ws.data.rows,
        }),
      );

      for (const msg of session.pendingMessages) sidecar.send(msg);
      session.pendingMessages = [];
      return;
    }

    const cwd = getHomeDir(email);
    const userRoot = dirname(cwd);
    mkdirSync(userRoot, { recursive: true });
    mkdirSync(cwd, { recursive: true });

    const session: BridgeSession = { client: ws, sidecar: null, dockerId: '', port: 0, pendingMessages: [] };
    sessions.set(ws, session);

    let sidecar: WebSocket | null = null;
    let info: ContainerInfo | undefined;
    try {
      info = await ensureDockerContainer(email, ws.data.userId, cwd, username, generateContainerContext(email), generateClaudeSettings(email, username));
      sidecar = await connectSidecar(info.port);
    } catch (err) {
      const message = err instanceof Error ? err.message : 'Failed to connect terminal sidecar';
      console.error(`[terminal] sidecar connection failed for ${email}:`, message);
      sendOutput(ws, `\r\n[Terminal error] ${message}\r\n`);
      if (info) {
        const logs = readDockerLogs(info.dockerId);
        if (logs) {
          sendOutput(ws, `\r\n[Docker logs]\r\n${logs}\r\n`);
        }
      }
      sendOutput(ws, '\r\n[Process exited]\r\n');
      if (info) stopDockerSidecar(info.dockerId);
      sessions.delete(ws);
      return;
    }

    session.sidecar = sidecar;
    session.dockerId = info.dockerId;
    session.port = info.port;

    sidecar.addEventListener('message', (ev) => {
      try {
        if (typeof ev.data === 'string') {
          ws.send(ev.data);
        } else {
          ws.send(new TextDecoder().decode(ev.data));
        }
      } catch {
        // ws already closed
      }
    });

    const containerHome = `/home/${username}`;
    sidecar.send(
      JSON.stringify({
        type: 'init',
        sessionId: ws.data.sessionId ?? `default-${ws.data.userId}`,
        shell: containerShell,
        cwd: resolveCwd(containerHome, ws.data.cwd),
        homeDir: containerHome,
        userLabel: email,
        cols: ws.data.cols,
        rows: ws.data.rows,
      }),
    );

    for (const msg of session.pendingMessages) sidecar.send(msg);
    session.pendingMessages = [];
  },

  message(ws: ServerWebSocket<WSData>, raw: string | Buffer) {
    const session = sessions.get(ws);
    if (!session) return;

    const payload = typeof raw === 'string' ? raw : raw.toString();

    if (!session.sidecar || session.sidecar.readyState !== WebSocket.OPEN) {
      session.pendingMessages.push(payload);
      return;
    }

    try {
      session.sidecar.send(payload);
    } catch {
      // ignore
    }
  },

  close(ws: ServerWebSocket<WSData>) {
    const session = sessions.get(ws);
    if (session?.sidecar) {
      try {
        session.sidecar.close();
      } catch {
        // ignore
      }
    }
    sessions.delete(ws);
  },

  drain() {},
};

export const broadcastPanelRefresh = (email: string) => {
  const msg = JSON.stringify({ type: 'panel-refresh' });
  for (const [ws, session] of sessions) {
    if (ws.data.email === email && session.sidecar) {
      try { ws.send(msg); } catch { /* ignore */ }
    }
  }
};

export const stopAllContainers = async () => {
  // Stop host sidecar
  if (hostSidecarProcess) {
    hostSidecarProcess.kill();
    await hostSidecarProcess.exited.catch(() => {});
    hostSidecarProcess = null;
    console.log('[terminal] host sidecar stopped');
  }

  // Stop all Docker containers
  const map = await loadContainerMap();
  const entries = Object.entries(map);
  if (entries.length === 0) return;

  const dockerPath = Bun.which('docker') ?? 'docker';
  for (const [email, info] of entries) {
    try {
      Bun.spawnSync({ cmd: [dockerPath, 'stop', '-t', '2', info.dockerId], stdout: 'ignore', stderr: 'ignore' });
      console.log(`[terminal] stopped container ${info.dockerId} (${email})`);
    } catch {
      // ignore
    }
  }
};