/* Copyright 2026 Marimo. All rights reserved. */

import type { AiModel } from "@marimo-team/llm-info";
import { beforeEach, describe, expect, it, vi } from "vitest";
import type { UserConfig } from "@/core/config/config-schema";

vi.mock("@marimo-team/llm-info/models.json", () => {
  const make = (
    overrides: Partial<AiModel> & Pick<AiModel, "name" | "model">,
  ): AiModel => ({
    description: "",
    roles: ["chat", "edit"],
    capabilities: [],
    input_types: [],
    output_types: [],
    release_date: "1970-01-01",
    ...overrides,
  });

  const models: Record<string, AiModel[]> = {
    openai: [
      make({
        name: "GPT-4",
        model: "gpt-4",
        description: "OpenAI GPT-4 model",
      }),
    ],
    anthropic: [
      make({
        name: "Claude 3",
        model: "claude-3-sonnet",
        description: "Anthropic Claude 3 Sonnet",
      }),
    ],
    google: [
      make({
        name: "Gemini Pro",
        model: "gemini-pro",
        description: "Google Gemini Pro model",
      }),
    ],
    ollama: [
      make({
        name: "Ollama Model",
        model: "llama2",
        description: "Ollama Llama 2 model",
      }),
    ],
    marimo: [
      make({
        name: "Best for chat",
        model: "best-for-chat",
        description: "marimo hosted chat model",
      }),
    ],
  };

  return { models };
});

// Must import after mock
import {
  autoPopulateModels,
  getRecommendedModel,
  listConfiguredProviders,
  listModelsForAiSettings,
} from "../ai-utils";

describe("ai-utils", () => {
  beforeEach(() => {
    vi.clearAllMocks();
  });

  describe("listConfiguredProviders", () => {
    it("returns known providers with credentials, then custom providers with a base URL", () => {
      expect(
        listConfiguredProviders({
          wandb: { api_key: "unused" },
          open_ai: { api_key: "sk-test" },
          open_ai_compatible: { api_key: "unused", base_url: "https://x" },
          custom_providers: {
            groq: { base_url: "https://api.groq.com" },
            empty: {},
          },
        }),
      ).toEqual(["openai", "wandb", "marimo", "groq"]);
    });

    it("treats marimo as configured when open_ai_compatible has a base URL", () => {
      expect(
        listConfiguredProviders({
          open_ai_compatible: { api_key: "unused", base_url: "https://x" },
        }),
      ).toEqual(["marimo"]);
    });

    it("omits openai-compatible as its own provider id", () => {
      expect(
        listConfiguredProviders({
          open_ai: {},
          open_ai_compatible: { api_key: "unused" },
        }),
      ).toEqual([]);
    });
  });

  describe("listModelsForAiSettings", () => {
    const entries: [string, number][] = [
      ["openai", 1],
      ["anthropic", 2],
      ["wandb", 3],
      ["marimo", 4],
    ];

    it("returns all groups when provider config is allowed", () => {
      expect(
        listModelsForAiSettings(entries, {
          allow_provider_config: true,
          wandb: { api_key: "unused" },
        }),
      ).toEqual(entries);
    });

    it("filters to configured providers when provider config is locked", () => {
      expect(
        listModelsForAiSettings(entries, {
          allow_provider_config: false,
          wandb: { api_key: "unused" },
        }),
      ).toEqual([["wandb", 3]]);
    });

    it("keeps marimo models when locked with an openai-compatible gateway", () => {
      expect(
        listModelsForAiSettings(entries, {
          allow_provider_config: false,
          open_ai_compatible: { base_url: "https://ai.marimo.app/molab/v1" },
        }),
      ).toEqual([["marimo", 4]]);
    });
  });

  describe("getRecommendedModel", () => {
    it("should return undefined when no provider is configured", () => {
      const config: UserConfig = {} as UserConfig;
      expect(getRecommendedModel(config.ai)).toBeUndefined();
    });

    it("should return openai model when OpenAI is configured", () => {
      const config: UserConfig = {
        ai: {
          open_ai: { api_key: "sk-test" },
        },
      } as UserConfig;
      expect(getRecommendedModel(config.ai)).toBe("openai/gpt-4");
    });

    it("should return anthropic model when Anthropic is configured", () => {
      const config: UserConfig = {
        ai: {
          anthropic: { api_key: "sk-ant-test" },
        },
      } as UserConfig;
      expect(getRecommendedModel(config.ai)).toBe("anthropic/claude-3-sonnet");
    });

    it("should return google model when Google is configured", () => {
      const config: UserConfig = {
        ai: {
          google: { api_key: "google-key" },
        },
      } as UserConfig;
      expect(getRecommendedModel(config.ai)).toBe("google/gemini-pro");
    });

    it("should return ollama model when Ollama is configured", () => {
      const config: UserConfig = {
        ai: {
          ollama: { base_url: "http://localhost:11434" },
        },
      } as UserConfig;
      expect(getRecommendedModel(config.ai)).toBe("ollama/llama2");
    });
  });

  describe("autoPopulateModels", () => {
    it("should return empty result when both models are already set", () => {
      const values = {
        ai: {
          open_ai: { api_key: "sk-test" },
          models: {
            chat_model: "openai/gpt-4",
            edit_model: "openai/gpt-4",
            custom_models: [],
            displayed_models: [],
          },
        },
      } as unknown as UserConfig;

      const result = autoPopulateModels(values.ai);

      expect(result.chatModel).toBeUndefined();
      expect(result.editModel).toBeUndefined();
    });

    it("should return empty result when no credentials are configured", () => {
      const values: UserConfig = {
        ai: {},
      } as UserConfig;

      const result = autoPopulateModels(values.ai);

      expect(result.chatModel).toBeUndefined();
      expect(result.editModel).toBeUndefined();
    });

    it("should auto-populate both models when neither is set", () => {
      const values: UserConfig = {
        ai: {
          open_ai: { api_key: "sk-test" },
        },
      } as UserConfig;

      const result = autoPopulateModels(values.ai);

      expect(result.chatModel).toBe("openai/gpt-4");
      expect(result.editModel).toBe("openai/gpt-4");
    });

    it("should only auto-populate chat_model when edit_model is set", () => {
      const values = {
        ai: {
          open_ai: { api_key: "sk-test" },
          models: {
            edit_model: "openai/gpt-3.5-turbo",
            custom_models: [],
            displayed_models: [],
          },
        },
      } as unknown as UserConfig;

      const result = autoPopulateModels(values.ai);

      expect(result.chatModel).toBe("openai/gpt-4");
      expect(result.editModel).toBeUndefined();
    });

    it("should only auto-populate edit_model when chat_model is set", () => {
      const values = {
        ai: {
          open_ai: { api_key: "sk-test" },
          models: {
            chat_model: "openai/gpt-3.5-turbo",
            custom_models: [],
            displayed_models: [],
          },
        },
      } as unknown as UserConfig;

      const result = autoPopulateModels(values.ai);

      expect(result.chatModel).toBeUndefined();
      expect(result.editModel).toBe("openai/gpt-4");
    });

    it("should return recommended model for anthropic provider", () => {
      const values: UserConfig = {
        ai: {
          anthropic: { api_key: "sk-ant-test" },
        },
      } as UserConfig;

      const result = autoPopulateModels(values.ai);

      expect(result.chatModel).toBe("anthropic/claude-3-sonnet");
      expect(result.editModel).toBe("anthropic/claude-3-sonnet");
    });
  });
});