import type { Dirent, ExecOptions, GuestWrite, HostFs, HostStat, Kernel, PtyExecOptions, PtyProcess, RunningProcess, Signal } from "@nibling/kernel";
import { fromRecord, owned, type Connect, type FromSandbox, type FsCall, type FsOp, type Subscription, type ToSandbox } from "./protocol.ts";

/** The part of a kernel a coordinator drives. A local `Kernel` satisfies it too. */
export type SandboxKernel = Pick<
  Kernel,
  "exec" | "fs" | "applyFiles" | "putFile" | "observeWrites" | "kill" | "onListening" | "onUnlistening" | "onProcesses" | "close"
>;

export interface ConnectOptions {
  /** The sandbox's origin, such as `https://<id>.nibling.run`. Its page boots the kernel. */
  origin: string;
  /** Where the hidden iframe goes. Default: `document.body`. */
  container?: Element;
}

/** How long the sandbox page has to accept the handshake. Booting the kernel afterwards is not bounded. */
const HANDSHAKE_TIMEOUT_MS = 15_000;

/**
 * Boots a kernel on `origin` in a hidden iframe and resolves with a facade over a MessagePort once it
 * is ready. Rejects when the page never accepts the handshake (it refused to be framed, say) or the
 * kernel fails to boot. `close()` removes the iframe, which ends the kernel.
 */
export function connectSandbox({ origin, container = document.body }: ConnectOptions): Promise<SandboxKernel> {
  const target = new URL(origin).origin;
  const iframe = document.createElement("iframe");
  iframe.hidden = true;
  iframe.title = "Sandbox";
  iframe.allow = "cross-origin-isolated";
  iframe.src = `${target}/`;
  const channel = new MessageChannel();
  const port = channel.port1;
  return new Promise((resolve, reject) => {
    const fail = (error: Error) => {
      clearTimeout(timer);
      port.close();
      iframe.remove();
      reject(error);
    };
    const timer = setTimeout(() => fail(new Error(`The sandbox at ${target} did not answer.`)), HANDSHAKE_TIMEOUT_MS);
    port.onmessage = (event: MessageEvent<FromSandbox>) => {
      const message = event.data;
      if (message.t === "accepted") clearTimeout(timer);
      else if (message.t === "bootFailed") fail(fromRecord(message.error));
      else if (message.t === "ready") resolve(facade(port, iframe, target));
    };
    iframe.addEventListener("load", () => {
      iframe.contentWindow?.postMessage({ t: "connect" } satisfies Connect, target, [channel.port2]);
    }, { once: true });
    container.append(iframe);
  });
}

type ProcessMessage = Extract<FromSandbox, { proc: number }>;
type TopicMessage = Extract<FromSandbox, { sub: number }>;
type Answer = Extract<FromSandbox, { id: number }>;
type Request = Extract<ToSandbox, { id: number }>;

const isWhole = (value: unknown): value is number => Number.isSafeInteger(value) && (value as number) >= 0;
const isSignal = (value: unknown): value is Signal => typeof value === "string" && /^SIG[A-Z0-9]{1,8}$/.test(value);
const isStream = (value: unknown): value is "stdout" | "stderr" => value === "stdout" || value === "stderr";
const fields = (value: unknown) => (typeof value === "object" && value !== null ? value : {}) as Record<string, unknown>;
const under = (path: string, prefixes: readonly string[]) =>
  prefixes.some((p) => path === p || path.startsWith(p.endsWith("/") ? p : `${p}/`));

/** The GuestWrite in `message`, or undefined when it is malformed or names a path the listener did not ask for. */
function guestWrite(message: Extract<FromSandbox, { t: "write" }>, prefixes: readonly string[]): GuestWrite | undefined {
  const { path, pid, argv, data, gen, readGen } = message;
  if (typeof path !== "string" || !under(path, prefixes)) return undefined;
  if (!isWhole(pid) || !Array.isArray(argv) || !argv.every((arg) => typeof arg === "string")) return undefined;
  if (!(data instanceof Uint8Array) || !isWhole(gen) || !(readGen === null || isWhole(readGen))) return undefined;
  return { path, pid, argv, data, gen, readGen };
}

/** Reads what an answer carries: the value, or undefined when it has the wrong shape. */
type Accept<T> = (value: unknown) => { value: T } | undefined;

const nothing: Accept<void> = () => ({ value: undefined });
const bytes: Accept<Uint8Array> = (value) => (value instanceof Uint8Array ? { value } : undefined);
const text: Accept<string> = (value) => (typeof value === "string" ? { value } : undefined);

function dirent(value: unknown): Dirent | undefined {
  const { name, kind } = fields(value);
  if (typeof name !== "string" || name === "" || name.includes("/")) return undefined;
  return kind === "file" || kind === "dir" || kind === "symlink" ? { name, kind } : undefined;
}

