// BedrockLLMClientAdapter.ts
import { invokeModel } from "../clients/bedrock";
import { BedrockMemoryBuilder } from "../builders/memory-builder";
import { BedrockPayloadBuilder } from "../builders/message-builder";
import {
  LLMInvoker,
  UserMessage,
  Memory,
  InvokeModelCommandOutput,
} from "../llm-invoker";

export class BedrockLLMClientAdapter implements LLMInvoker {
  private memoryBuilder: BedrockMemoryBuilder;
  private messageBuilder: BedrockPayloadBuilder;

  constructor() {
    this.memoryBuilder = new BedrockMemoryBuilder();
    this.messageBuilder = new BedrockPayloadBuilder();
  }

  async invoke(
    messages: UserMessage[],
    memories: Memory[]
  ): Promise<InvokeModelCommandOutput> {
    messages.forEach((message) => {
      this.messageBuilder.setRole(message.role);
      this.messageBuilder.addContent(message.content);
    });

    const bedrockMessage = this.messageBuilder.build();

    memories.forEach((memory) => {
      this.memoryBuilder.addMemory(memory.content);
    });
    const bedrockMemories = this.memoryBuilder.build();

    const rawResponse = await invokeModel(bedrockMessage, bedrockMemories);

    /** Assuming the payload */
    const bodyString: string = new TextDecoder("utf-8").decode(
      rawResponse.body
    );

    if (!bodyString) {
      console.error("Empty response body from Bedrock LLM");
      return {} as InvokeModelCommandOutput;
    }

    const parsedResponse = JSON.parse(bodyString);

    return {
      message: {
        role: parsedResponse.output.message.role as "assistant",
        content: parsedResponse.output.message.content[0].text,
      },
      usage: {
        promptTokens: parsedResponse.usage.promptTokens,
        completionTokens: parsedResponse.usage.completionTokens,
        totalTokens: parsedResponse.usage.totalTokens,
      },
      raw: parsedResponse,
    };
  }
}