import type { LlmAdapter } from "./types.js";

interface GeminiResponse {
  candidates: Array<{ content: { parts: Array<{ text: string }> } }>;
  usageMetadata: {
    promptTokenCount: number;
    candidatesTokenCount: number;
    totalTokenCount: number;
  };
}

export class GeminiAdapter implements LlmAdapter {
  private apiKey: string;
  private model: string;

  constructor(apiKey: string, model: string) {
    this.apiKey = apiKey;
    this.model = model;
  }

  async call(
    words: string[],
    systemPrompt: string,
  ): Promise<{
    content: string;
    promptTokens: number;
    completionTokens: number;
    totalTokens: number;
    promptTimeMs: number;
    completionTimeMs: number;
  }> {
    const url = `https://generativelanguage.googleapis.com/v1beta/models/${this.model}:generateContent?key=${this.apiKey}`;

    const payload = {
      contents: [
        {
          role: "user",
          parts: [
            { text: systemPrompt + "\n\nWords: " + JSON.stringify(words) },
          ],
        },
      ],
      generationConfig: {
        temperature: 0.1,
        topP: 0.9,
        maxOutputTokens: Math.ceil(words.length * 250 * 1.2),
      },
    };

    const startTime = Date.now();

    const response = await fetch(url, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(payload),
    });

    const totalTimeMs = Date.now() - startTime;

    if (!response.ok) {
      throw new Error(`Gemini API responded with status: ${response.status}`);
    }

    const json = (await response.json()) as GeminiResponse;

    const content = json.candidates[0]?.content?.parts[0]?.text;
    if (!content) {
      throw new Error("Gemini response content is empty");
    }

    const promptTokens = json.usageMetadata.promptTokenCount;
    const completionTokens = json.usageMetadata.candidatesTokenCount;

    return {
      content,
      promptTokens,
      completionTokens,
      totalTokens: json.usageMetadata.totalTokenCount,
      promptTimeMs: totalTimeMs * 0.3,
      completionTimeMs: totalTimeMs * 0.7,
    };
  }
}