import { ChainCommand, Command } from "./command";
import { LLMInvoker, UserMessage, Memory } from "./llm-invoker";
import { GenericMemoryBuilder } from "./builders/memory-builder";
import { LlmResponseAdapter } from "./adapters/llm-response-adapter";

// Represents a task to be executed by an agent.
export interface AgentTask {
  // Identifier for the agent that should handle this task.
  agent: string;
  // The command or payload for the task.
  command: string;
  // Additional task-specific properties.
  originalInput: UserMessage[];
  [key: string]: any;
}

// Additional configuration options can be specified here if needed.
export interface AgentConfig {
  // Optionally, if you need to pass extra configuration
  // For example, if you want to pass a specific model identifier etc.
  defaultMemory: Memory[];
  [key: string]: any;
}

/**
 * Example usage:
 *
 * This code might typically be placed in an application bootstrap file.
 *
 * @example
 * ```typescript
 * async function main() {
 * // Configuration options - these can be loaded from environment variables or a config file.
 * const config: AgentConfig = {
 * defaultMemory: [<AI system memory>]
 * };
 * 
 * const agent = new Agent(config);
 * 
 * // Create and register an Email command.
 * const emailCommand = new Command<AgentTask, void>();
 * emailCommand.setTask(async (task: AgentTask) => {
 * console.log(
 * `[EmailCommand] Processing email command: "${task.command}"`
 * );
 * 
 * // Implement email logic here (e.g., trigger an email sending service).
 * });
 * 
 * agent.registerAgent("email", emailCommand);
 * 
 * // Always running agent
 * agent.registerAlwaysRunAgent("analytic", analyticCommand);
 * 
 * // Optionally, you can register other commands here by creating new Command instances
 * // and assigning them tasks that match your application's behavior.
 * 
 * // Process an input prompt. The LLM is expected to choose an appropriate agent.
 * const testInput = "Initiate onboarding email sequence for new users";
 * await agent.processInput(testInput);
 * }
 * ```

------------------------------
Agent class.
This class encapsulates communication with AWS Bedrock,
synthesizes the LLM response into an AgentTask,
and dispatches the task to the appropriate command from a registry.
------------------------------
*/
export class Agent {
  private llmInvoker: LLMInvoker;
  // Map: agent name → Command instance.
  private registry: Map<
    string,
    Command<AgentTask, any> | ChainCommand<AgentTask | unknown, any>
  >;
  // Registry for agents that should run regardless, in parallel.
  private alwaysRunAgents: Array<{
    agentName: string;
    command: Command<AgentTask, any> | ChainCommand<AgentTask | unknown, any>;
  }>;
  // The fallback command, used if no agent is matched.
  private defaultCommand: Command<AgentTask, any>;
  private defaultMemory: Memory[] = [];

  /**
   * @param llmInvoker An instance that conforms to the LLMInvoker interface.
   * @param config Optional configuration.
   */
  constructor(llmInvoker: LLMInvoker, config: AgentConfig) {
    this.llmInvoker = llmInvoker;
    this.registry = new Map<string, Command<AgentTask, any>>();
    this.alwaysRunAgents = [];

    if (config) {
      this.defaultMemory = config.defaultMemory;
    }

    // Initialize the default command.
    this.defaultCommand = new Command<AgentTask, any>("default");
    this.defaultCommand.setTask(async (task: AgentTask | undefined) => {
      if (!task) {
        console.log(`[DefaultCommand] No task provided.`);
        return `Default action executed with no command.`;
      }
      console.log(`[DefaultCommand] Executing default action: ${task.command}`);

      // Invoke the LLM client (adapter) using our supplied messages and memories.
      const response = await this.llmInvoker.invoke(
        task.originalInput,
        this.defaultMemory
      );

      const llmResponse = new LlmResponseAdapter(response);
      const content = llmResponse.extractContent();

      return {
        message: {
          content: (content as {command: unknown})?.command ?? content,
          role: response.message.role
        },
        usage: response.usage
      };
    }, "Default command executed");

    // Register default command under "default".
    this.registry.set("default", this.defaultCommand);
  }

  /**
   * Registers a new command (agent) to handle tasks.
   *
   * @param agentName Unique key identifying the command.
   * @param command A Command instance that encapsulates the behavior.
   */
  public registerAgent(
    agentName: string,
    command: Command<AgentTask, any> | ChainCommand<AgentTask | unknown, any>
  ): void {
    this.registry.set(agentName, command);
    console.log(`Agent registered: ${agentName}`);
  }

  /**
   * Register an always-run (parallel) agent that executes on every input.
   * @param agentName Unique identifier (for logging) and the command instance.
   * @param command A Command instance to run regardless of LLM routing.
   */
  public registerAlwaysRunAgent(
    agentName: string,
    command: Command<AgentTask, any> | ChainCommand<AgentTask | unknown, any>
  ): void {
    this.alwaysRunAgents.push({ agentName, command });
    console.log(`Always-run agent registered: ${agentName}`);
  }

