/**
 * @license
 * Copyright 2026 Google Inc.
 * SPDX-License-Identifier: Apache-2.0
 */

import type {Protocol} from 'devtools-protocol';

import type {CDPSession} from '../api/CDPSession.js';
import type {ElementHandle} from '../api/ElementHandle.js';
import type {Frame} from '../api/Frame.js';
import type {ConsoleMessageLocation} from '../common/ConsoleMessage.js';
import {DEBUG_PREFIXES, type Logger} from '../common/Debug.js';
import {EventEmitter} from '../common/EventEmitter.js';
import {DisposableStack} from '../util/disposable.js';

import type {CdpFrame} from './Frame.js';
import type {FrameManager} from './FrameManager.js';
import {MAIN_WORLD} from './IsolatedWorlds.js';

/**
 * Represents a registered WebMCP tool available on the page.
 *
 * @public
 */
export class WebMCPTool extends EventEmitter<{
  /** Emitted when invocation starts. */
  toolinvoked: WebMCPToolCall;
}> {
  #webmcp: WebMCP;
  #backendNodeId?: number;
  #formElement?: ElementHandle<HTMLFormElement>;

  /**
   * Tool name.
   */
  name: string;
  /**
   * Tool description.
   */
  description: string;
  /**
   * Schema for the tool's input parameters.
   */
  inputSchema?: object;
  /**
   * Optional annotations for the tool.
   */
  annotations?: Protocol.WebMCP.Annotation;
  /**
   * Frame the tool was defined for.
   */
  frame: Frame;
  /**
   * Source location that defined the tool (if available).
   */
  location?: ConsoleMessageLocation;
  /**
   * @internal
   */
  rawStackTrace?: Protocol.Runtime.StackTrace;

  /**
   * @internal
   */
  constructor(webmcp: WebMCP, tool: Protocol.WebMCP.Tool, frame: Frame) {
    super();
    this.#webmcp = webmcp;
    this.name = tool.name;
    this.description = tool.description;
    this.inputSchema = tool.inputSchema;
    this.annotations = tool.annotations;
    this.frame = frame;
    this.#backendNodeId = tool.backendNodeId;
    if (tool.stackTrace?.callFrames.length) {
      this.location = {
        url: tool.stackTrace.callFrames[0]!.url,
        lineNumber: tool.stackTrace.callFrames[0]!.lineNumber,
        columnNumber: tool.stackTrace.callFrames[0]!.columnNumber,
      };
    }
    this.rawStackTrace = tool.stackTrace;
  }

  /**
   * The corresponding ElementHandle when tool was registered via a form.
   */
  get formElement(): Promise<ElementHandle<HTMLFormElement> | undefined> {
    return (async () => {
      if (this.#formElement && !this.#formElement.disposed) {
        return this.#formElement;
      }
      if (!this.#backendNodeId) {
        return undefined;
      }
      this.#formElement = (await (this.frame as CdpFrame).worlds[
        MAIN_WORLD
      ].adoptBackendNode(
        this.#backendNodeId,
      )) as ElementHandle<HTMLFormElement>;
      return this.#formElement;
    })();
  }

  /**
   * Executes tool with input parameters, matching tool's `inputSchema`.
   */
  async execute(
    input: object = {},
    options: WebMCPToolExecuteOptions = {},
  ): Promise<WebMCPToolCallResult> {
    const {invocationId} = await this.#webmcp.invokeTool(this, input);
    return await new Promise<WebMCPToolCallResult>(resolve => {
      const onAbort = () => {
        void this.#webmcp.cancelInvocation(invocationId);
      };
      const handler = (event: WebMCPToolCallResult) => {
        if (event.id === invocationId) {
          options.signal?.removeEventListener('abort', onAbort);
          this.#webmcp.off('toolresponded', handler);
          resolve(event);
        }
      };
      this.#webmcp.on('toolresponded', handler);
      if (options.signal?.aborted) {
        onAbort();
      } else {
        options.signal?.addEventListener('abort', onAbort, {once: true});
      }
    });
  }
}

