import * as cut from "./agent";
import { BedrockLLMClientAdapter } from "./adapters/bedrock-llm-client-adapter";

jest.mock("./adapters/bedrock-llm-client-adapter", () => ({
  BedrockLLMClientAdapter: jest.fn().mockImplementation(() => ({
    invoke: jest.fn(),
  })),
}));

describe("Agent", () => {
  it("should be implemented", () => {
    expect(cut.Agent).toBeDefined();
    expect(cut.Agent.prototype.dispatchTask).toBeDefined();
    expect(typeof cut.Agent.prototype.dispatchTask).toBe("function");
    expect(cut.Agent.prototype.processInput).toBeDefined();
    expect(typeof cut.Agent.prototype.processInput).toBe("function");
    expect(cut.Agent.prototype.registerAgent).toBeDefined();
    expect(typeof cut.Agent.prototype.registerAgent).toBe("function");
    expect(cut.Agent.prototype.registerAlwaysRunAgent).toBeDefined();
    expect(typeof cut.Agent.prototype.registerAlwaysRunAgent).toBe("function");
  });

  it("should add commands correctly", () => {
    const llm = new BedrockLLMClientAdapter();
    const agent = new cut.Agent(llm, {
      defaultMemory: [{ content: "default memory" }],
    });

    // Minimal mock Command implementation
    const command: any = { execute: jest.fn() };

    agent.registerAgent("testCommand", command);

    expect(agent.getRegistry()).toEqual(
      new Map([
        ["default", expect.anything()],
        ["testCommand", command],
      ])
    );
  });

  it("should process input and dispatch tasks", async () => {
    const llm = new BedrockLLMClientAdapter();

    llm.invoke = jest.fn().mockResolvedValue({
      message: {
        role: "assistant",
        content: JSON.stringify({
          agent: "testCommand",
          command: "task result",
        }),
      },
      usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 },
      raw: {},
    });

    const agent = new cut.Agent(llm, {
      defaultMemory: [{ content: "default memory" }],
    });

    const command: any = {
      execute: jest.fn().mockResolvedValue({
        message: {
          content: "task result",
          role: "assistant"
        }
      }),
    };

    agent.registerAgent("testCommand", command);

    const task = {
      agent: "testCommand",
      originalInput: [{ content: "test input", role: "user" }],
      command: "task result"
    };

    const result = await agent.processInput(
      [{ content: "test input", role: "user" }],
      [{ content: "test memory" }]
    );

    expect(result).toEqual({
        testCommand: {
          message: {
            content: "task result",
            role: "assistant"
        }
        }
    });
    expect(command.execute).toHaveBeenCalledWith(task);
  });
});