  /**
   * Dispatches an AgentTask to the appropriate command.
   * Returns the output of the executed command.
   *
   * @param task The task generated from the LLM response.
   */
  public async dispatchTask(task: AgentTask): Promise<any> {
    const command = this.registry.get(task.agent) || this.defaultCommand;
    console.log(
      `Dispatching task to: ${
        this.registry.get(task.agent) ? task.agent : "default"
      }`
    );
    return await command.execute(task);
  }

  /**
   * Processes a user input prompt:
   * • Converts it into a UserMessage.
   * • Calls the injected LLM client.
   * • Decodes the response to synthesize an AgentTask.
   * • Dispatches the task and returns its output.
   *
   * @param input The user input string.
   * @param memories Optionally, a list of Memory objects.
   */
  public async processInput(
    input: UserMessage[],
    memories: Memory[] = []
  ): Promise<any> {
    if (!input) {
      throw new Error(`Input argument is missing from processInput()`);
    }

    const memoryBuilder = new GenericMemoryBuilder();

    let systemMemory = JSON.parse(JSON.stringify(memories))

    try {
      // Build the required UserMessage structure.
      if (!systemMemory || (Array.isArray(systemMemory) && memories.length <= 0)) {
        let availableAgents: string = "";

        this.registry.forEach((command) => {
          availableAgents += `${command.getDescription()}`;
        });

        console.log(
          `Available agents: ${availableAgents || "none registered"}`
        );

        systemMemory = memoryBuilder
          .addMemory(
            `You are an advanced task routing system that receives user instructions and 
            determines which specialized agent should handle the task. You have access to 
            the following agents: ${availableAgents}, 
            and 'default' for all other tasks. Your response must be a valid JSON object with 
            exactly two properties: 'agent' and 'command'. For instance, if the input is 
            'Send a welcome email to new users', your output should be: 
            { \"agent\": \"email\", \"command\": \"Send a welcome email to new users\" }. 
             Follow these instructions exactly and output only valid JSON.`
          )
          .build() as Memory[];
      }

      // Invoke the LLM client (adapter) using our supplied messages and memories.
      const response = await this.llmInvoker.invoke(input, systemMemory);

      // Generate the agent routes
      const llmResponseAdapter = new LlmResponseAdapter(response);
      const route = llmResponseAdapter.extractContent();
      console.log("[input routing] route: ", JSON.stringify(route))

      let tasksOutput: { [key: string]: any } = {};
      if (
        typeof route === "object" &&
        route !== null &&
        (route as { [key: string]: unknown }).agent &&
        (route as { [key: string]: unknown }).command
      ) {
        // Synthesize the AgentTask from the route object.
        const agentTask: AgentTask = {
          agent: (route as any).agent || "default",
          command: (route as any).command || input,
          originalInput: input,
        };

        // Dispatch the task and return the output.
        const output = await this.dispatchTask(agentTask);

        tasksOutput[agentTask.agent] = {
          ...output,
        };
      }

      // Initiate all always-run agents to execute in parallel.
      const alwaysRunPromises = this.alwaysRunAgents.map(async (cmd) => {
        const result = await cmd.command.execute({
          agent: cmd.agentName,
          command: "run",
          originalInput: input,
        });

        if (cmd.agentName) {
          return {
            [cmd.agentName]: result,
          };
        }

        return;
      });

      // Await both main task and always-run agents concurrently.
      // (Using Promise.allSettled so that errors in always-run agents do not affect the main task.)
      const [alwaysRunResults] = await Promise.all([
        Promise.allSettled(alwaysRunPromises),
      ]);

      alwaysRunResults.map((result) => {
        if (result.status === "fulfilled" && result.value) {
          const keys = Object.keys(result.value);

          if (!tasksOutput.default) {
            tasksOutput.default = {};
          }
          console.log("keys: ", JSON.stringify(keys));
          console.log("key value: ", JSON.stringify(result));
          tasksOutput.default[keys[0]] = result.value[keys[0]];
        }
      });

      return tasksOutput;
    } catch (error) {
      console.error("Error processing input:", error);
      throw error;
    }
  }

  /**
   * Returns the registry of commands.
   * Each command is represented by its name and the command instance.
   */
  public getRegistry(): Map<
    string,
    Command<AgentTask, any> | ChainCommand<AgentTask | unknown, any>
  > {
    return this.registry;
  }

  /**
   * Returns the list of always-run agents.
   * Each agent is represented by its name and the command to execute.
   */
  public getAlwaysRunAgents(): Array<{
    agentName: string;
    command: Command<AgentTask, any> | ChainCommand<AgentTask | unknown, any>;
  }> {
    return this.alwaysRunAgents;
  }
}