import { describe, it, expect, vi, beforeEach } from 'vitest';
import { OllamaClient, normalizeMessagesForOllama } from './ollama.js';
import type { Message } from '../types.js';

const mockChat = vi.fn();
const mockShow = vi.fn();

vi.mock('ollama', () => ({
  Ollama: vi.fn().mockImplementation(() => ({
    chat: mockChat,
    show: mockShow,
  })),
}));

describe('OllamaClient', () => {
  beforeEach(() => {
    mockChat.mockReset();
    mockShow.mockReset();
    // Default: model supports tools
    mockShow.mockResolvedValue({ capabilities: ['completion', 'tools'] });
  });

  it('sends messages and returns response', async () => {
    mockChat.mockResolvedValue({
      message: { content: 'Hello from Ollama!' },
      done_reason: 'stop',
      prompt_eval_count: 10,
      eval_count: 5,
    });

    const client = new OllamaClient({ model: 'llama3.2' });

    const response = await client.chat({
      messages: [{ role: 'user', content: 'Hello' }],
    });

    expect(response.content).toBe('Hello from Ollama!');
    expect(response.stopReason).toBe('stop');
    expect(response.usage.inputTokens).toBe(10);
    expect(response.usage.outputTokens).toBe(5);
  });

  it('passes tools to Ollama API in correct format', async () => {
    mockChat.mockResolvedValue({
      message: { content: 'I can help with that.' },
      done_reason: 'stop',
      prompt_eval_count: 15,
      eval_count: 8,
    });

    const client = new OllamaClient({ model: 'llama3.2' });

    await client.chat({
      messages: [{ role: 'user', content: 'List files' }],
      tools: [
        {
          name: 'shell.exec',
          description: 'Execute a shell command',
          input_schema: {
            type: 'object',
            properties: {
              command: { type: 'string', description: 'The command to run' },
            },
            required: ['command'],
          },
        },
      ],
    });

    expect(mockChat).toHaveBeenCalledWith(
      expect.objectContaining({
        tools: [
          {
            type: 'function',
            function: {
              name: 'shell.exec',
              description: 'Execute a shell command',
              parameters: {
                type: 'object',
                required: ['command'],
                properties: {
                  command: { type: 'string', description: 'The command to run' },
                },
              },
            },
          },
        ],
      }),
    );
  });

  it('parses tool_calls from response', async () => {
    mockChat.mockResolvedValue({
      message: {
        content: '',
        tool_calls: [
          { function: { name: 'shell.exec', arguments: { command: 'ls' } } },
        ],
      },
      done_reason: 'stop',
      prompt_eval_count: 12,
      eval_count: 6,
    });

    const client = new OllamaClient({ model: 'llama3.2' });

    const response = await client.chat({
      messages: [{ role: 'user', content: 'List files' }],
    });

    expect(response.stopReason).toBe('tool_use');
    expect(response.toolCalls).toHaveLength(1);
    expect(response.toolCalls![0]).toEqual({
      id: 'ollama_tc_0',
      name: 'shell.exec',
      args: { command: 'ls' },
    });
  });

  it('handles thinking field from reasoning models', async () => {
    mockChat.mockResolvedValue({
      message: { content: '', thinking: 'Let me think...' },
      done_reason: 'stop',
      prompt_eval_count: 20,
      eval_count: 15,
    });

    const client = new OllamaClient({ model: 'deepseek-r1' });

    const response = await client.chat({
      messages: [{ role: 'user', content: 'Solve this problem' }],
    });

    // When content is empty, thinking is used as fallback
    expect(response.content).toBe('Let me think...');
    expect(response.thinkingContent).toBe('Let me think...');
  });

  it('thinking field does not override existing content', async () => {
    mockChat.mockResolvedValue({
      message: { content: 'Final answer', thinking: 'Reasoning...' },
      done_reason: 'stop',
      prompt_eval_count: 20,
      eval_count: 15,
    });

    const client = new OllamaClient({ model: 'deepseek-r1' });

    const response = await client.chat({
      messages: [{ role: 'user', content: 'Solve this problem' }],
    });

    expect(response.content).toBe('Final answer');
    expect(response.thinkingContent).toBe('Reasoning...');
  });

  it('does not send tools when none provided', async () => {
    mockChat.mockResolvedValue({
      message: { content: 'No tools needed.' },
      done_reason: 'stop',
      prompt_eval_count: 5,
      eval_count: 3,
    });

    const client = new OllamaClient({ model: 'llama3.2' });

    await client.chat({
      messages: [{ role: 'user', content: 'Hello' }],
    });

    const callArgs = mockChat.mock.calls[0][0];
    expect(callArgs.tools).toBeUndefined();
  });

  it('streaming: yields content events', async () => {
    mockChat.mockResolvedValue(
      (async function* () {
        yield { message: { content: 'Hello' }, done: false };
        yield { message: { content: ' world' }, done: false };
        yield { message: { content: '' }, done: true, prompt_eval_count: 10, eval_count: 5 };
      })(),
    );

    const client = new OllamaClient({ model: 'llama3.2' });

    const events: Array<{ type: string; content?: string; usage?: { inputTokens: number; outputTokens: number } }> = [];
    for await (const event of client.chatStream({
      messages: [{ role: 'user', content: 'Hello' }],
    })) {
      events.push(event);
    }

    expect(events).toHaveLength(3);
    expect(events[0]).toEqual({ type: 'content', content: 'Hello' });
    expect(events[1]).toEqual({ type: 'content', content: ' world' });
    expect(events[2]).toEqual({
      type: 'done',
      usage: { inputTokens: 10, outputTokens: 5 },
    });
  });

  it('streaming: yields tool_use events from final chunk', async () => {
    mockChat.mockResolvedValue(
      (async function* () {
        yield {
          message: {
            content: '',
            tool_calls: [
              { function: { name: 'system.info', arguments: {} } },
            ],
          },
          done: true,
          prompt_eval_count: 5,
          eval_count: 3,
        };
      })(),
    );

    const client = new OllamaClient({ model: 'llama3.2' });

    const events: Array<{ type: string; toolCall?: { id: string; name: string; args: unknown }; usage?: { inputTokens: number; outputTokens: number } }> = [];
    for await (const event of client.chatStream({
      messages: [{ role: 'user', content: 'Get system info' }],
    })) {
      events.push(event);
    }

    // Should have tool_use event followed by done
    expect(events).toHaveLength(2);
    expect(events[0]).toEqual({
      type: 'tool_use',
      toolCall: {
        id: 'ollama_tc_0',
        name: 'system.info',
        args: {},
      },
    });
    expect(events[1]).toEqual({
      type: 'done',
      usage: { inputTokens: 5, outputTokens: 3 },
    });
  });

  it('omits tools when model does not support them', async () => {
    mockShow.mockResolvedValue({ capabilities: ['completion'] });
    mockChat.mockResolvedValue({
      message: { content: 'I cannot use tools but I can answer.' },
      done_reason: 'stop',
      prompt_eval_count: 10,
      eval_count: 5,
    });

    const client = new OllamaClient({ model: 'qwen2.5-14b-instruct' });

    const response = await client.chat({
      messages: [{ role: 'user', content: 'What time is it?' }],
      tools: [
        {
          name: 'system.info',
          description: 'Get system info',
          input_schema: { type: 'object', properties: {}, required: [] },
        },
      ],
    });

    expect(response.content).toBe('I cannot use tools but I can answer.');
    const callArgs = mockChat.mock.calls[0][0];
    expect(callArgs.tools).toBeUndefined();
  });

  it('passes tools when model supports them', async () => {
    mockShow.mockResolvedValue({ capabilities: ['completion', 'tools'] });
    mockChat.mockResolvedValue({
      message: { content: 'Using tools.' },
      done_reason: 'stop',
      prompt_eval_count: 10,
      eval_count: 5,
    });

    const client = new OllamaClient({ model: 'llama3.2' });

    await client.chat({
      messages: [{ role: 'user', content: 'List files' }],
      tools: [
        {
          name: 'shell.exec',
          description: 'Run command',
          input_schema: { type: 'object', properties: { command: { type: 'string' } }, required: ['command'] },
        },
      ],
    });

    const callArgs = mockChat.mock.calls[0][0];
    expect(callArgs.tools).toBeDefined();
    expect(callArgs.tools[0].function.name).toBe('shell.exec');
  });

  it('caches capability check across multiple calls', async () => {
    mockShow.mockResolvedValue({ capabilities: ['completion', 'tools'] });
    mockChat.mockResolvedValue({
      message: { content: 'OK' },
      done_reason: 'stop',
      prompt_eval_count: 5,
      eval_count: 3,
    });

    const client = new OllamaClient({ model: 'llama3.2' });

    const tools = [
      {
        name: 'test.tool',
        description: 'A test tool',
        input_schema: { type: 'object' as const, properties: {} },
      },
    ];

    await client.chat({ messages: [{ role: 'user', content: 'a' }], tools });
    await client.chat({ messages: [{ role: 'user', content: 'b' }], tools });
    await client.chat({ messages: [{ role: 'user', content: 'c' }], tools });

    // show() should only be called once despite 3 chat calls
    expect(mockShow).toHaveBeenCalledTimes(1);
  });

  it('defaults to tools supported when show() fails', async () => {
    mockShow.mockRejectedValue(new Error('connection refused'));
    mockChat.mockResolvedValue({
      message: { content: 'OK' },
      done_reason: 'stop',
      prompt_eval_count: 5,
      eval_count: 3,
    });

    const client = new OllamaClient({ model: 'old-model' });

    await client.chat({
      messages: [{ role: 'user', content: 'hello' }],
      tools: [
        {
          name: 'test.tool',
          description: 'A test tool',
          input_schema: { type: 'object', properties: {} },
        },
      ],
    });

    // Tools should still be passed (optimistic default)
    const callArgs = mockChat.mock.calls[0][0];
    expect(callArgs.tools).toBeDefined();
  });
});