const dirents: Accept<Dirent[]> = (value) => {
  if (!Array.isArray(value)) return undefined;
  const entries = value.map(dirent);
  return entries.every((entry) => entry !== undefined) ? { value: entries } : undefined;
};

const stat: Accept<HostStat> = (value) => {
  const { kind, size, mtimeMs, mode } = fields(value);
  if (kind !== "file" && kind !== "dir") return undefined;
  return isWhole(size) && typeof mtimeMs === "number" && Number.isFinite(mtimeMs) && isWhole(mode) ? { value: { kind, size, mtimeMs, mode } } : undefined;
};

/** The kernel behind `port`, whose other end is the page on `origin` that `frame` holds. */
export function facade(port: MessagePort, frame: Pick<HTMLIFrameElement, "remove">, origin: string): SandboxKernel {
  let next = 1;
  const requests = new Map<number, { settle: (answer: Answer) => void; reject: (error: Error) => void }>();
  const procs = new Map<number, (message: ProcessMessage) => void>();
  const subs = new Map<number, (message: TopicMessage) => void>();
  const send = (message: ToSandbox, transfer: Transferable[] = []) => port.postMessage(message, transfer);

  port.onmessage = (event: MessageEvent<FromSandbox>) => {
    const message = event.data;
    if (typeof message !== "object" || message === null) return;
    if ("id" in message) {
      const request = requests.get(message.id);
      requests.delete(message.id);
      request?.settle(message);
    } else if ("proc" in message) procs.get(message.proc)?.(message);
    else if ("sub" in message) subs.get(message.sub)?.(message);
  };

  /** Sends a request and resolves with what `accept` reads from its answer; an answer it rejects is an error. */
  const call = <T>(message: Request, accept: (answer: Answer) => { value: T } | undefined, transfer: Transferable[] = []) =>
    new Promise<T>((resolve, reject) => {
      requests.set(message.id, {
        settle(answer) {
          if (answer.t === "failed") return reject(fromRecord(answer.error));
          const accepted = accept(answer);
          if (accepted) resolve(accepted.value);
          else reject(new Error(`The sandbox answered ${message.t === "fs" ? message.op : message.t} with ${answer.t}.`));
        },
        reject,
      });
      send(message, transfer);
    });

  const subscribe = (subscription: Subscription, listener: (message: TopicMessage) => void) => {
    const sub = next++;
    subs.set(sub, listener);
    send({ t: "subscribe", sub, ...subscription });
    return () => {
      if (subs.delete(sub)) send({ t: "unsubscribe", sub });
    };
  };

  const fsCall = <K extends FsOp>(op: K, args: Parameters<HostFs[K]>, accept: Accept<Awaited<ReturnType<HostFs[K]>>>, transfer: Transferable[] = []) =>
    call({ t: "fs", id: next++, op, args } as FsCall, (answer) => (answer.t === "fsDone" ? accept(answer.value) : undefined), transfer);

  const fs: HostFs = {
    readFile: (path) => fsCall("readFile", [path], bytes),
    writeFile(path, data) {
      const copy = typeof data === "string" ? data : owned(data);
      return fsCall("writeFile", [path, copy], nothing, typeof copy === "string" ? [] : [copy.buffer]);
    },
    mkdir: (path, opts) => fsCall("mkdir", [path, opts], nothing),
    readdir: (path) => fsCall("readdir", [path], dirents),
    stat: (path) => fsCall("stat", [path], stat),
    unlink: (path) => fsCall("unlink", [path], nothing),
    rm: (path, opts) => fsCall("rm", [path, opts], nothing),
    rename: (from, to) => fsCall("rename", [from, to], nothing),
    symlink: (target, path) => fsCall("symlink", [target, path], nothing),
    readlink: (path) => fsCall("readlink", [path], text),
  };

  function exec(entry: string, options: PtyExecOptions): PtyProcess;
  function exec(entry: string, options?: ExecOptions): RunningProcess;
  function exec(entry: string, options: ExecOptions | PtyExecOptions = {}): RunningProcess | PtyProcess {
    const proc = next++;
    const process = remoteProcess(proc, send, Boolean((options as Partial<PtyExecOptions>).pty));
    procs.set(proc, process.receive);
    void process.done.then(() => procs.delete(proc));
    send({ t: "exec", proc, entry, options });
    return process.handle;
  }

  return {
    fs,
    exec,
    applyFiles(files) {
      return call({ t: "applyFiles", id: next++, files }, (answer) => (answer.t === "done" ? { value: undefined } : undefined));
    },
    putFile(path, data) {
      const bytes = typeof data === "string" ? data : owned(data);
      return call(
        { t: "putFile", id: next++, path, data: bytes },
        (answer) => (answer.t === "written" && isWhole(answer.gen) ? { value: answer.gen } : undefined),
        typeof bytes === "string" ? [] : [bytes.buffer],
      );
    },
    kill(pid, signal) {
      send({ t: "kill", pid, signal });
    },
    onListening(listener) {
      return subscribe({ topic: "listening" }, (message) => {
        // Anything on the sandbox's origin can post on this port, so a URL that leaves it is dropped.
        if (message.t === "listening" && URL.parse(message.url)?.origin === origin) listener(message.port, message.url);
      });
    },
    onUnlistening(listener) {
      return subscribe({ topic: "unlistening" }, (message) => {
        if (message.t === "unlistening") listener(message.port);
      });
    },
    onProcesses(listener) {
      return subscribe({ topic: "processes" }, (message) => {
        if (message.t === "processes") listener(message.table);
      });
    },
    observeWrites(prefixes, listener) {
      const asked = [...prefixes];
      return subscribe({ topic: "writes", prefixes: asked }, (message) => {
        const write = message.t === "write" ? guestWrite(message, asked) : undefined;
        if (write) listener(write);
      });
    },
    close() {
      port.close();
      frame.remove();
      for (const request of requests.values()) request.reject(new Error("The sandbox closed."));
      requests.clear();
      return Promise.resolve();
    },
  };
}

