// Copyright (C) 2026 Rob Colbert <rob.colbert@openplatform.us>
// Licensed under the Apache License, Version 2.0

import type { IAiTool } from "@gadget/ai";
import type { IProject, ChatSessionMode } from "@gadget/api";

export interface DroneToolboxEnvironment {
  NODE_ENV: string;
  services?: {
    google?: {
      cse?: {
        apiKey?: string;
        engineId?: string;
      };
    };
  };
  workspace?: {
    workspaceDir?: string;
    projectDir?: string;
    cacheDir?: string;
  };
  project?: IProject;
  chatSessionMode?: ChatSessionMode;
}

export type ToolMap = Map<string, IAiTool>;
export type ToolSet = Set<IAiTool>;

export class AiToolbox {
  private _env: DroneToolboxEnvironment;
  private tools: ToolMap = new Map<string, IAiTool>();
  private modeSets: Map<string, ToolSet> = new Map<string, Set<IAiTool>>();

  constructor(env: DroneToolboxEnvironment) {
    this._env = env;
  }

  get env(): DroneToolboxEnvironment {
    return this._env;
  }

  updateWorkspace(workspace: NonNullable<DroneToolboxEnvironment["workspace"]>): void {
    this._env.workspace = workspace;
  }

  updateProjectContext(project: IProject, mode: ChatSessionMode): void {
    this._env.project = project;
    this._env.chatSessionMode = mode;
  }

  register(tool: IAiTool, modes?: string[]): void {
    if (this.tools.has(tool.name)) {
      throw new Error(`tool already registered: ${tool.name}`);
    }
    this.tools.set(tool.name, tool);

    if (!modes) {
      return;
    }

    for (const mode of modes) {
      let set = this.modeSets.get(mode);
      if (!set) {
        set = new Set<IAiTool>();
        this.modeSets.set(mode, set);
      }
      set.add(tool);
    }
  }

  getTool(name: string): IAiTool | undefined {
    return this.tools.get(name);
  }

  getModeSet(mode: string): ToolSet | undefined {
    return this.modeSets.get(mode);
  }

  getToolNamesForMode(mode: string): string[] {
    return Array.from(this.getModeSet(mode) || []).map((tool) => tool.name);
  }
}