/**
 * @public
 */
export interface WebMCPToolExecuteOptions {
  /**
   * A signal object that allows you to cancel the tool execution.
   */
  signal?: AbortSignal;
}

/**
 * @public
 */
export interface WebMCPToolsAddedEvent {
  /**
   * Array of tools that were added.
   */
  tools: WebMCPTool[];
}

/**
 * @public
 */
export interface WebMCPToolsRemovedEvent {
  /**
   * Array of tools that were removed.
   */
  tools: WebMCPTool[];
}

/**
 * @public
 */
export class WebMCPToolCall {
  /**
   * Tool invocation identifier.
   */
  id: string;
  /**
   * Tool that was called.
   */
  tool: WebMCPTool;
  /**
   * The input parameters used for the call.
   */
  input: object;

  #logger?: Logger;

  /**
   * @internal
   */
  constructor(
    invocationId: string,
    tool: WebMCPTool,
    input: string,
    logger?: Logger,
  ) {
    this.id = invocationId;
    this.tool = tool;
    this.#logger = logger;
    try {
      this.input = JSON.parse(input);
    } catch (error) {
      this.input = {};
      this.#logger?.(DEBUG_PREFIXES.error)?.(error);
    }
  }
}

/**
 * @public
 */
export interface WebMCPToolCallResult {
  /**
   * Tool invocation identifier.
   */
  id: string;
  /**
   * The corresponding tool call if available.
   */
  call?: WebMCPToolCall;
  /**
   * Status of the invocation.
   */
  status: Protocol.WebMCP.InvocationStatus;
  /**
   * Output or error delivered as delivered to the agent. Missing if `status` is anything
   * other than Completed.
   */
  output?: any;
  /**
   * Error text.
   */
  errorText?: string;
  /**
   * The exception object, if the javascript tool threw an error.
   */
  exception?: Protocol.Runtime.RemoteObject;
}

/**
 * The experimental WebMCP class provides an API for the WebMCP API.
 *
 * See the
 * {@link https://pptr.dev/guides/webmcp|WebMCP guide}
 * for more details.
 *
 * @example
 *
 * ```ts
 * await page.goto('https://www.example.com');
 * const tools = page.webmcp.tools();
 * for (const tool of tools) {
 *   console.log(`Tool found: ${tool.name} - ${tool.description}`);
 * }
 * ```
 *
 * @experimental
 * @public
 */
