import type { Environment } from "@berry/api";
import * as posix from "@std/path/posix";
import * as windows from "@std/path/windows";
import {
  McpClient,
  McpConnectionClosedError,
  McpSessionExpiredError,
  StreamableHttpTransport,
} from "@earendil-works/pi-mcp";
import type { Server } from "./config.ts";
import { EnvironmentStdio } from "./stdio.ts";

type Connection = {
  id: string;
  initialization: ReturnType<typeof setTimeout>;
  client: McpClient;
  ready: Promise<unknown>;
  stop: AbortController;
  backend?: AbortSignal;
  disconnected: () => void;
  stdio?: EnvironmentStdio;
  failure?: Error;
  retired: boolean;
  users: number;
  idle?: ReturnType<typeof setTimeout>;
  closing?: Promise<void>;
};

/** Lazy clients shared only within the same machine, working directory, and resolved configuration. */
export class Connections {
  private readonly clients = new Map<string, Connection>();
  private readonly owned = new Set<Connection>();
  private disposed = false;

  private close(connection: Connection): Promise<void> {
    return connection.closing ??= (async () => {
      connection.retired = true;
      clearTimeout(connection.idle);
      clearTimeout(connection.initialization);
      if (this.clients.get(connection.id) === connection) this.clients.delete(connection.id);
      connection.backend?.removeEventListener("abort", connection.disconnected);
      try {
        await connection.client.close();
      } finally {
        connection.stop.abort();
        this.owned.delete(connection);
      }
    })();
  }

  async use<T>(
    env: Environment,
    name: string,
    config: Server,
    signal: AbortSignal,
    run: (client: McpClient) => Promise<T>,
  ): Promise<T> {
    signal.throwIfAborted();
    env.lifetime?.throwIfAborted();
    if (this.disposed) throw new Error("MCP extension has been disposed");
    if (!config.enabled) throw new Error(`MCP server "${name}" is disabled`);
    const id = JSON.stringify([env.id, env.cwd, name, config]);
    let connection = this.clients.get(id);
    if (
      !connection || connection.retired || connection.client.connectionState === "closed" ||
      connection.backend !== env.lifetime
    ) {
      const previous = connection;
      if (previous) previous.retired = true;
      const stop = new AbortController();
      const path = env.platform === "win32" ? windows : posix;
      const client = new McpClient({
        name: "berry",
        version: "1",
        requestTimeoutMs: config.timeout,
        roots: [{ uri: path.toFileUrl(config.cwd).href, name: "workspace" }],
      });
      let stdio: EnvironmentStdio | undefined;
      const transport = "command" in config
        ? (stdio = new EnvironmentStdio(env, config, stop.signal))
        : new StreamableHttpTransport({
          url: config.url,
          headers: config.headers,
          fetch: async (input, init) => {
            try {
              return await env.fetch(input, {
                ...init,
                redirect: "error",
                signal: AbortSignal.any([
                  stop.signal,
                  ...(env.lifetime ? [env.lifetime] : []),
                  ...(init?.signal ? [init.signal] : []),
                ]),
              });
            } catch (error) {
              if (error instanceof TypeError) created.retired = true;
              throw error;
            }
          },
        });
      const initializationTimeout = Math.min(config.timeout, 30_000);
      const timer = setTimeout(() => {
        created.failure = new Error(
          `MCP server "${name}" did not initialize within ${initializationTimeout} ms`,
        );
        void this.close(created);
      }, initializationTimeout);
      const ready = (previous && previous.users === 0 ? this.close(previous) : Promise.resolve())
        .then(() => {
          stop.signal.throwIfAborted();
          return client.connect(transport);
        }).finally(() => clearTimeout(timer));
      const disconnected = () => {
        created.failure = env.lifetime?.reason;
        void this.close(created);
      };
      const created: Connection = {
        id,
        initialization: timer,
        client,
        ready,
        stop,
        backend: env.lifetime,
        disconnected,
        stdio,
        retired: false,
        users: 0,
      };
      connection = created;
      this.clients.set(id, created);
      this.owned.add(created);
      env.lifetime?.addEventListener("abort", disconnected, { once: true });
      if (env.lifetime?.aborted) disconnected();
      void ready.catch(() => {});
    }
    const current = connection;
    clearTimeout(current.idle);
    current.users++;
    try {
      await new Promise<void>((resolve, reject) => {
        const finish = (error?: unknown) => {
          signal.removeEventListener("abort", abort);
          current.stop.signal.removeEventListener("abort", stopped);
          if (error) reject(error);
          else resolve();
        };
        const abort = () => finish(signal.reason);
        const stopped = () => finish(current.failure ?? new McpConnectionClosedError());
        signal.addEventListener("abort", abort, { once: true });
        current.stop.signal.addEventListener("abort", stopped, { once: true });
        void current.ready.then(() => finish(), finish);
        if (signal.aborted) abort();
        else if (current.stop.signal.aborted) stopped();
      });
      signal.throwIfAborted();
      return await run(current.client);
    } catch (error) {
      // retire a failed session without replaying calls or interrupting its other active users
      if (error instanceof McpSessionExpiredError) current.retired = true;
      if (error instanceof McpConnectionClosedError && current.stdio) await current.stdio.close();
      signal.throwIfAborted();
      if (current.failure) throw current.failure;
      if (error instanceof McpConnectionClosedError && current.stdio?.failure) {
        throw current.stdio.failure;
      }
      throw error;
    } finally {
      current.users--;
      if (current.users === 0) {
        if (current.retired || current.client.connectionState === "closed") {
          void this.close(current);
        } else if (!this.disposed && !current.stop.signal.aborted) {
          current.idle = setTimeout(() => {
            if (this.clients.get(id) === current) this.clients.delete(id);
            void this.close(current);
          }, 5 * 60_000);
        }
      }
    }
  }

  async dispose() {
    this.disposed = true;
    this.clients.clear();
    await Promise.all([...this.owned].map((connection) => this.close(connection)));
  }
}