import type { Environment, Process } from "@berry/api";
import { BACKGROUND_CONTEXT, withAbortSignal } from "@earendil-works/chord/context";
import {
  type JsonRpcMessage,
  McpConnectionClosedError,
  type McpTransport,
  parseJsonRpcMessage,
} from "@earendil-works/pi-mcp";
import type { Server } from "./config.ts";
import { lines } from "./lines.ts";

async function settleWithin(promise: Promise<unknown>, ms: number) {
  let timer: ReturnType<typeof setTimeout> | undefined;
  try {
    await Promise.race([
      promise.catch(() => {}),
      new Promise<void>((resolve) => timer = setTimeout(resolve, ms)),
    ]);
  } finally {
    clearTimeout(timer);
  }
}

/** MCP framing over a process in the conversation's environment. */
export class EnvironmentStdio implements McpTransport {
  private readonly messages = new Set<(message: JsonRpcMessage) => void>();
  private readonly errors = new Set<(error: Error) => void>();
  private readonly closes = new Set<() => void>();
  private readonly lifetime = new AbortController();
  private process?: Process;
  private writer?: WritableStreamDefaultWriter<Uint8Array>;
  private readers: ReadableStreamDefaultReader<Uint8Array>[] = [];
  private closed = false;
  private closing?: Promise<void>;
  private stderr = "";
  private exitCode?: number | null;
  private reason?: Error;
  private reading?: Promise<unknown>;

  constructor(
    private readonly env: Environment,
    private readonly config: Extract<Server, { command: string }>,
    signal: AbortSignal,
  ) {
    signal.addEventListener("abort", () => void this.close(), { once: true });
    if (signal.aborted) this.closed = true;
  }

  get failure(): Error | undefined {
    if (!this.reason) return undefined;
    return new McpConnectionClosedError(
      [
        this.reason.message,
        this.exitCode === undefined
          ? ""
          : this.exitCode === null
          ? "terminated by signal"
          : `exit code ${this.exitCode}`,
        this.stderr.trim(),
      ].filter(Boolean).join("\n"),
    );
  }

  onMessage(listener: (message: JsonRpcMessage) => void) {
    this.messages.add(listener);
    return () => void this.messages.delete(listener);
  }

  onError(listener: (error: Error) => void) {
    this.errors.add(listener);
    return () => void this.errors.delete(listener);
  }

  onClose(listener: () => void) {
    this.closes.add(listener);
    return () => void this.closes.delete(listener);
  }

  private finish(error: unknown) {
    if (this.closed) return;
    this.reason = error instanceof Error ? error : new Error(String(error));
    this.closed = true;
    for (const listener of this.errors) listener(this.reason);
    for (const listener of this.closes) listener();
    this.lifetime.abort();
    void this.close();
  }

  async start() {
    if (this.closed) throw new McpConnectionClosedError();
    if (this.process) throw new Error("MCP transport already started");
    const child = await this.env.spawn(
      [this.config.command, ...this.config.args],
      { cwd: this.config.cwd, env: this.config.env, group: true },
      withAbortSignal(this.lifetime.signal, BACKGROUND_CONTEXT),
    );
    this.process = child;
    if (this.closed) {
      child.kill("SIGKILL");
      await settleWithin(
        Promise.allSettled([
          child.exited,
          child.stdout.cancel(),
          child.stderr.cancel(),
          child.stdin.abort(),
        ]),
        1000,
      );
      throw new McpConnectionClosedError();
    }
    this.writer = child.stdin.getWriter();
    this.readers = [child.stdout.getReader(), child.stderr.getReader()];
    const errors = (async () => {
      const reader = this.readers[1];
      const decoder = new TextDecoder();
      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
          this.stderr = (this.stderr + decoder.decode(value, { stream: true })).slice(-4000);
        }
        this.stderr = (this.stderr + decoder.decode()).slice(-4000);
      } finally {
        reader.releaseLock();
      }
    })();
    const messages = (async () => {
      const reader = this.readers[0];
      try {
        for await (const line of lines(reader)) {
          let message: JsonRpcMessage;
          try {
            message = parseJsonRpcMessage(JSON.parse(line));
          } catch (error) {
            const normalized = error instanceof Error ? error : new Error(String(error));
            for (const listener of this.errors) listener(normalized);
            continue;
          }
          for (const listener of this.messages) listener(message);
        }
      } finally {
        reader.releaseLock();
      }
      this.finish(new McpConnectionClosedError("MCP server closed stdout"));
    })();
    this.reading = Promise.allSettled([errors, messages]);
    void messages.catch((error) => this.finish(error));
    void errors.catch((error) => this.finish(error));
    void child.exited.then(
      (code) => this.exitCode = code,
      (error) => this.finish(error),
    );
  }

  async send(message: JsonRpcMessage) {
    if (this.closed || !this.writer) throw new McpConnectionClosedError();
    await this.writer.write(new TextEncoder().encode(JSON.stringify(message) + "\n"));
  }

  close(): Promise<void> {
    return this.closing ??= (async () => {
      if (!this.closed) {
        this.closed = true;
        for (const listener of this.closes) listener();
      }
      const inputClosed = this.writer?.close().catch(() => {});
      const child = this.process;
      if (!child) {
        this.lifetime.abort();
        return;
      }
      await settleWithin(child.exited, 500);
      this.lifetime.abort();
      child.kill();
      await settleWithin(child.exited, 1000);
      child.kill("SIGKILL");
      await settleWithin(Promise.allSettled([child.exited, inputClosed, this.reading]), 500);
      // escaped descendants may retain inherited pipes even after the process group is dead
      await settleWithin(
        Promise.allSettled([
          ...this.readers.map((reader) => reader.cancel().catch(() => {})),
          this.writer?.abort().catch(() => {}),
          this.reading,
        ]),
        500,
      );
    })();
  }
}