import type { ServerWebSocket } from 'bun';
import { join, isAbsolute } from 'node:path';
import { mkdirSync, writeFileSync, chmodSync, rmSync, readdirSync, readFileSync } from 'node:fs';
import { getTaskByDirName } from './task-files';
import { getOwnerHomeDir } from '../../data-path';

type WSData = {
  userId: number;
  email: string;
  username: string;
};

type RunMessage = {
  type: 'run';
  taskDirName: string;
  inputs: Record<string, string>;
  cwd?: string;
};

type StopMessage = {
  type: 'stop';
};

type ClientMessage = RunMessage | StopMessage;

type OutMessage =
  | { type: 'started'; taskName: string }
  | { type: 'stdout'; data: string }
  | { type: 'stderr'; data: string }
  | { type: 'exit'; code: number }
  | { type: 'error'; message: string };

// Active processes per WebSocket
const activeProcs = new WeakMap<ServerWebSocket<WSData>, { proc: ReturnType<typeof Bun.spawn>; kill: () => void }>();

// Every descendant PID of `root`, from a single /proc walk (Linux). Bun's proc.kill() signals only
// the direct child (bash), so a running grandchild like ffmpeg is left orphaned and keeps going.
function descendantPids(root: number): number[] {
  const childrenOf = new Map<number, number[]>();
  let entries: string[];
  try {
    entries = readdirSync('/proc');
  } catch {
    return [];
  }
  for (const entry of entries) {
    const pid = Number(entry);
    if (!Number.isInteger(pid)) continue;
    try {
      const stat = readFileSync(`/proc/${pid}/stat`, 'utf8');
      // "pid (comm) state ppid …" — comm may contain spaces/parens, so parse after the last ')'.
      const fields = stat.slice(stat.lastIndexOf(')') + 2).split(' ');
      const ppid = Number(fields[1]);
      if (!Number.isInteger(ppid)) continue;
      const list = childrenOf.get(ppid);
      if (list) list.push(pid);
      else childrenOf.set(ppid, [pid]);
    } catch {
      // process vanished mid-scan
    }
  }
  const out: number[] = [];
  const stack = [root];
  while (stack.length > 0) {
    const parent = stack.pop()!;
    for (const child of childrenOf.get(parent) ?? []) {
      out.push(child);
      stack.push(child);
    }
  }
  return out;
}

// Terminate a process and its whole subtree: SIGTERM everything, then SIGKILL stragglers after a
// grace period. Killing the root (bash) too stops it from spawning the next file mid-batch.
function killTree(root: number) {
  const pids = [root, ...descendantPids(root)];
  for (const pid of pids) {
    try {
      process.kill(pid, 'SIGTERM');
    } catch {
      /* already gone */
    }
  }
  setTimeout(() => {
    for (const pid of pids) {
      try {
        process.kill(pid, 'SIGKILL');
      } catch {
        /* gone */
      }
    }
  }, 2000);
}

import { tmpdir } from 'node:os';

function send(ws: ServerWebSocket<WSData>, msg: OutMessage) {
  if (ws.readyState === 1) ws.send(JSON.stringify(msg));
}

function getRunner(language: string): string[] {
  switch (language) {
    case 'bash':
      return ['bash'];
    case 'python':
      return ['python3'];
    case 'typescript':
      return ['bun', 'run'];
    case 'javascript':
      return ['node'];
    default:
      return ['bash'];
  }
}

function getFileName(language: string): string {
  switch (language) {
    case 'bash':
      return 'run.sh';
    case 'python':
      return 'run.py';
    case 'typescript':
      return 'index.ts';
    case 'javascript':
      return 'index.js';
    default:
      return 'run.sh';
  }
}

// Write implementation to a temp file for execution, cleaned up after
function materializeScript(language: string, implementation: string): string {
  const dir = join(tmpdir(), `officer-task-${Date.now()}-${Math.random().toString(36).slice(2)}`);
  mkdirSync(dir, { recursive: true });

  const fileName = getFileName(language);
  const filePath = join(dir, fileName);

  writeFileSync(filePath, implementation);
  chmodSync(filePath, 0o755);

  return filePath;
}

function buildInputEnv(inputs: Record<string, string>): Record<string, string> {
  const env: Record<string, string> = {};
  for (const [key, value] of Object.entries(inputs)) {
    env[`INPUT_${key.toUpperCase()}`] = value;
  }
  return env;
}

function buildArgs(inputs: Record<string, string>, argsOrder?: string[] | null): string[] {
  if (!argsOrder || argsOrder.length === 0) return [];
  return argsOrder.map((name) => inputs[name] ?? '');
}

