import { randomUUID } from "node:crypto";
import type { EventBus } from "@earendil-works/pi-coding-agent";
import {
  isOfferFrame,
  isResponseFrame,
  PEER_CHANNEL_OFFER,
  PEER_CHANNEL_PROBE,
  peerCancelChannel,
  peerRequestChannel,
  peerResponseChannel,
} from "./protocol.ts";

export interface PeerToolClientOptions {
  readonly events: EventBus;
  /** Window during which cooperating servers may answer the session probe. */
  readonly probeTimeoutMs?: number;
  /** Upper bound for one peer tool call; the kernel bridge timeout still bounds the whole call. */
  readonly requestTimeoutMs?: number;
}

export interface PeerCallResult {
  readonly message: string;
  readonly ok: false;
}

export interface PeerCallOkResult {
  readonly ok: true;
  readonly value: unknown;
}

export type PeerToolCallResult = PeerCallResult | PeerCallOkResult;

export const PEER_DEFAULT_PROBE_TIMEOUT_MS = 250;
export const PEER_DEFAULT_REQUEST_TIMEOUT_MS = 30_000;

interface PendingPeerCall {
  readonly requestId: string;
  readonly resolve: (result: PeerToolCallResult) => void;
  readonly responseChannel: string;
  readonly serverId: string;
  readonly settleTimer: NodeJS.Timeout;
  readonly signalUnsubscribe: (() => void) | undefined;
  readonly toolName: string;
  readonly unsubscribe: () => void;
}

/**
 * Session-scoped client side of the cooperative tool protocol. Bindings are
 * (re)established by `probe()` and cleared by `clear()`. A call only reaches
 * a server that explicitly offered the tool name, so unbound tools keep the
 * shim's ordinary refusal with no added latency.
 */
export class PeerToolClient {
  readonly #events: EventBus;
  readonly #probeTimeoutMs: number;
  readonly #requestTimeoutMs: number;
  readonly #bindings = new Map<string, string>();
  readonly #pending = new Map<string, PendingPeerCall>();
  #offerUnsubscribe: (() => void) | undefined;
  #probing: Promise<void> | undefined;

  constructor(options: PeerToolClientOptions) {
    this.#events = options.events;
    this.#probeTimeoutMs =
      options.probeTimeoutMs ?? PEER_DEFAULT_PROBE_TIMEOUT_MS;
    this.#requestTimeoutMs =
      options.requestTimeoutMs ?? PEER_DEFAULT_REQUEST_TIMEOUT_MS;
  }

  isBound(toolName: string): boolean {
    return this.#bindings.has(toolName);
  }

  /**
   * Probe for cooperating servers and bind the tools they own. First offer
   * per name wins and stays stable for the session; later probes replace the
   * whole binding set. Concurrent probes serialize on the same promise.
   */
  probe(activeToolNames: readonly string[]): Promise<void> {
    if (this.#probing !== undefined) {
      return this.#probing;
    }
    const probing = this.#probeOnce(activeToolNames);
    this.#probing = probing;
    return probing;
  }