export class WebMCP extends EventEmitter<{
  /** Emitted when tools are added. */
  toolsadded: WebMCPToolsAddedEvent;
  /** Emitted when tools are removed. */
  toolsremoved: WebMCPToolsRemovedEvent;
  /** Emitted when a tool invocation starts. */
  toolinvoked: WebMCPToolCall;
  /** Emitted when a tool invocation completes or fails. */
  toolresponded: WebMCPToolCallResult;
}> {
  #client: CDPSession;
  #frameManager: FrameManager;
  #tools = new Map<string, Map<string, WebMCPTool>>();
  #pendingCalls = new Map<string, WebMCPToolCall>();
  #logger?: Logger;
  #subscriptions = new DisposableStack();

  #onToolsAdded = (event: Protocol.WebMCP.ToolsAddedEvent) => {
    const tools: WebMCPTool[] = [];
    for (const tool of event.tools) {
      const frame = this.#frameManager.frame(tool.frameId);
      if (!frame) {
        continue;
      }

      const frameTools = this.#tools.get(tool.frameId) ?? new Map();
      if (!this.#tools.has(tool.frameId)) {
        this.#tools.set(tool.frameId, frameTools);
        this.#listenToContextDestroyed(frame as CdpFrame);
      }

      const addedTool = new WebMCPTool(this, tool, frame);
      frameTools.set(tool.name, addedTool);
      tools.push(addedTool);
    }

    this.emit('toolsadded', {tools});
  };

  #onToolsRemoved = (event: Protocol.WebMCP.ToolsRemovedEvent) => {
    const tools: WebMCPTool[] = [];
    event.tools.forEach(tool => {
      const removedTool = this.#tools.get(tool.frameId)?.get(tool.name);
      if (removedTool) {
        tools.push(removedTool);
      }
      this.#tools.get(tool.frameId)?.delete(tool.name);
    });
    this.emit('toolsremoved', {tools});
  };

  #onToolInvoked = (event: Protocol.WebMCP.ToolInvokedEvent) => {
    const tool = this.#tools.get(event.frameId)?.get(event.toolName);
    if (!tool) {
      return;
    }
    const call = new WebMCPToolCall(
      event.invocationId,
      tool,
      event.input,
      this.#logger,
    );
    this.#pendingCalls.set(call.id, call);
    tool.emit('toolinvoked', call);
    this.emit('toolinvoked', call);
  };

  #onToolResponded = (event: Protocol.WebMCP.ToolRespondedEvent) => {
    const call = this.#pendingCalls.get(event.invocationId);
    if (call) {
      this.#pendingCalls.delete(event.invocationId);
    }
    const response: WebMCPToolCallResult = {
      id: event.invocationId,
      call: call,
      status: event.status,
      output: event.output,
      errorText: event.errorText,
      exception: event.exception,
    };
    this.emit('toolresponded', response);
  };

  #onContextDisposed = (frame: CdpFrame) => {
    this.#pendingCalls.clear();
    const frameTools = this.#tools.get(frame._id);
    if (!frameTools) {
      return;
    }
    const tools = Array.from(frameTools.values());
    this.#tools.delete(frame._id);
    if (tools.length) {
      this.emit('toolsremoved', {tools});
    }
  };

  #listenToContextDestroyed(frame: CdpFrame): void {
    frame.mainRealm().context?.once('disposed', () => {
      this.#onContextDisposed(frame);
    });
  }

  /**
   * @internal
   */
  constructor(client: CDPSession, frameManager: FrameManager, logger?: Logger) {
    super(undefined, logger);
    this.#client = client;
    this.#frameManager = frameManager;
    this.#logger = logger;
    this.#bindListeners();
  }

  /**
   * @internal
   */
  async initialize(): Promise<void> {
    return await this.#client.send('WebMCP.enable').catch(err => {
      this.#logger?.(DEBUG_PREFIXES.error)?.(err);
    });
  }

  /**
   * @internal
   */
  async invokeTool(
    tool: WebMCPTool,
    input: object,
  ): Promise<{invocationId: string}> {
    return await this.#client.send('WebMCP.invokeTool', {
      frameId: tool.frame._id,
      toolName: tool.name,
      input,
    });
  }

  /**
   * @internal
   */
  async cancelInvocation(invocationId: string): Promise<void> {
    return await this.#client
      .send('WebMCP.cancelInvocation', {
        invocationId,
      })
      .catch(err => {
        this.#logger?.(DEBUG_PREFIXES.error)?.(err);
      });
  }

  /**
   * Gets all WebMCP tools defined by the page.
   */
  tools(): WebMCPTool[] {
    return Array.from(this.#tools.values()).flatMap(toolMap => {
      return Array.from(toolMap.values());
    });
  }

  #bindListeners(): void {
    const clientEmitter = this.#subscriptions.use(
      new EventEmitter(this.#client),
    );
    clientEmitter.on('WebMCP.toolsAdded', this.#onToolsAdded);
    clientEmitter.on('WebMCP.toolsRemoved', this.#onToolsRemoved);
    clientEmitter.on('WebMCP.toolInvoked', this.#onToolInvoked);
    clientEmitter.on('WebMCP.toolResponded', this.#onToolResponded);
  }

  /**
   * @internal
   */
  updateClient(client: CDPSession): void {
    this.#subscriptions.dispose();
    this.#subscriptions = new DisposableStack();
    this.#client = client;
    this.#bindListeners();
  }
}
