import { BACKGROUND_CONTEXT } from "@earendil-works/chord/context";
import { Type } from "@earendil-works/pi-ai";
import {
  type Channel,
  type ChannelMessage,
  defineExtension,
  type Environment,
  type Pty,
} from "@berry/api";
import { Ghostty, type Terminal } from "./ghostty.ts";
import { COLORS } from "./palette.ts";
import { shell } from "./shell.ts";
import { type ClientMessage, OUTPUT, type ServerMessage, SNAPSHOT, tagged } from "./protocol.ts";

const WASM = new URL("./ghostty-vt.wasm", import.meta.url);
const SIZE = { min: 1, max: 1000 };

/** A conversation's shell, its terminal kept here too, so whoever opens the view later sees the screen so far. */
type Session = {
  readonly pty: Pty;
  readonly input: WritableStreamDefaultWriter<Uint8Array>;
  readonly terminal: Terminal;
  /** Once its output has ended, which its pty sees to shortly after the shell exits. */
  readonly pumped: Promise<void>;
  exit: { code: number | null } | null;
};

export default defineExtension({
  name: "terminal",
  description: "libghostty-backed terminal on the target machine",
  alwaysOn: true,
  initial: () => null,
  setup(ext) {
    let loaded: Promise<Ghostty> | undefined;
    const ghostty = () =>
      loaded ??= Deno.readFile(WASM).then(Ghostty.load, (error) => {
        loaded = undefined;
        throw error instanceof Deno.errors.NotFound
          ? new Error("extensions/terminal/ghostty-vt.wasm is missing: run `deno task ghostty`")
          : error;
      });
    /** One per conversation, until its shell exits and another is asked for. */
    const sessions = new Map<string, Session>();
    const starting = new Set<string>();
    /** Whoever has a conversation's terminal open. */
    const viewers = new Map<string, Set<Channel>>();
    const broadcast = (tid: string, message: ChannelMessage | ServerMessage) => {
      for (const viewer of viewers.get(tid) ?? []) viewer.send(message);
    };

    async function start(tid: string, env: Environment, cols: number, rows: number) {
      const terminal = (await ghostty()).terminal(cols, rows);
      terminal.colors(COLORS);
      let pty: Pty;
      try {
        pty = await env.pty(await shell(env), {
          cols,
          rows,
          env: { TERM: "xterm-256color", COLORTERM: "truecolor" },
        }, BACKGROUND_CONTEXT);
      } catch (error) {
        terminal.free();
        throw error;
      }
      sessions.get(tid)?.terminal.free();
      const pumped = (async () => {
        for await (const chunk of pty.readable) {
          terminal.write(chunk);
          broadcast(tid, tagged(OUTPUT, chunk));
        }
      })().catch(() => {});
      const input = pty.writable.getWriter();
      const session: Session = { pty, input, terminal, pumped, exit: null };
      sessions.set(tid, session);
      // the program's queries, such as where the cursor is, are answered once, here, rather than by every viewer
      terminal.answer((data) => void input.write(data).catch(() => {}));
      broadcast(tid, tagged(SNAPSHOT, terminal.snapshot()));

      // only once its output has ended, so nothing of this shell reaches the next one's screen
      void Promise.all([pty.exited, pumped]).then(([code]) => {
        session.exit = { code };
        if (sessions.get(tid) === session) broadcast(tid, { t: "exit", code });
      });
    }

    ext.onDispose(async () => {
      const ending = [...sessions.values()];
      sessions.clear();
      await Promise.all(ending.map(async (session) => {
        session.pty.kill();
        await session.pumped;
        session.terminal.free();
      }));
    });

    ext.channel({
      name: "session",
      parameters: Type.Object({}),
      async open(_, channel, ctx) {
        // a missing module is the reason to give the view, rather than a terminal that never draws
        await ghostty();
        const { tid } = ctx;
        const watching = viewers.get(tid) ?? viewers.set(tid, new Set()).get(tid)!;
        watching.add(channel);
        ctx.signal.addEventListener("abort", () => watching.delete(channel));
        const session = sessions.get(tid);
        if (session === undefined) channel.send({ t: "none" } satisfies ServerMessage);
        else {
          channel.send(tagged(SNAPSHOT, session.terminal.snapshot()));
          if (session.exit) channel.send({ t: "exit", ...session.exit } satisfies ServerMessage);
        }

        void (async () => {
          for await (const message of channel.messages) {
            const session = sessions.get(tid);
            const running = session && !session.exit ? session : undefined;
            if (message instanceof Uint8Array) {
              await running?.input.write(message).catch(() => {});
              continue;
            }
            const { t, cols, rows } = message as ClientMessage;
            const sized = [cols, rows].every((n) =>
              Number.isInteger(n) && n >= SIZE.min && n <= SIZE.max
            );
            if (!sized) continue;
            if (t === "start" && !running && !starting.has(tid)) {
              starting.add(tid);
              start(tid, ctx.env, cols, rows)
                .catch((error) =>
                  channel.send(
                    {
                      t: "error",
                      message: error instanceof Error ? error.message : String(error),
                    } satisfies ServerMessage,
                  )
                )
                .finally(() => starting.delete(tid));
            } else if (t === "resize" && running) {
              running.pty.resize(cols, rows);
              running.terminal.resize(cols, rows);
              broadcast(tid, { t: "size", cols, rows });
            }
          }
        })();
      },
    });
  },
});