import OpenAI from "openai";
import type { EmbeddingProvider } from "./embedding-provider.js";

export type OpenAIEmbeddingModel =
	| "text-embedding-3-small"
	| "text-embedding-3-large";

export type OpenAIEmbeddingProviderConfiguration = {
	apiKey: string;
	model?: OpenAIEmbeddingModel;
};

const MODEL_DIMENSIONS: Record<OpenAIEmbeddingModel, number> = {
	"text-embedding-3-small": 1536,
	"text-embedding-3-large": 3072,
};

export class OpenAIEmbeddingProvider implements EmbeddingProvider {
	private client: OpenAI;
	private model: OpenAIEmbeddingModel;
	readonly dimensions: number;

	constructor(configuration: OpenAIEmbeddingProviderConfiguration) {
		this.client = new OpenAI({ apiKey: configuration.apiKey });
		this.model = configuration.model ?? "text-embedding-3-small";
		this.dimensions = MODEL_DIMENSIONS[this.model];
	}

	async generate(text: string): Promise<number[]> {
		const response = await this.client.embeddings.create({
			model: this.model,
			input: text,
		});

		return response.data[0]?.embedding ?? [];
	}

	async generateBatch(texts: string[]): Promise<number[][]> {
		const response = await this.client.embeddings.create({
			model: this.model,
			input: texts,
		});

		return response.data.map((item) => item.embedding);
	}
}