import { GoogleGenAI } from "@google/genai";
import { describe, it, expect, beforeAll, afterAll } from "vitest";

import { startServer, stopServer, getBaseUrl } from "./setup";

function getGeminiClient(): GoogleGenAI {
  return new GoogleGenAI({
    vertexai: false,
    apiKey: "test-key",
    httpOptions: {
      baseUrl: getBaseUrl(),
    },
  });
}

describe("Google GenAI (Gemini) API", () => {
  beforeAll(async () => {
    await startServer();
  });

  afterAll(async () => {
    await stopServer();
  });

  describe("Non-streaming", () => {
    it("should return a valid generateContent response", async () => {
      const ai = getGeminiClient();

      const response = await ai.models.generateContent({
        model: "gemini-2.0-flash",
        contents: "Hello!",
      });

      expect(response.responseId).toBeDefined();
      expect(response.modelVersion).toBe("gemini-2.0-flash");
      expect(response.candidates).toHaveLength(1);
      expect(response.candidates![0].content.role).toBe("model");
      expect(response.candidates![0].content.parts).toHaveLength(1);
      expect(response.text).toBeTruthy();
      expect(response.usageMetadata).toBeDefined();
      expect(response.usageMetadata!.promptTokenCount).toBeGreaterThan(0);
      expect(response.usageMetadata!.candidatesTokenCount).toBeGreaterThan(0);
    });

    it("should handle multiple messages", async () => {
      const ai = getGeminiClient();

      const response = await ai.models.generateContent({
        model: "gemini-2.0-flash",
        contents: [
          { role: "user", parts: [{ text: "Hello!" }] },
          { role: "model", parts: [{ text: "Hi there!" }] },
          { role: "user", parts: [{ text: "How are you?" }] },
        ],
      });

      expect(response.text).toBeTruthy();
    });

    it("should handle system instruction", async () => {
      const ai = getGeminiClient();

      const response = await ai.models.generateContent({
        model: "gemini-2.0-flash",
        contents: "Hello!",
        config: {
          systemInstruction: "You are a helpful assistant.",
        },
      });

      expect(response.text).toBeTruthy();
    });

    it("should reject requests without contents", async () => {
      const response = await fetch(
        `${getBaseUrl()}/v1beta/models/gemini-2.0-flash:generateContent`,
        {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
            "x-goog-api-key": "test-key",
          },
          body: JSON.stringify({}),
        }
      );

      expect(response.status).toBe(400);
      const data = await response.json();
      expect(data.error).toBeDefined();
      expect(data.error.status).toBe("INVALID_ARGUMENT");
    });
  });

  describe("Streaming", () => {
    it("should stream a valid generateContent response", async () => {
      const ai = getGeminiClient();

      const response = await ai.models.generateContentStream({
        model: "gemini-2.0-flash",
        contents: "Hello!",
      });

      const chunks: string[] = [];
      let hasFinishReason = false;

      for await (const chunk of response) {
        if (chunk.text) {
          chunks.push(chunk.text);
        }
        if (chunk.candidates?.[0]?.finishReason === "STOP") {
          hasFinishReason = true;
        }
      }

      expect(chunks.length).toBeGreaterThan(0);
      expect(hasFinishReason).toBe(true);

      const fullContent = chunks.join("");
      expect(fullContent).toBeTruthy();
      expect(fullContent.length).toBeGreaterThan(10);
    });

    it("should include usage metadata in streaming response", async () => {
      const ai = getGeminiClient();

      const response = await ai.models.generateContentStream({
        model: "gemini-2.0-flash",
        contents: "Hello!",
      });

      let hasUsage = false;

      for await (const chunk of response) {
        if (chunk.usageMetadata) {
          hasUsage = true;
          expect(chunk.usageMetadata.promptTokenCount).toBeGreaterThan(0);
        }
      }

      expect(hasUsage).toBe(true);
    });
  });

  describe("Function Calling", () => {
    it("should return function call in non-streaming mode", async () => {
      // Start server with 100% tool call probability
      await stopServer();
      await startServer({ TOOL_CALL_PROBABILITY: "1.0" });

      const ai = getGeminiClient();

      const response = await ai.models.generateContent({
        model: "gemini-2.0-flash",
        contents: "What's the weather in San Francisco?",
        config: {
          tools: [
            {
              functionDeclarations: [
                {
                  name: "get_weather",
                  description: "Get the weather for a location",
                  parameters: {
                    type: "object",
                    properties: {
                      location: { type: "string", description: "City name" },
                      unit: { type: "string", enum: ["celsius", "fahrenheit"] },
                    },
                    required: ["location"],
                  },
                },
              ],
            },
          ],
        },
      });

      expect(response.functionCalls).toBeDefined();
      expect(response.functionCalls!.length).toBeGreaterThan(0);

      const functionCall = response.functionCalls![0];
      expect(functionCall.name).toBe("get_weather");
      expect(functionCall.args).toBeDefined();
      expect(typeof functionCall.args).toBe("object");
    });

    it("should stream function call correctly", async () => {
      // Ensure we have 100% tool call probability
      await stopServer();
      await startServer({ TOOL_CALL_PROBABILITY: "1.0" });

      const ai = getGeminiClient();

      const response = await ai.models.generateContentStream({
        model: "gemini-2.0-flash",
        contents: "What's the weather?",
        config: {
          tools: [
            {
              functionDeclarations: [
                {
                  name: "get_weather",
                  description: "Get the weather for a location",
                  parameters: {
                    type: "object",
                    properties: {
                      location: { type: "string" },
                    },
                    required: ["location"],
                  },
                },
              ],
            },
          ],
        },
      });

      let hasFunctionCall = false;
      let functionName: string | null = null;

      for await (const chunk of response) {
        if (chunk.functionCalls && chunk.functionCalls.length > 0) {
          hasFunctionCall = true;
          functionName = chunk.functionCalls[0].name;
        }
      }

      expect(hasFunctionCall).toBe(true);
      expect(functionName).toBe("get_weather");
    });

    it("should respect toolConfig mode NONE", async () => {
      const ai = getGeminiClient();

      const response = await ai.models.generateContent({
        model: "gemini-2.0-flash",
        contents: "What's the weather?",
        config: {
          tools: [
            {
              functionDeclarations: [
                {
                  name: "get_weather",
                  description: "Get the weather",
                  parameters: {
                    type: "object",
                    properties: {
                      location: { type: "string" },
                    },
                  },
                },
              ],
            },
          ],
          toolConfig: {
            functionCallingConfig: {
              mode: "NONE",
            },
          },
        },
      });

      // Should return text, not function call
      expect(response.text).toBeTruthy();
      expect(response.functionCalls).toBeUndefined();
    });
  });
});