function deferred<T>() {
  let resolve!: (value: T) => void;
  let reject!: (error: Error) => void;
  const promise = new Promise<T>((res, rej) => ((resolve = res), (reject = rej)));
  return { promise, resolve, reject };
}

/**
 * A process whose other end is in the sandbox, fed by `receive`: a RunningProcess, or a PtyProcess
 * when `pty`. `done` settles once it exited and both streams ended. A message that repeats an end or
 * an exit, or has the wrong shape, is dropped.
 */
function remoteProcess(proc: number, send: (message: ToSandbox, transfer?: Transferable[]) => void, pty: boolean) {
  const streams = {} as Record<"stdout" | "stderr", ReadableStreamDefaultController<Uint8Array>>;
  const stdout = new ReadableStream<Uint8Array>({ start: (controller) => void (streams.stdout = controller) });
  const stderr = new ReadableStream<Uint8Array>({ start: (controller) => void (streams.stderr = controller) });
  const pid = deferred<number>();
  const exited = deferred<number>();
  const signal = deferred<Signal | null>();
  const done = deferred<void>();
  pid.promise.catch(() => {});
  signal.promise.catch(() => {});
  const open = new Set(["stdout", "stderr", "exit"]);
  const settle = (part: string) => {
    open.delete(part);
    if (open.size === 0) done.resolve();
  };

  const receive = (message: ProcessMessage) => {
    switch (message.t) {
      case "pid":
        if (isWhole(message.pid)) pid.resolve(message.pid);
        return;
      case "output":
        if (isStream(message.stream) && open.has(message.stream) && message.bytes instanceof Uint8Array) streams[message.stream].enqueue(message.bytes);
        return;
      case "end":
        if (!isStream(message.stream) || !open.has(message.stream)) return;
        streams[message.stream].close();
        return settle(message.stream);
      case "exit":
        if (!open.has("exit") || !isWhole(message.code) || !(message.signal === null || isSignal(message.signal))) return;
        exited.resolve(message.code);
        signal.resolve(message.signal);
        return settle("exit");
      case "spawnFailed": {
        if (open.size < 3) return;
        const error = fromRecord(message.error);
        for (const pending of [pid, exited, signal]) pending.reject(error);
        streams.stdout.error(error);
        streams.stderr.error(error);
        open.clear();
        return done.resolve();
      }
    }
  };

  const base = {
    pid: pid.promise,
    stdout,
    stderr,
    exited: exited.promise,
    signal: signal.promise,
    kill: (sig?: Signal) => send({ t: "signal", proc, signal: sig }),
  };
  const handle: RunningProcess | PtyProcess = pty
    ? {
        ...base,
        pty: {
          write(data) {
            const copy = typeof data === "string" ? data : owned(data);
            send({ t: "ptyWrite", proc, data: copy }, typeof copy === "string" ? [] : [copy.buffer]);
          },
          resize: (cols, rows) => send({ t: "ptyResize", proc, cols, rows }),
        },
      }
    : {
        ...base,
        stdin: new WritableStream<Uint8Array>({
          write(chunk) {
            const copy = owned(chunk);
            send({ t: "stdin", proc, bytes: copy }, [copy.buffer]);
          },
          close: () => send({ t: "stdinEnd", proc }),
          abort: () => send({ t: "stdinEnd", proc }),
        }),
      };
  return { handle, receive, done: done.promise };
}