import { UserMessage } from "../clients/bedrock";
import { UserMessage as GenericUserMessage } from "../llm-invoker";

class BasePayloadBuilder {
  public role: "user" | "assistant" = "user";

  constructor(role?: "user" | "assistant") {
    this.role = role ?? "user";
  }

  setRole(role: "user" | "assistant") {
    this.role = role;
  }
}

export class BedrockPayloadBuilder extends BasePayloadBuilder {
  private messages: UserMessage[] = [
    {
      content: [],
      role: this.role,
    },
  ];

  constructor(role?: "user" | "assistant") {
    super(role);
  }

  addContent(text: string): BedrockPayloadBuilder {
    this.messages[0].role = this.role;
    this.messages[0].content.push({ text });
    return this;
  }

  addContents(texts: string[]): BedrockPayloadBuilder {
    texts.forEach((text) => this.addContent(text));
    return this;
  }

  build(): UserMessage[] {
    return this.messages;
  }
}

export class PayloadBuilder extends BasePayloadBuilder {
  private messages: GenericUserMessage[] = [];

  constructor(role?: "user" | "assistant") {
    super(role);
  }

  addMessage(text: string): PayloadBuilder {
    this.messages.push({
      content: text,
      role: this.role,
    });

    return this;
  }

  addMessages(texts: string[]): PayloadBuilder {
    texts.forEach((text) => this.addMessage(text));
    return this;
  }

  build(): GenericUserMessage[] {
    return this.messages;
  }
}