import { describe, it, expect, vi, beforeEach } from 'vitest';
import type { ChatResponseStream } from 'ollama';

// Mock the Ollama client BEFORE importing the module
const mockOllamaClient = {
  chat: vi.fn(),
  generate: vi.fn(),
  list: vi.fn(),
  show: vi.fn(),
};

vi.mock('ollama', () => {
  return {
    Ollama: class MockOllama {
      constructor() {
        return mockOllamaClient;
      }
    },
  };
});

import { OllamaAiApi } from './ollama';

// Mock logger
const mockLogger = {
  debug: vi.fn(),
  info: vi.fn(),
  warn: vi.fn(),
  error: vi.fn(),
};

// Mock environment and provider
const mockEnv = {
  NODE_ENV: 'test',
  services: {
    google: {
      cse: {
        apiKey: 'test-key',
        engineId: 'test-engine',
      },
    },
  },
};

const mockProvider = {
  _id: 'test-provider',
  name: 'Test Ollama',
  sdk: 'ollama' as const,
  baseUrl: 'http://localhost:11434',
  apiKey: 'test-key',
};

describe('OllamaAiApi', () => {
  let api: OllamaAiApi;

  beforeEach(() => {
    vi.clearAllMocks();
    api = new OllamaAiApi(mockEnv as any, mockProvider as any, mockLogger as any);
  });

  describe('chat', () => {
    it('should handle normal response streaming', async () => {
      // Mock streaming response
      const mockStream = async function* () {
        yield {
          message: { content: 'Hello' },
          done: false,
        };
        yield {
          message: { content: ' world' },
          done: false,
        };
        yield {
          message: { content: '!' },
          done: true,
          done_reason: 'stop',
          total_duration: 100,
          prompt_eval_count: 10,
          eval_count: 3,
        };
      };

      mockOllamaClient.chat.mockResolvedValue(mockStream());

      const streamCallback = vi.fn();
      const response = await api.chat(
        {
          provider: mockProvider as any,
          modelId: 'test-model',
          params: { reasoning: false, temperature: 0.8, topP: 0.9, topK: 40 },
        },
        {
          userPrompt: 'Test prompt',
          context: [],
        },
        streamCallback,
      );

      // Verify stream callback was called for each chunk
      expect(streamCallback).toHaveBeenCalledTimes(3);
      expect(streamCallback).toHaveBeenCalledWith({
        type: 'response',
        data: 'Hello',
      });
      expect(streamCallback).toHaveBeenCalledWith({
        type: 'response',
        data: ' world',
      });
      expect(streamCallback).toHaveBeenCalledWith({
        type: 'response',
        data: '!',
      });

      // Verify response
      expect(response.done).toBe(true);
      expect(response.doneReason).toBe('stop');
      expect(response.response).toBe('!');
    });

    it('should handle tool calls', async () => {
      // Mock streaming response with tool call
      const mockStream = async function* () {
        yield {
          message: {
            content: '',
            tool_calls: [
              {
                function: {
                  name: 'search_google',
                  arguments: { query: 'test query' },
                },
              },
            ],
          },
          done: false,
        };
        yield {
          message: { content: '' },
          done: true,
          done_reason: 'stop',
          total_duration: 100,
          prompt_eval_count: 10,
          eval_count: 1,
        };
      };

      mockOllamaClient.chat.mockResolvedValue(mockStream());

      const mockTool = {
        definition: {
          type: 'function',
          function: {
            name: 'search_google',
            description: 'Search Google',
            parameters: { type: 'object', properties: {} },
          },
        },
        execute: vi.fn().mockResolvedValue({ result: 'search results' }),
      };

      const streamCallback = vi.fn();
      const response = await api.chat(
        {
          provider: mockProvider as any,
          modelId: 'test-model',
          params: { reasoning: false, temperature: 0.8, topP: 0.9, topK: 40 },
        },
        {
          userPrompt: 'Test prompt',
          context: [],
          tools: [mockTool as any],
        },
        streamCallback,
      );

      // Verify tool call was emitted via stream callback
      expect(streamCallback).toHaveBeenCalledWith(
        expect.objectContaining({
          type: 'toolCall',
          toolName: 'search_google',
        }),
      );

      // Verify tool was executed
      expect(mockTool.execute).toHaveBeenCalled();

      // Verify response indicates tool calls were processed
      expect(response.toolCalls).toBeDefined();
    });

    it('should handle thinking content when reasoning is enabled', async () => {
      // Mock streaming response with thinking
      const mockStream = async function* () {
        yield {
          message: {
            thinking: 'Let me think about this...',
            content: '',
          },
          done: false,
        };
        yield {
          message: {
            thinking: ' The answer is',
            content: '',
          },
          done: false,
        };
        yield {
          message: { content: '42' },
          done: true,
          done_reason: 'stop',
          total_duration: 100,
          prompt_eval_count: 10,
          eval_count: 1,
        };
      };

      mockOllamaClient.chat.mockResolvedValue(mockStream());

      const streamCallback = vi.fn();
      const response = await api.chat(
        {
          provider: mockProvider as any,
          modelId: 'test-model',
          params: { reasoning: true, temperature: 0.8, topP: 0.9, topK: 40 },
        },
        {
          userPrompt: 'What is the answer?',
          context: [],
        },
        streamCallback,
      );

      // Verify thinking was emitted
      expect(streamCallback).toHaveBeenCalledWith({
        type: 'thinking',
        data: 'Let me think about this...',
      });
      expect(streamCallback).toHaveBeenCalledWith({
        type: 'thinking',
        data: ' The answer is',
      });

      // Verify response was emitted
      expect(streamCallback).toHaveBeenCalledWith({
        type: 'response',
        data: '42',
      });

      // Verify final response includes thinking
      expect(response.thinking).toBe('Let me think about this... The answer is');
    });

    it('should handle empty response on load failure', async () => {
      // Mock streaming response with load failure
      const mockStream = async function* () {
        yield {
          message: { content: '' },
          done: true,
          done_reason: 'load',
          total_duration: 5000,
          prompt_eval_count: 0,
          eval_count: 0,
        };
      };

      mockOllamaClient.chat.mockResolvedValue(mockStream());

      const streamCallback = vi.fn();
      const response = await api.chat(
        {
          provider: mockProvider as any,
          modelId: 'test-model',
          params: { reasoning: false, temperature: 0.8, topP: 0.9, topK: 40 },
        },
        {
          userPrompt: 'Test prompt',
          context: [],
        },
        streamCallback,
      );

      // Verify response indicates load failure
      expect(response.done).toBe(true);
      expect(response.doneReason).toBe('load');
      expect(response.response).toBe('');

      // Verify no stream callbacks for empty content
      expect(streamCallback).not.toHaveBeenCalled();
    });

    it('should iterate tool calling loop when tools are present', async () => {
      let callCount = 0;
      
      // Mock streaming response that requires tool call then returns
      const mockStream = async function* () {
        callCount++;
        if (callCount === 1) {
          // First call: return tool call
          yield {
            message: {
              content: '',
              tool_calls: [
                {
                  function: {
                    name: 'search_google',
                    arguments: { query: 'test' },
                  },
                },
              ],
            },
            done: false,
          };
          yield {
            message: { content: '' },
            done: true,
            done_reason: 'stop',
            total_duration: 100,
            prompt_eval_count: 10,
            eval_count: 1,
          };
        } else {
          // Second call: return final response
          yield {
            message: { content: 'Here are the results' },
            done: true,
            done_reason: 'stop',
            total_duration: 100,
            prompt_eval_count: 15,
            eval_count: 5,
          };
        }
      };

      mockOllamaClient.chat.mockImplementation(() => mockStream());

      const mockTool = {
        definition: {
          type: 'function',
          function: {
            name: 'search_google',
            description: 'Search Google',
            parameters: { type: 'object', properties: {} },
          },
        },
        execute: vi.fn().mockResolvedValue({ result: 'search results' }),
      };

      const streamCallback = vi.fn();
      const response = await api.chat(
        {
          provider: mockProvider as any,
          modelId: 'test-model',
          params: { reasoning: false, temperature: 0.8, topP: 0.9, topK: 40 },
        },
        {
          userPrompt: 'Test prompt',
          context: [],
          tools: [mockTool as any],
        },
        streamCallback,
      );

      // Verify chat was called twice (once for tool call, once for response)
      expect(mockOllamaClient.chat).toHaveBeenCalledTimes(2);

      // Verify tool was executed
      expect(mockTool.execute).toHaveBeenCalled();

      // Verify final response
      expect(response.done).toBe(true);
      expect(response.response).toBe('Here are the results');
    });
  });
});
