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 | null;
    completionTimeMs: number | null;
    totalTimeMs: number;
  }> {
    const url = `https://generativelanguage.googleapis.com/v1beta/models/${this.model}:generateContent?key=${this.apiKey}`;

    const payload = {
      systemInstruction: { parts: [{ text: systemPrompt }] },
      contents: [
        { role: "user", parts: [{ text: "Words: " + 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: null,
      completionTimeMs: null,
      totalTimeMs,
    };
  }
}