import { describe, it, expect, beforeEach, vi, type Mock } from "vitest";
import { OllamaProvider, type OllamaProviderConfig } from "./ollama.provider";
import type { ChatRequestDto, EmbedRequestDto } from "../dto";

// Mock the ollama module
vi.mock("ollama", () => {
  return {
    Ollama: vi.fn().mockImplementation(function (this: unknown) {
      return {
        list: vi.fn(),
        chat: vi.fn(),
        embed: vi.fn(),
      };
    }),
  };
});

describe("OllamaProvider", () => {
  let provider: OllamaProvider;
  let config: OllamaProviderConfig;
  let mockOllamaInstance: {
    list: Mock;
    chat: Mock;
    embed: Mock;
  };

  beforeEach(() => {
    // Reset all mocks
    vi.clearAllMocks();

    // Setup test configuration
    config = {
      endpoint: "http://localhost:11434",
      timeout: 30000,
    };

    provider = new OllamaProvider(config);

    // Get the mock instance created by the constructor
    mockOllamaInstance = (provider as any).client;
  });

  describe("constructor and initialization", () => {
    it("should create provider with correct name and type", () => {
      expect(provider.name).toBe("Ollama");
      expect(provider.type).toBe("ollama");
    });

    it("should initialize successfully", async () => {
      await expect(provider.initialize()).resolves.toBeUndefined();
    });
  });

  describe("checkHealth", () => {
    it("should return healthy status when Ollama is reachable", async () => {
      const mockModels = [{ name: "llama2" }, { name: "mistral" }];
      mockOllamaInstance.list.mockResolvedValue({ models: mockModels });

      const health = await provider.checkHealth();

      expect(health).toEqual({
        healthy: true,
        provider: "ollama",
        endpoint: config.endpoint,
        models: ["llama2", "mistral"],
      });
      expect(mockOllamaInstance.list).toHaveBeenCalledOnce();
    });

    it("should return unhealthy status when Ollama is unreachable", async () => {
      const error = new Error("Connection refused");
      mockOllamaInstance.list.mockRejectedValue(error);

      const health = await provider.checkHealth();

      expect(health).toEqual({
        healthy: false,
        provider: "ollama",
        endpoint: config.endpoint,
        error: "Connection refused",
      });
    });

    it("should handle non-Error exceptions", async () => {
      mockOllamaInstance.list.mockRejectedValue("string error");

      const health = await provider.checkHealth();

      expect(health.healthy).toBe(false);
      expect(health.error).toBe("string error");
    });
  });

  describe("listModels", () => {
    it("should return array of model names", async () => {
      const mockModels = [{ name: "llama2" }, { name: "mistral" }, { name: "codellama" }];
      mockOllamaInstance.list.mockResolvedValue({ models: mockModels });

      const models = await provider.listModels();

      expect(models).toEqual(["llama2", "mistral", "codellama"]);
      expect(mockOllamaInstance.list).toHaveBeenCalledOnce();
    });

    it("should throw error when listing models fails", async () => {
      const error = new Error("Failed to connect");
      mockOllamaInstance.list.mockRejectedValue(error);

      await expect(provider.listModels()).rejects.toThrow("Failed to list models");
    });
  });

  describe("chat", () => {
    it("should perform chat completion successfully", async () => {
      const request: ChatRequestDto = {
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
      };

      const mockResponse = {
        model: "llama2",
        message: { role: "assistant", content: "Hi there!" },
        done: true,
        total_duration: 1000000,
        prompt_eval_count: 10,
        eval_count: 5,
      };

      mockOllamaInstance.chat.mockResolvedValue(mockResponse);

      const response = await provider.chat(request);

      expect(response).toEqual({
        model: "llama2",
        message: { role: "assistant", content: "Hi there!" },
        done: true,
        totalDuration: 1000000,
        promptEvalCount: 10,
        evalCount: 5,
      });

      expect(mockOllamaInstance.chat).toHaveBeenCalledWith({
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
        stream: false,
        options: {},
      });
    });

    it("should include system prompt in messages", async () => {
      const request: ChatRequestDto = {
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
        systemPrompt: "You are a helpful assistant",
      };

      mockOllamaInstance.chat.mockResolvedValue({
        model: "llama2",
        message: { role: "assistant", content: "Hi!" },
        done: true,
      });

      await provider.chat(request);

      expect(mockOllamaInstance.chat).toHaveBeenCalledWith({
        model: "llama2",
        messages: [
          { role: "system", content: "You are a helpful assistant" },
          { role: "user", content: "Hello" },
        ],
        stream: false,
        options: {},
      });
    });

    it("should not duplicate system prompt when already in messages", async () => {
      const request: ChatRequestDto = {
        model: "llama2",
        messages: [
          { role: "system", content: "Existing system prompt" },
          { role: "user", content: "Hello" },
        ],
        systemPrompt: "New system prompt (should be ignored)",
      };

      mockOllamaInstance.chat.mockResolvedValue({
        model: "llama2",
        message: { role: "assistant", content: "Hi!" },
        done: true,
      });

      await provider.chat(request);

      expect(mockOllamaInstance.chat).toHaveBeenCalledWith({
        model: "llama2",
        messages: [
          { role: "system", content: "Existing system prompt" },
          { role: "user", content: "Hello" },
        ],
        stream: false,
        options: {},
      });
    });

    it("should pass temperature and maxTokens as options", async () => {
      const request: ChatRequestDto = {
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
        temperature: 0.7,
        maxTokens: 100,
      };

      mockOllamaInstance.chat.mockResolvedValue({
        model: "llama2",
        message: { role: "assistant", content: "Hi!" },
        done: true,
      });

      await provider.chat(request);

      expect(mockOllamaInstance.chat).toHaveBeenCalledWith({
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
        stream: false,
        options: {
          temperature: 0.7,
          num_predict: 100,
        },
      });
    });

    it("should throw error when chat fails", async () => {
      const request: ChatRequestDto = {
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
      };

      mockOllamaInstance.chat.mockRejectedValue(new Error("Model not found"));

      await expect(provider.chat(request)).rejects.toThrow("Chat completion failed");
    });
  });

  describe("chatStream", () => {
    it("should stream chat completion chunks", async () => {
      const request: ChatRequestDto = {
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
      };

      const mockChunks = [
        { model: "llama2", message: { role: "assistant", content: "Hi" }, done: false },
        { model: "llama2", message: { role: "assistant", content: " there" }, done: false },
        { model: "llama2", message: { role: "assistant", content: "!" }, done: true },
      ];

      // Mock async generator
      async function* mockStreamGenerator() {
        for (const chunk of mockChunks) {
          yield chunk;
        }
      }

      mockOllamaInstance.chat.mockResolvedValue(mockStreamGenerator());

      const chunks = [];
      for await (const chunk of provider.chatStream(request)) {
        chunks.push(chunk);
      }

      expect(chunks).toHaveLength(3);
      expect(chunks[0]).toEqual({
        model: "llama2",
        message: { role: "assistant", content: "Hi" },
        done: false,
      });
      expect(chunks[2].done).toBe(true);

      expect(mockOllamaInstance.chat).toHaveBeenCalledWith({
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
        stream: true,
        options: {},
      });
    });

    it("should pass options in streaming mode", async () => {
      const request: ChatRequestDto = {
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
        temperature: 0.5,
        maxTokens: 50,
      };

      async function* mockStreamGenerator() {
        yield { model: "llama2", message: { role: "assistant", content: "Hi" }, done: true };
      }

      mockOllamaInstance.chat.mockResolvedValue(mockStreamGenerator());

      const generator = provider.chatStream(request);
      await generator.next();

      expect(mockOllamaInstance.chat).toHaveBeenCalledWith({
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
        stream: true,
        options: {
          temperature: 0.5,
          num_predict: 50,
        },
      });
    });

    it("should throw error when streaming fails", async () => {
      const request: ChatRequestDto = {
        model: "llama2",
        messages: [{ role: "user", content: "Hello" }],
      };

      mockOllamaInstance.chat.mockRejectedValue(new Error("Stream error"));

      const generator = provider.chatStream(request);

      await expect(generator.next()).rejects.toThrow("Streaming failed");
    });
  });

  describe("embed", () => {
    it("should generate embeddings successfully", async () => {
      const request: EmbedRequestDto = {
        model: "nomic-embed-text",
        input: ["Hello world", "Test embedding"],
      };

      const mockResponse = {
        model: "nomic-embed-text",
        embeddings: [
          [0.1, 0.2, 0.3],
          [0.4, 0.5, 0.6],
        ],
        total_duration: 500000,
      };

      mockOllamaInstance.embed.mockResolvedValue(mockResponse);

      const response = await provider.embed(request);

      expect(response).toEqual({
        model: "nomic-embed-text",
        embeddings: [
          [0.1, 0.2, 0.3],
          [0.4, 0.5, 0.6],
        ],
        totalDuration: 500000,
      });

      expect(mockOllamaInstance.embed).toHaveBeenCalledWith({
        model: "nomic-embed-text",
        input: ["Hello world", "Test embedding"],
        truncate: true,
      });
    });

    it("should handle truncate option", async () => {
      const request: EmbedRequestDto = {
        model: "nomic-embed-text",
        input: ["Test"],
        truncate: "start",
      };

      mockOllamaInstance.embed.mockResolvedValue({
        model: "nomic-embed-text",
        embeddings: [[0.1, 0.2]],
      });

      await provider.embed(request);

      expect(mockOllamaInstance.embed).toHaveBeenCalledWith({
        model: "nomic-embed-text",
        input: ["Test"],
        truncate: true,
      });
    });

    it("should respect truncate none option", async () => {
      const request: EmbedRequestDto = {
        model: "nomic-embed-text",
        input: ["Test"],
        truncate: "none",
      };

      mockOllamaInstance.embed.mockResolvedValue({
        model: "nomic-embed-text",
        embeddings: [[0.1, 0.2]],
      });

      await provider.embed(request);

      expect(mockOllamaInstance.embed).toHaveBeenCalledWith({
        model: "nomic-embed-text",
        input: ["Test"],
        truncate: false,
      });
    });

    it("should throw error when embedding fails", async () => {
      const request: EmbedRequestDto = {
        model: "nomic-embed-text",
        input: ["Test"],
      };

      mockOllamaInstance.embed.mockRejectedValue(new Error("Embedding error"));

      await expect(provider.embed(request)).rejects.toThrow("Embedding failed");
    });
  });

  describe("getConfig", () => {
    it("should return copy of configuration", () => {
      const returnedConfig = provider.getConfig();

      expect(returnedConfig).toEqual(config);
      expect(returnedConfig).not.toBe(config); // Should be a copy, not reference
    });

    it("should prevent external modification of config", () => {
      const returnedConfig = provider.getConfig();
      returnedConfig.endpoint = "http://modified:11434";

      const secondCall = provider.getConfig();
      expect(secondCall.endpoint).toBe("http://localhost:11434"); // Original unchanged
    });
  });
});
