import { randomUUID } from "node:crypto";
import type {
  AgentToolResult,
  EventBus,
  ExtensionContext,
  ToolDefinition,
} from "@earendil-works/pi-coding-agent";
import type { TSchema } from "typebox";
import { Check } from "typebox/value";
import {
  isCancelFrame,
  isProbeFrame,
  isRequestFrame,
  PEER_CHANNEL_OFFER,
  PEER_CHANNEL_PROBE,
  peerCancelChannel,
  peerRequestChannel,
  peerResponseChannel,
} from "./protocol.ts";

export interface ToolPeerServerOptions {
  /** The extension's shared event bus (`pi.events`). */
  readonly events: EventBus;
  /** Resolves the extension's active ExtensionContext for tool execution. */
  readonly getContext: () => ExtensionContext | undefined;
  /** The tool definitions this extension owns and exposes to eval cells. */
  readonly tools: readonly ToolDefinition<TSchema, unknown, unknown>[];
}

/**
 * Opt-in server side of the cooperative tool protocol (Plan 003). An
 * extension calls this with its own tool definitions; eval cells in the same
 * pi session can then invoke them through `tool.<name>()`. The server runs
 * the OWNER's `prepareArguments`, schema validation, and `execute` with the
 * OWNER's context — pi-codemode never reimplements host tool semantics.
 *
 * Streaming updates (`onUpdate`) are not carried in v1; calls are result-only.
 */
export function createToolPeerServer(options: ToolPeerServerOptions): {
  readonly dispose: () => void;
} {
  const serverId = randomUUID();
  const tools = new Map(
    options.tools.map((tool) => [tool.name, eraseTool(tool)])
  );
  const inFlight = new Map<string, AbortController>();

  const probeUnsubscribe = options.events.on(PEER_CHANNEL_PROBE, (data) => {
    if (!isProbeFrame(data)) {
      return;
    }
    options.events.emit(PEER_CHANNEL_OFFER, {
      v: 1,
      serverId,
      toolNames: [...tools.keys()],
    });
  });

  const requestUnsubscribe = options.events.on(
    peerRequestChannel(serverId),
    (data) => {
      if (!isRequestFrame(data)) {
        return;
      }
      runRequest(options, tools, serverId, data, inFlight).catch(
        (error: unknown) => {
          options.events.emit(peerResponseChannel(data.requestId), {
            v: 1,
            ok: false,
            requestId: data.requestId,
            error: {
              message: error instanceof Error ? error.message : String(error),
            },
          });
        }
      );
    }
  );

  const cancelUnsubscribe = options.events.on(
    peerCancelChannel(serverId),
    (data) => {
      if (!isCancelFrame(data)) {
        return;
      }
      inFlight.get(data.requestId)?.abort();
    }
  );

  return {
    dispose: () => {
      probeUnsubscribe();
      requestUnsubscribe();
      cancelUnsubscribe();
      for (const controller of inFlight.values()) {
        controller.abort();
      }
      inFlight.clear();
    },
  };
}

interface ErasedPeerTool {
  readonly execute: (
    toolCallId: string,
    params: unknown,
    signal: AbortSignal,
    ctx: ExtensionContext
  ) => Promise<AgentToolResult<unknown>>;
  readonly name: string;
}

function eraseTool(
  tool: ToolDefinition<TSchema, unknown, unknown>
): ErasedPeerTool {
  return {
    name: tool.name,
    execute: async (toolCallId, params, signal, ctx) => {
      let prepared: unknown = params;
      if (tool.prepareArguments !== undefined) {
        prepared = tool.prepareArguments(params);
      }
      if (!Check(tool.parameters, prepared)) {
        throw new Error(`Invalid arguments for tool ${tool.name}`);
      }
      return await tool.execute(toolCallId, prepared, signal, undefined, ctx);
    },
  };
}

interface RequestFrameLike {
  readonly params: unknown;
  readonly requestId: string;
  readonly toolName: string;
}

async function runRequest(
  options: ToolPeerServerOptions,
  tools: ReadonlyMap<string, ErasedPeerTool>,
  serverId: string,
  request: RequestFrameLike,
  inFlight: Map<string, AbortController>
): Promise<void> {
  const tool = tools.get(request.toolName);
  if (tool === undefined) {
    return;
  }
  const ctx = options.getContext();
  if (ctx === undefined) {
    options.events.emit(peerResponseChannel(request.requestId), {
      v: 1,
      ok: false,
      requestId: request.requestId,
      error: {
        message: `${request.toolName} peer has no active session context`,
      },
    });
    return;
  }
  const controller = new AbortController();
  inFlight.set(request.requestId, controller);
  try {
    const result = await tool.execute(
      `peer-${serverId.slice(0, 8)}-${request.requestId}`,
      request.params,
      controller.signal,
      ctx
    );
    options.events.emit(peerResponseChannel(request.requestId), {
      v: 1,
      ok: true,
      requestId: request.requestId,
      value: result,
    });
  } catch (error) {
    const message =
      error instanceof Error && error.name !== "AbortError"
        ? error.message
        : "peer tool call was cancelled";
    options.events.emit(peerResponseChannel(request.requestId), {
      v: 1,
      ok: false,
      requestId: request.requestId,
      error: { message },
    });
  } finally {
    inFlight.delete(request.requestId);
  }
}