describe('normalizeMessagesForOllama', () => {
  it('passes plain text messages through', () => {
    const messages: Message[] = [
      { role: 'user', content: 'Hello' },
      { role: 'assistant', content: 'Hi there' },
    ];

    const result = normalizeMessagesForOllama('System prompt', messages);

    expect(result).toEqual([
      { role: 'system', content: 'System prompt' },
      { role: 'user', content: 'Hello' },
      { role: 'assistant', content: 'Hi there' },
    ]);
  });

  it('converts assistant tool_use blocks to Ollama tool_calls format', () => {
    const messages: Message[] = [
      { role: 'user', content: 'Search for news' },
      {
        role: 'assistant',
        content: [
          { type: 'text', text: 'Let me search for that.' },
          { type: 'tool_use', id: 'call_1', name: 'web.search', input: { query: 'latest news' } },
        ] as any,
      },
    ];

    const result = normalizeMessagesForOllama(undefined, messages);

    expect(result).toHaveLength(2);
    expect(result[0]).toEqual({ role: 'user', content: 'Search for news' });
    expect(result[1]).toEqual({
      role: 'assistant',
      content: 'Let me search for that.',
      tool_calls: [{
        function: {
          name: 'web.search',
          arguments: { query: 'latest news' },
        },
      }],
    });
  });

  it('converts user tool_result blocks to Ollama tool role messages with tool_name', () => {
    const messages: Message[] = [
      { role: 'user', content: 'Search for news' },
      {
        role: 'assistant',
        content: [
          { type: 'tool_use', id: 'call_1', name: 'web.search', input: { query: 'news' } },
        ] as any,
      },
      {
        role: 'user',
        content: [
          { type: 'tool_result', tool_use_id: 'call_1', content: 'Breaking news: ...', is_error: false },
        ] as any,
      },
    ];

    const result = normalizeMessagesForOllama(undefined, messages);

    expect(result).toHaveLength(3);
    expect(result[2]).toEqual({
      role: 'tool',
      content: 'Breaking news: ...',
      tool_name: 'web.search',
    });
  });

  it('handles multiple tool calls in a single assistant message', () => {
    const messages: Message[] = [
      { role: 'user', content: 'Search two things' },
      {
        role: 'assistant',
        content: [
          { type: 'tool_use', id: 'call_1', name: 'web.search', input: { query: 'a' } },
          { type: 'tool_use', id: 'call_2', name: 'web.search', input: { query: 'b' } },
        ] as any,
      },
      {
        role: 'user',
        content: [
          { type: 'tool_result', tool_use_id: 'call_1', content: 'Result A' },
          { type: 'tool_result', tool_use_id: 'call_2', content: 'Result B' },
        ] as any,
      },
    ];

    const result = normalizeMessagesForOllama(undefined, messages);

    expect(result[1].tool_calls).toHaveLength(2);
    // Each tool_result becomes a separate tool message with tool_name
    expect(result[2]).toEqual({ role: 'tool', content: 'Result A', tool_name: 'web.search' });
    expect(result[3]).toEqual({ role: 'tool', content: 'Result B', tool_name: 'web.search' });
  });

  it('handles full tool round-trip conversation', () => {
    const messages: Message[] = [
      { role: 'user', content: 'What is the weather?' },
      {
        role: 'assistant',
        content: [
          { type: 'text', text: 'Let me check.' },
          { type: 'tool_use', id: 'tc_0', name: 'weather.get', input: { city: 'NYC' } },
        ] as any,
      },
      {
        role: 'user',
        content: [
          { type: 'tool_result', tool_use_id: 'tc_0', content: 'Sunny, 72F' },
        ] as any,
      },
      { role: 'assistant', content: 'The weather in NYC is sunny, 72F.' },
    ];

    const result = normalizeMessagesForOllama('You are helpful.', messages);

    expect(result).toHaveLength(5);
    expect(result[0]).toEqual({ role: 'system', content: 'You are helpful.' });
    expect(result[1]).toEqual({ role: 'user', content: 'What is the weather?' });
    expect(result[2]).toEqual({
      role: 'assistant',
      content: 'Let me check.',
      tool_calls: [{ function: { name: 'weather.get', arguments: { city: 'NYC' } } }],
    });
    expect(result[3]).toEqual({ role: 'tool', content: 'Sunny, 72F', tool_name: 'weather.get' });
    expect(result[4]).toEqual({ role: 'assistant', content: 'The weather in NYC is sunny, 72F.' });
  });
});
