import { Cause, Duration, Effect, Exit, Predicate } from "effect";

import type { McpConnection, McpConnector } from "./connection";
import type { McpInvocationError } from "./errors";

// The pool preserves sessions for sessionful legacy servers. Stateless
// 2026-07-28 servers do not need it, but retaining a cheap idle client is
// harmless and keeps one lifecycle for both protocol eras.

const IDLE_TTL_MS = 5 * 60 * 1_000;

/** How long a `close()` is waited on before the connection is abandoned.
 *
 *  Eviction is driven by live traffic — the invocation that acquires a lease is
 *  the one that runs the sweep — so an unbounded close is that caller's problem:
 *  a server that accepts the close and then goes quiet would hold up a request
 *  that has nothing to do with the connection being reclaimed. Teardown is
 *  milliseconds' work when it works at all, so anything past this window is a
 *  socket that is not coming back. */
const CLOSE_TIMEOUT = Duration.seconds(2);

type IdleConnection = {
  readonly connection: McpConnection;
  readonly idleSince: number;
};

type ConnectionLease = {
  readonly connection: McpConnection;
  readonly reused: boolean;
};

const closeQuietly = (connection: McpConnection): Effect.Effect<void> =>
  Effect.tryPromise(() => connection.close()).pipe(Effect.timeout(CLOSE_TIMEOUT), Effect.ignore);

const isMcpInvocationError = (error: unknown): error is McpInvocationError =>
  Predicate.isTagged(error, "McpInvocationError");

const isDeadConnectionFailure = (error: unknown): boolean => {
  if (Predicate.isTagged(error, "McpConnectionError")) return true;
  if (!isMcpInvocationError(error)) return false;
  return (
    error.transportFailure === true ||
    error.status === 400 ||
    // A 401 means the bearer this session was dialled with is no longer
    // accepted. The session is bound to that token for its lifetime, so
    // retaining it would hand the same dead credential to every caller for the
    // full idle window — and would defeat core's post-401 token refresh, whose
    // freshly minted token can only reach the server over a NEW session. Drop
    // it so the retry dials with the new credential.
    error.status === 401 ||
    error.status === 404 ||
    error.status === 408
  );
};

const shouldDropConnection = <A, E>(exit: Exit.Exit<A, E>): boolean => {
  if (Exit.isSuccess(exit)) return false;
  const failures = exit.cause.reasons.filter(Cause.isFailReason);
  if (failures.length === 0) return true;
  return failures.some((failure) => isDeadConnectionFailure(failure.error));
};

/** A per-plugin-instance pool that gives each invocation an exclusive MCP
 * connection lease while retaining at most one idle session per identity. */
export interface McpConnectionPool {
  readonly withConnection: <A, E, R>(
    key: string,
    connector: McpConnector,
    use: (connection: McpConnection) => Effect.Effect<A, E, R>,
  ) => Effect.Effect<A, E | Effect.Error<McpConnector>, R>;
  /** Closes and removes every currently idle connection. Leased connections
   * remain owned by their in-flight invocation and follow normal release. */
  readonly close: () => Effect.Effect<void>;
}

/** Creates an MCP connection pool with lazy five-minute idle eviction and one
 * automatic fresh-dial retry for a reused session rejected with HTTP 404.
 *
 * "Lazy" means activity-driven — there is no timer and no background fiber — but
 * it applies to EVERY parked connection, not only the identity being asked for.
 * A pooled session holds the credential it was dialled with, so an identity that
 * is never requested again must still age out. */
export const createMcpConnectionPool = (): McpConnectionPool => {
  const idle = new Map<string, IdleConnection>();

  /** Close and drop every entry past the idle window, not just the one being
   *  asked for.
   *
   *  The TTL used to be consulted only against `idle.get(key)`, so an identity
   *  that was never dialled again was never examined again: its session stayed
   *  open and authenticated indefinitely, holding the bearer it was dialled
   *  with. The advertised bound only held for connections that happened to be
   *  reused.
   *
   *  Still lazy — activity drives it, there is no timer and no background fiber.
   *  The map holds at most one entry per identity, so scanning it is trivial.
   *
   *  The entries leave the map synchronously, before any close is awaited, so a
   *  slow teardown can never hand the same connection to a second caller. The
   *  closes themselves run concurrently and each is bounded by `CLOSE_TIMEOUT`,
   *  the same shape `close()` below uses: the sweep is work the acquiring
   *  invocation pays for, and one unresponsive server must not be able to stall
   *  it, let alone stall the connections queued behind it. */
  const sweepExpired = Effect.suspend(() => {
    const now = Date.now();
    const expired: McpConnection[] = [];
    for (const [key, entry] of idle) {
      if (now - entry.idleSince < IDLE_TTL_MS) continue;
      idle.delete(key);
      expired.push(entry.connection);
    }
    return Effect.forEach(expired, closeQuietly, { concurrency: "unbounded", discard: true });
  });

  const acquire = (key: string, connector: McpConnector, forceFresh: boolean) =>
    Effect.gen(function* () {
      yield* sweepExpired;
      if (forceFresh) {
        const connection = yield* connector;
        return { connection, reused: false } satisfies ConnectionLease;
      }

      const taken = yield* Effect.sync(() => {
        const entry = idle.get(key);
        if (!entry) return { connection: undefined, expired: undefined };
        idle.delete(key);
        if (Date.now() - entry.idleSince >= IDLE_TTL_MS) {
          return { connection: undefined, expired: entry.connection };
        }
        return { connection: entry.connection, expired: undefined };
      });

      if (taken.expired) yield* closeQuietly(taken.expired);
      if (taken.connection) {
        return { connection: taken.connection, reused: true } satisfies ConnectionLease;
      }

      const connection = yield* connector;
      return { connection, reused: false } satisfies ConnectionLease;
    });

  const release = <A, E>(key: string, lease: ConnectionLease, exit: Exit.Exit<A, E>) =>
    Effect.gen(function* () {
      if (shouldDropConnection(exit)) {
        yield* closeQuietly(lease.connection);
        return;
      }

      const displaced = yield* Effect.sync(() => {
        if (idle.has(key)) return lease.connection;
        idle.set(key, { connection: lease.connection, idleSince: Date.now() });
        return undefined;
      });
      if (displaced) yield* closeQuietly(displaced);
    });

  const withConnection: McpConnectionPool["withConnection"] = (key, connector, use) => {
    let reused = false;
    const run = (forceFresh: boolean) =>
      Effect.acquireUseRelease(
        acquire(key, connector, forceFresh),
        (lease) => {
          if (!forceFresh) reused = lease.reused;
          return use(lease.connection);
        },
        (lease, exit) => release(key, lease, exit),
      );

    return run(false).pipe(
      Effect.catch((error) =>
        reused && isMcpInvocationError(error) && error.status === 404
          ? run(true)
          : Effect.fail(error),
      ),
    );
  };

  const close = () =>
    Effect.gen(function* () {
      const connections = yield* Effect.sync(() => {
        const connections = [...idle.values()].map((entry) => entry.connection);
        idle.clear();
        return connections;
      });
      yield* Effect.forEach(connections, closeQuietly, {
        concurrency: "unbounded",
        discard: true,
      });
    });

  return { withConnection, close };
};