import {
  BedrockRuntimeClient,
  InvokeModelCommand,
  InvokeModelCommandOutput,
} from "@aws-sdk/client-bedrock-runtime";
import { config } from "dotenv";
import { getAwsCredentials } from "../utils/aws-credential";
import { NodeHttpHandler } from "@aws-sdk/node-http-handler";

config();

const MODEL_ID: string = process.env.AI_MODEL || "amazon.nova-lite-v1:0";

let inferenceConfig: InferenceConfig = {
  maxTokens: 5000,
  topP: 0.9,
  topK: 20,
  temperature: 0.7,
};

// Override inference configuration from .env if provided
if (process.env.INFERENCE_CONFIG) {
  try {
    inferenceConfig = JSON.parse(process.env.INFERENCE_CONFIG);
  } catch (error) {
    console.error("Failed to parse INFERENCE_CONFIG:", error);
  }
}

/** @ts-ignore */
const runtimeClient = new BedrockRuntimeClient({
  ...getAwsCredentials(),
  requestHandler: new NodeHttpHandler({
    connectionTimeout: 3600000, // 60 minutes
    socketTimeout: 3600000, // 60 minutes
  }),
});

export async function invokeModel(
  messages: UserMessage[],
  memories: Memory[]
): Promise<InvokeModelCommandOutput> {
  // Build the request payload according to the Amazon Nova Lite schema
  const requestBody: RequestPayload = {
    schemaVersion: "messages-v1",
    system: memories,
    messages,
    inferenceConfig,
  };

  const params = {
    modelId: MODEL_ID,
    contentType: "application/json",
    accept: "application/json",
    body: JSON.stringify(requestBody),
  };

  try {
    const command = new InvokeModelCommand(params);

    return await runtimeClient.send(command);
  } catch (error) {
    console.error(`Something is wrong with the model invoke, ${error}`);
    throw error; // Ensure a value is always returned or an error is thrown
  }
}

interface InferenceConfig {
  maxTokens: number;
  topP: number;
  topK: number;
  temperature: number;
}

export interface Message {
  text: string;
}

export interface UserMessage {
  role: "user" | "assistant";
  content: Message[];
}

interface RequestPayload {
  schemaVersion: string;
  system: Memory[];
  messages: UserMessage[];
  inferenceConfig: InferenceConfig;
}

export interface Memory {
  text: string;
}

export type BedrockInvokeModelCommandOutput = InvokeModelCommandOutput;