import { AgentRuntime } from "./AgentRuntime";
import { AgentState, ToolDefinition } from "./types";

export interface AgentWorkerConfig {
  id: string;
  name: string;
  role: string;
  systemPrompt: string;
  tools: ToolDefinition<any>[];
}

export class AgentWorker {
  public id: string;
  public name: string;
  private runtime: AgentRuntime;
  private config: AgentWorkerConfig;

  constructor(config: AgentWorkerConfig) {
    this.id = config.id;
    this.name = config.name;
    this.config = config;
    this.runtime = new AgentRuntime(6);

    for (const tool of config.tools) {
      this.runtime.registerTool(tool);
    }
  }

  public async executeSubtask(taskPrompt: string): Promise<AgentState> {
    return this.runtime.run(taskPrompt, this.config.systemPrompt);
  }
}

import { AgentWorker } from "./AgentWorker";
import { AgentState } from "./types";

export class SwarmOrchestrator {
  private workers: Map<string, AgentWorker> = new Map();

  public registerWorker(worker: AgentWorker): void {
    this.workers.set(worker.id, worker);
  }

  /**
   * Sequential Pipeline Execution with Context Handoff
   */
  public async executePipeline(
    pipeline: string[],
    initialTask: string
  ): Promise<Record<string, AgentState>> {
    const stageResults: Record<string, AgentState> = {};
    let activePrompt = initialTask;

    for (const agentId of pipeline) {
      const worker = this.workers.get(agentId);
      if (!worker) {
        throw new Error(`Orchestrator: Worker '${agentId}' not registered.`);
      }

      console.log(`[Swarm Orchestrator] Invoking stage: ${worker.name}...`);
      const state = await worker.executeSubtask(activePrompt);
      stageResults[agentId] = state;

      // Extract output to feed into next agent stage
      const lastMessage = state.messages[state.messages.length - 1];
      if (lastMessage && lastMessage.content) {
        activePrompt = `### Output from ${worker.name}:\n${lastMessage.content}\n\n### Next Directive:\nProcess and finalize the next step.`;
      }
    }

    return stageResults;
  }

  /**
   * Parallel Fan-Out Execution
   */
  public async executeParallelFanOut(
    agentIds: string[],
    tasks: Record<string, string>
  ): Promise<Record<string, AgentState>> {
    const promises = agentIds.map(async (id) => {
      const worker = this.workers.get(id);
      if (!worker) throw new Error(`Worker '${id}' not found`);
      const task = tasks[id] || "Process stage input";
      const result = await worker.executeSubtask(task);
      return { id, result };
    });

    const results = await Promise.all(promises);
    const outputMap: Record<string, AgentState> = {};
    for (const item of results) {
      outputMap[item.id] = item.result;
    }
    return outputMap;
  }
}