async function handleRun(ws: ServerWebSocket<WSData>, msg: RunMessage) {
  const { email } = ws.data;

  // Resolve task from the file-backed store
  const task = await getTaskByDirName(msg.taskDirName);
  if (!task) {
    send(ws, { type: 'error', message: `Task not found: ${msg.taskDirName}` });
    return;
  }

  if (task.mode !== 'script') {
    send(ws, { type: 'error', message: 'Task is not a script-mode task' });
    return;
  }

  if (!task.implementation) {
    send(ws, { type: 'error', message: `Task ${msg.taskDirName} has no implementation` });
    return;
  }

  const language = task.language ?? 'bash';

  // Write script to temp dir for execution
  const scriptPath = materializeScript(language, task.implementation);

  // Build env vars from inputs
  const inputEnv = buildInputEnv(msg.inputs);

  // Build positional args
  const positionalArgs = buildArgs(msg.inputs, task.args);

  // Build the command
  const runner = getRunner(language);
  const cmd = [...runner, scriptPath, ...positionalArgs];

  // msg.cwd arrives from the file browser relative to the user's home; Bun.spawn needs it absolute
  // (a missing cwd surfaces as ENOENT naming the binary, not the directory)
  const homeDir = getOwnerHomeDir(email);
  const cwd = msg.cwd ? (isAbsolute(msg.cwd) ? msg.cwd : join(homeDir, msg.cwd)) : homeDir;

  const spawnCmd = cmd;
  const spawnEnv = { ...(process.env as Record<string, string>), ...inputEnv };
  const spawnCwd = cwd;

  const cleanup = () => {
    try {
      rmSync(join(scriptPath, '..'), { recursive: true, force: true });
    } catch {
      /* best effort */
    }
  };

  send(ws, { type: 'started', taskName: task.name });

  try {
    const proc = Bun.spawn(spawnCmd, {
      cwd: spawnCwd,
      env: spawnEnv,
      stdout: 'pipe',
      stderr: 'pipe',
    });

    activeProcs.set(ws, {
      proc,
      kill: () => {
        try {
          killTree(proc.pid);
        } catch {
          /* already dead */
        }
      },
    });

    // Keep the WebSocket alive during long silent phases (e.g. ffmpeg's faststart pass rewrites a huge
    // file for minutes with no output). Bun's default 120s idle timeout would otherwise close the
    // socket → close(ws) → killTree kills the task mid-run. A ping resets the idle timer.
    const keepAlive = setInterval(() => {
      try {
        ws.ping();
      } catch {
        /* socket gone */
      }
    }, 30_000);

    const stdoutReader = proc.stdout.getReader();
    const stderrReader = proc.stderr.getReader();
    const decoder = new TextDecoder();

    const readStream = async (reader: ReadableStreamDefaultReader<Uint8Array>, type: 'stdout' | 'stderr') => {
      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
          send(ws, { type, data: decoder.decode(value) });
        }
      } catch {
        // stream closed
      }
    };

    const [, , exitCode] = await Promise.all([
      readStream(stdoutReader, 'stdout'),
      readStream(stderrReader, 'stderr'),
      proc.exited,
    ]);

    clearInterval(keepAlive);
    activeProcs.delete(ws);
    cleanup();
    send(ws, { type: 'exit', code: exitCode });
  } catch (err) {
    activeProcs.delete(ws);
    cleanup();
    send(ws, { type: 'error', message: `Failed to spawn: ${err instanceof Error ? err.message : String(err)}` });
  }
}

export function open(_ws: ServerWebSocket<WSData>) {
  // nothing to do
}

export function message(ws: ServerWebSocket<WSData>, raw: string | Buffer) {
  const data = typeof raw === 'string' ? raw : raw.toString();

  try {
    const msg = JSON.parse(data) as ClientMessage;

    if (msg.type === 'run') {
      handleRun(ws, msg);
    } else if (msg.type === 'stop') {
      const active = activeProcs.get(ws);
      if (active) {
        active.kill();
        activeProcs.delete(ws);
        send(ws, { type: 'exit', code: -1 });
      }
    }
  } catch {
    send(ws, { type: 'error', message: 'Failed to parse message' });
  }
}

export function close(ws: ServerWebSocket<WSData>) {
  const active = activeProcs.get(ws);
  if (active) {
    active.kill();
    activeProcs.delete(ws);
  }
}

export const taskRunnerWebsocket = {
  open,
  message,
  close,
  drain() {},
};