import { OllamaClient } from '../src/ollama-client';
import { OllamaMessage, OllamaTool } from '../src/types';

describe('OllamaClient', () => {
  let client: OllamaClient;
  let mockFetch: jest.Mock;

  function createMockReader(data: string) {
    const encoder = new TextEncoder();
    const encoded = encoder.encode(data);
    let called = false;
    return {
      read: () => {
        if (!called) {
          called = true;
          return Promise.resolve({ done: false, value: encoded });
        }
        return Promise.resolve({ done: true, value: new Uint8Array(0) });
      },
      releaseLock: jest.fn(),
    };
  }

  const mockMessages: OllamaMessage[] = [
    { role: 'system', content: 'You are helpful.' },
    { role: 'user', content: 'Hello' },
  ];

  const mockTools: OllamaTool[] = [
    {
      type: 'function',
      function: {
        name: 'test_tool',
        description: 'A test tool',
        parameters: {
          type: 'object',
          properties: { input: { type: 'string' } },
          required: ['input'],
        },
      },
    },
  ];

  beforeEach(() => {
    mockFetch = jest.fn();
    client = new OllamaClient('http://localhost:11434', 'llama3', mockFetch);
  });

  afterEach(() => {
    jest.clearAllMocks();
    client.cancelStream();
  });

  describe('chat (non-streaming)', () => {
    it('should send a non-streaming request and return the response', async () => {
      const mockResponse = {
        ok: true,
        json: () => Promise.resolve({ message: { content: 'Hello back!' } }),
      };
      mockFetch.mockResolvedValue(mockResponse);

      const result = await client.chat(mockMessages, mockTools);

      expect(result.content).toBe('Hello back!');
      expect(mockFetch).toHaveBeenCalledWith(
        'http://localhost:11434/api/chat',
        expect.objectContaining({
          method: 'POST',
          body: JSON.stringify({
            model: 'llama3',
            messages: mockMessages,
            tools: mockTools,
            stream: false,
          }),
        })
      );
    });

    it('should throw on non-OK response', async () => {
      mockFetch.mockResolvedValue({ ok: false, status: 500 });

      await expect(client.chat(mockMessages, mockTools)).rejects.toThrow('Ollama API error: 500');
    });

    it('should handle missing message content gracefully', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        json: () => Promise.resolve({}),
      });

      const result = await client.chat(mockMessages, mockTools);
      expect(result.content).toBe('');
      expect(result.tool_calls).toEqual([]);
    });

    it('should include abort signal in fetch options', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        json: () => Promise.resolve({ message: { content: 'ok' } }),
      });

      await client.chat(mockMessages, mockTools);

      const fetchOptions = mockFetch.mock.calls[0][1];
      expect(fetchOptions.signal).toBeInstanceOf(AbortSignal);
    });

    it('should forward tool_calls from response when present', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        json: () =>
          Promise.resolve({
            message: {
              content: 'result',
              tool_calls: [{ function: { name: 'create_file', arguments: '{}' } }],
            },
          }),
      });

      const result = await client.chat(mockMessages, mockTools);
      expect(result.tool_calls).toEqual([{ function: { name: 'create_file', arguments: '{}' } }]);
    });
  });

  describe('streamChat', () => {
    it('should send a streaming request and yield chunks', async () => {
      const streamData = [
        JSON.stringify({ message: { content: 'He' } }),
        JSON.stringify({ message: { content: 'llo' } }),
        JSON.stringify({ message: { content: '!' } }),
        '',
      ].join('\n');

      const mockReader = createMockReader(streamData);
      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);
      const chunks: string[] = [];

      for await (const chunk of stream) {
        chunks.push(chunk.content);
      }

      expect(chunks).toEqual(['He', 'llo', '!']);
      expect(mockReader.releaseLock).toHaveBeenCalled();
    });

    it('should skip malformed JSON chunks and log a warning', async () => {
      const streamData = [
        JSON.stringify({ message: { content: 'valid' } }),
        'this is not json',
        JSON.stringify({ message: { content: 'also valid' } }),
        '',
      ].join('\n');

      const mockReader = createMockReader(streamData);
      const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();

      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);
      const chunks: string[] = [];

      for await (const chunk of stream) {
        chunks.push(chunk.content);
      }

      expect(chunks).toEqual(['valid', 'also valid']);
      expect(consoleWarnSpy).toHaveBeenCalledWith(
        expect.stringContaining(
          '[ollama-client] WARN: Skipped malformed chunk: this is not json... - Unexpected token \'h\', "this is not json" is not valid JSON'
        )
      );
      consoleWarnSpy.mockRestore();
    });

    it('should throw when too many chunks are malformed', async () => {
      const streamData = Array(51).fill('invalid json').join('\n') + '\n';

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);

      await expect(
        (async () => {
          for await (const _ of stream) {
            /* consume */
          }
        })()
      ).rejects.toThrow(/malformed/);
    });

    it('should throw on non-OK response', async () => {
      mockFetch.mockResolvedValue({ ok: false, status: 404 });

      await expect(client.streamChatAsPromise(mockMessages, mockTools)).rejects.toThrow(
        'Ollama API error: 404'
      );
    });

    it('should throw when response has no body', async () => {
      mockFetch.mockResolvedValue({ ok: true, body: undefined });

      await expect(client.streamChatAsPromise(mockMessages, mockTools)).rejects.toThrow(
        'No response body'
      );
    });

    it('should throw on invalid content type', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        body: {
          getReader: () => ({
            read: () => Promise.resolve({ done: true, value: new Uint8Array(0) }),
          }),
        },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'text/html' : null),
        },
      });

      await expect(client.streamChatAsPromise(mockMessages, mockTools)).rejects.toThrow(
        'Invalid response format'
      );
    });

    it('should propagate Ollama error messages from the stream', async () => {
      const streamData = JSON.stringify({ error: 'model not found' }) + '\n';

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);

      await expect(
        (async () => {
          for await (const _ of stream) {
            /* consume */
          }
        })()
      ).rejects.toThrow('Ollama error: model not found');
    });

    it('should yield tool_calls when present in streamed response', async () => {
      const streamData = [
        JSON.stringify({
          message: {
            content: '',
            tool_calls: [{ function: { name: 'create_file', arguments: '{"path":"a.md"}' } }],
          },
        }),
        '',
      ].join('\n');

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);
      let lastChunk: any;

      for await (const chunk of stream) {
        lastChunk = chunk;
      }

      expect(lastChunk.tool_calls).toEqual([
        { function: { name: 'create_file', arguments: '{"path":"a.md"}' } },
      ]);
    });

    it('should default tool_calls to empty array when not present', async () => {
      const streamData = JSON.stringify({ message: { content: 'hello' } }) + '\n';

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);
      let lastChunk: any;

      for await (const chunk of stream) {
        lastChunk = chunk;
      }

      expect(lastChunk.tool_calls).toEqual([]);
    });

    it('should send correct request body with stream:true', async () => {
      const streamData = JSON.stringify({ message: { content: 'ok' } }) + '\n';
      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);
      for await (const _ of stream) {
        /* consume */
      }

      expect(mockFetch).toHaveBeenCalledWith(
        'http://localhost:11434/api/chat',
        expect.objectContaining({
          method: 'POST',
          body: JSON.stringify({
            model: 'llama3',
            messages: mockMessages,
            tools: mockTools,
            stream: true,
          }),
          signal: expect.any(AbortSignal),
        })
      );
    });
  });

  describe('streamChat with retry logic', () => {
    it('should retry on 5xx errors and eventually succeed', async () => {
      let callCount = 0;
      mockFetch.mockImplementation(async () => {
        callCount++;
        if (callCount === 1) {
          return { ok: false, status: 500 };
        }
        if (callCount === 2) {
          return { ok: false, status: 502 };
        }
        return {
          ok: true,
          body: {
            getReader: () => ({
              read: () => Promise.resolve({ done: true, value: new Uint8Array(0) }),
              releaseLock: () => {},
            }),
          },
          headers: {
            get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
          },
        };
      });

      const stream = await client.streamChat(mockMessages, mockTools);
      const chunks: string[] = [];

      for await (const _ of stream) {
        chunks.push('chunk');
      }

      expect(callCount).toBe(3);
      expect(chunks.length).toBe(0);
    });

    it('should give up after maxRetries attempts', async () => {
      mockFetch.mockResolvedValue({ ok: false, status: 500 });

      const stream = await client.streamChat(mockMessages, mockTools);

      await expect(
        (async () => {
          for await (const _ of stream) {
            /* consume */
          }
        })()
      ).rejects.toThrow('Ollama API error: 500');
    });

    it('should not retry on 4xx errors', async () => {
      mockFetch.mockResolvedValue({ ok: false, status: 404 });

      const stream = await client.streamChat(mockMessages, mockTools);

      await expect(
        (async () => {
          for await (const _ of stream) {
            /* consume */
          }
        })()
      ).rejects.toThrow('Ollama API error: 404');
    });
  });

  describe('cancelStream', () => {
    it('should abort the current request', () => {
      client.cancelStream();
      expect(client['abortController']).toBeNull();
    });

    it('should handle cancel when no active stream', () => {
      expect(() => client.cancelStream()).not.toThrow();
      expect(client['abortController']).toBeNull();
    });
  });

  describe('streamChat final buffer parsing', () => {
    it('should parse final buffer content when stream ends with partial line', async () => {
      const streamData = [
        JSON.stringify({ message: { content: 'First' } }),
        JSON.stringify({ message: { content: 'Second' } }),
        JSON.stringify({ message: { content: 'Third' } }),
        '', // Final line should be empty to signal end
      ].join('\n');

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);
      const chunks: string[] = [];

      for await (const chunk of stream) {
        chunks.push(chunk.content);
      }

      expect(chunks).toEqual(['First', 'Second', 'Third']);
      expect(mockReader.releaseLock).toHaveBeenCalled();
    });

    it('should handle malformed final buffer content gracefully', async () => {
      const streamData = [
        JSON.stringify({ message: { content: 'Valid' } }),
        'malformed json',
        '', // Final line should be empty to signal end
      ].join('\n');

      const mockReader = createMockReader(streamData);
      const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();

      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);
      const chunks: string[] = [];

      for await (const chunk of stream) {
        chunks.push(chunk.content);
      }

      expect(chunks).toEqual(['Valid']);
      expect(consoleWarnSpy).toHaveBeenCalledWith(
        expect.stringContaining('Skipped malformed chunk')
      );
      consoleWarnSpy.mockRestore();
    });

    it('should parse final buffer content even when it contains message data', async () => {
      const streamData = [
        JSON.stringify({ message: { content: 'First' } }),
        '', // Final line should be empty to signal end
        JSON.stringify({ message: { content: 'Final' } }),
      ].join('\n');

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValue({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
        },
      });

      const stream = await client.streamChat(mockMessages, mockTools);
      const chunks: string[] = [];

      for await (const chunk of stream) {
        chunks.push(chunk.content);
      }

      expect(chunks).toEqual(['First', 'Final']);
      expect(mockReader.releaseLock).toHaveBeenCalled();
    });
  });
});