  async #probeOnce(activeToolNames: readonly string[]): Promise<void> {
    const allowed = new Set(activeToolNames);
    const bindings = new Map<string, string>();
    const offerUnsubscribe = this.#events.on(PEER_CHANNEL_OFFER, (data) => {
      if (!isOfferFrame(data)) {
        return;
      }
      for (const toolName of data.toolNames) {
        if (allowed.has(toolName) && !bindings.has(toolName)) {
          bindings.set(toolName, data.serverId);
        }
      }
    });
    this.#offerUnsubscribe?.();
    this.#offerUnsubscribe = offerUnsubscribe;
    this.#events.emit(PEER_CHANNEL_PROBE, { v: 1 });
    await new Promise<void>((resolve) => {
      setTimeout(resolve, this.#probeTimeoutMs);
    });
    offerUnsubscribe();
    if (this.#offerUnsubscribe === offerUnsubscribe) {
      this.#offerUnsubscribe = undefined;
    }
    if (this.#probing === undefined) {
      // A clear() superseded this probe; drop the collected offers.
      return;
    }
    this.#bindings.clear();
    for (const [toolName, serverId] of bindings) {
      this.#bindings.set(toolName, serverId);
    }
    this.#probing = undefined;
  }

  /**
   * Await any in-flight session probe, then dispatch only when the tool is
   * bound. Returns `undefined` for unbound tools so the shim can keep its
   * ordinary refusal with no added latency once the probe has settled.
   */
  async callIfBound(
    toolName: string,
    params: unknown,
    options?: { signal?: AbortSignal }
  ): Promise<PeerToolCallResult | undefined> {
    const probing = this.#probing;
    if (probing !== undefined) {
      await probing;
    }
    if (!this.#bindings.has(toolName)) {
      return;
    }
    return await this.call(toolName, params, options);
  }

  async call(
    toolName: string,
    params: unknown,
    options?: { signal?: AbortSignal }
  ): Promise<PeerToolCallResult> {
    const serverId = this.#bindings.get(toolName);
    if (serverId === undefined) {
      return {
        ok: false,
        message: `Tool ${toolName} has no cooperating peer in this session`,
      };
    }
    const requestId = randomUUID();
    const responseChannel = peerResponseChannel(requestId);
    const timeoutMs = this.#requestTimeoutMs;
    return await new Promise<PeerToolCallResult>((resolve) => {
      let settled = false;
      const settle = (result: PeerToolCallResult): void => {
        if (settled) {
          return;
        }
        settled = true;
        clearTimeout(settleTimer);
        unsubscribe();
        signalUnsubscribe?.();
        this.#pending.delete(requestId);
        resolve(result);
      };
      const unsubscribe = this.#events.on(responseChannel, (data) => {
        if (!isResponseFrame(data) || data.requestId !== requestId) {
          return;
        }
        if (data.ok) {
          settle({ ok: true, value: data.value });
        } else {
          settle({ ok: false, message: data.error.message });
        }
      });
      const settleTimer = setTimeout(() => {
        settle({
          ok: false,
          message: `Tool ${toolName} peer call timed out after ${timeoutMs}ms`,
        });
        this.#emitCancel(serverId, requestId);
      }, timeoutMs);
      const onAbort = (): void => {
        settle({
          ok: false,
          message: `Tool ${toolName} peer call aborted`,
        });
        this.#emitCancel(serverId, requestId);
      };
      const signalUnsubscribe =
        options?.signal === undefined
          ? undefined
          : () => options.signal?.removeEventListener("abort", onAbort);
      if (options?.signal?.aborted) {
        onAbort();
        return;
      }
      options?.signal?.addEventListener("abort", onAbort, { once: true });
      const pending: PendingPeerCall = {
        requestId,
        resolve: settle,
        settleTimer,
        responseChannel,
        unsubscribe,
        signalUnsubscribe,
        serverId,
        toolName,
      };
      this.#pending.set(requestId, pending);
      this.#events.emit(peerRequestChannel(serverId), {
        v: 1,
        requestId,
        toolName,
        params,
      });
    });
  }

  #emitCancel(serverId: string, requestId: string): void {
    this.#events.emit(peerCancelChannel(serverId), { v: 1, requestId });
  }

  /**
   * Drop all bindings and settle every in-flight call as failed. Called on
   * session shutdown/switch/fork so no pending peer work survives a session.
   */
  clear(): void {
    this.#probing = undefined;
    this.#offerUnsubscribe?.();
    this.#offerUnsubscribe = undefined;
    this.#bindings.clear();
    for (const pending of [...this.#pending.values()]) {
      pending.resolve({
        ok: false,
        message: `Tool ${pending.toolName} peer call cancelled: session ended`,
      });
      pending.signalUnsubscribe?.();
    }
  }
}