declare module "agents-js" {
  export interface AgentOptions {
    name?: string;
    model?: string;
    instructions?: string | ((context: Record<string, any>) => string);
    functions?: (Function | ToolConfirmation)[];
    toolChoice?: string | null;
    parallelToolCalls?: boolean;
    beforeModel?: (context: CallbackContext) => Promise<any> | any;
    afterModel?: (
      context: CallbackContext & { response: any }
    ) => Promise<any> | any;
    beforeTool?: (
      toolName: string,
      args: Record<string, any>,
      contextVariables: Record<string, any>
    ) => Promise<boolean | void> | boolean | void;
    afterTool?: (
      toolName: string,
      args: Record<string, any>,
      result: Result,
      contextVariables: Record<string, any>
    ) => Promise<void> | void;
    subAgents?: Agent[];
    description?: string;
  }

  export interface CallbackContext {
    agent: Agent;
    messages: Message[];
    contextVariables: Record<string, any>;
  }

  export interface Message {
    role: "system" | "user" | "assistant" | "tool";
    content: string;
    tool_calls?: ToolCall[];
    tool_call_id?: string;
    tool_name?: string;
    sender?: string;
  }

  export interface ToolCall {
    id: string;
    type: string;
    function: {
      name: string;
      arguments: string;
    };
  }

  export class Agent {
    name: string;
    model: string;
    instructions: string | ((context: Record<string, any>) => string);
    functions: (Function | ToolConfirmation)[];
    toolChoice: string | null;
    parallelToolCalls: boolean;
    beforeModel: ((context: CallbackContext) => Promise<any> | any) | null;
    afterModel:
      | ((context: CallbackContext & { response: any }) => Promise<any> | any)
      | null;
    beforeTool:
      | ((
          toolName: string,
          args: Record<string, any>,
          contextVariables: Record<string, any>
        ) => Promise<boolean | void> | boolean | void)
      | null;
    afterTool:
      | ((
          toolName: string,
          args: Record<string, any>,
          result: Result,
          contextVariables: Record<string, any>
        ) => Promise<void> | void)
      | null;
    subAgents: Agent[];
    description: string;
    _originalSubAgents: Agent[];

    constructor(options?: AgentOptions);
  }

  export interface ResultOptions {
    value?: string;
    agent?: Agent | null;
    agentName?: string | null;
    contextVariables?: Record<string, any>;
  }

  export class Result {
    value: string;
    agent: Agent | null;
    agentName: string | null;
    contextVariables: Record<string, any>;

    constructor(options?: ResultOptions);
  }

  export interface TokenUsage {
    totalInputTokens: number;
    totalOutputTokens: number;
    inputTokens?: number;
    outputTokens?: number;
    totalTokens?: number;
  }

  export interface ResponseOptions {
    messages?: Message[];
    agent?: Agent | null;
    contextVariables?: Record<string, any>;
    tokenUsage?: TokenUsage | null;
  }

  export class Response {
    messages: Message[];
    agent: Agent | null;
    contextVariables: Record<string, any>;
    tokenUsage: TokenUsage | null;

    constructor(options?: ResponseOptions);
  }

  export interface SessionOptions {
    id?: string | null;
    storage?: Storage | null;
  }

  export class Session {
    id: string;
    storage: Storage | null;
    data: Record<string, any>;

    constructor(options?: SessionOptions);
    get(key: string): Promise<any>;
    set(key: string, value: any): Promise<void>;
    delete(key: string): Promise<void>;
    clear(): Promise<void>;
  }

  export interface Storage {
    get(sessionId: string, key: string): Promise<any>;
    set(sessionId: string, key: string, value: any): Promise<void>;
    delete(sessionId: string, key: string): Promise<void>;
    clear(sessionId: string): Promise<void>;
  }

  export class MemoryStorage implements Storage {
    store: Record<string, Record<string, any>>;

    constructor();
    get(sessionId: string, key: string): Promise<any>;
    set(sessionId: string, key: string, value: any): Promise<void>;
    delete(sessionId: string, key: string): Promise<void>;
    clear(sessionId: string): Promise<void>;
  }

  export interface ToolConfirmationOptions {
    tool: Function;
    message?: string | null;
    autoConfirm?: boolean;
  }

  export class ToolConfirmation {
    tool: Function;
    message: string | null;
    autoConfirm: boolean;
    requiresConfirmation: boolean;

    constructor(options: ToolConfirmationOptions);
  }

  export interface AgentControllerOptions {
    confirmationHandler?:
      | ((
          toolName: string,
          args: Record<string, any>,
          message: string
        ) => Promise<boolean>)
      | null;
    apiKeys?: Record<string, string>;
    enableLogging?: boolean;
  }

  export class AgentController {
    client: any;
    confirmationHandler:
      | ((
          toolName: string,
          args: Record<string, any>,
          message: string
        ) => Promise<boolean>)
      | null;
    apiKeys: Record<string, string>;
    tokenUsage: TokenUsage;
    enableLogging: boolean;

    constructor(client: any, options?: AgentControllerOptions);

    setApiKey(service: string, key: string): void;
    getApiKey(service: string): string | undefined;

    updateTokenUsage(model: string, usage: any): TokenUsage;
    resetTokenUsage(): void;
    getTokenUsage(): TokenUsage;
    log(
      message: string,
      type?:
        | "info"
        | "success"
        | "warning"
        | "error"
        | "transfer"
        | "tool"
        | "agent"
    ): void;

    getChatCompletion(
      agent: Agent,
      history: Message[],
      contextVariables: Record<string, any>,
      modelOverride: string | null,
      stream: boolean,
      debug: boolean
    ): Promise<any>;

    handleFunctionResult(result: any, debug: boolean): Result;

    handleToolCalls(
      toolCalls: ToolCall[],
      functions: (Function | ToolConfirmation)[],
      contextVariables: Record<string, any>,
      debug: boolean,
      agent: Agent
    ): Promise<Response>;

    runAndStream(
      agent: Agent,
      messages: Message[],
      contextVariables?: Record<string, any>,
      modelOverride?: string | null,
      debug?: boolean,
      maxTurns?: number,
      executeTools?: boolean
    ): AsyncGenerator<any>;

    run(
      agent: Agent,
      messages: Message[],
      contextVariables?: Record<string, any>,
      modelOverride?: string | null,
      stream?: boolean,
      debug?: boolean,
      maxTurns?: number,
      executeTools?: boolean
    ): Promise<Response> | AsyncGenerator<any>;
  }
}

declare module "agents-js/tools" {
  import { Result } from "agents-js";

  export function serperSearch(params: { query: string }): Promise<Result>;

  export function firecrawlScrape(params: { url: string }): Promise<Result>;

  export function firecrawlSearch(params: { query: string }): Promise<Result>;

  export function firecrawlCrawl(params: {
    url: string;
    limit?: number;
  }): Promise<Result>;

  export function calculate(params: { expression: string }): Result;
  export function add(params: { a: number; b: number }): Result;
  export function subtract(params: { a: number; b: number }): Result;
  export function multiply(params: { a: number; b: number }): Result;
  export function divide(params: { a: number; b: number }): Result;

  export function getCurrentTime(params: { timezone?: string }): Result;
  export function getTimestamp(): Result;
  export function formatDate(params: {
    timestamp: number;
    format?: string;
  }): Result;
}