import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { LlmRouter } from '../router.js';
import { createLocalOllamaProvider } from '../registry.js';
import type { ProviderConfig, ChatCompletionResponse } from '../types.js';
import * as client from '../client.js';

// Mock the HTTP client
vi.mock('../client.js', () => ({
  sendChatCompletion: vi.fn(),
}));

const MOCK_RESPONSE: ChatCompletionResponse = {
  id: 'chatcmpl-test',
  object: 'chat.completion',
  created: Date.now(),
  model: 'test-model',
  choices: [
    {
      index: 0,
      message: { role: 'assistant', content: 'Hello!' },
      finish_reason: 'stop',
    },
  ],
  usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
};

const TEST_PROVIDERS: ProviderConfig[] = [
  {
    name: 'test-fast',
    baseUrl: 'https://fast.test/v1',
    apiKeyEnv: 'TEST_FAST_KEY',
    rpmLimit: 30,
    tpmLimit: 10_000,
    models: [
      {
        id: 'fast-model',
        label: 'Fast',
        contextWindow: 8_192,
        strengths: ['general'],
        speedTier: 1,
      },
    ],
  },
  {
    name: 'test-quality',
    baseUrl: 'https://quality.test/v1',
    apiKeyEnv: 'TEST_QUALITY_KEY',
    rpmLimit: 20,
    tpmLimit: 0,
    models: [
      {
        id: 'quality-model',
        label: 'Quality',
        contextWindow: 128_000,
        strengths: ['code', 'reasoning'],
        speedTier: 2,
      },
    ],
  },
];

describe('LlmRouter', () => {
  beforeEach(() => {
    vi.resetAllMocks();
    // Set fake API keys
    process.env.TEST_FAST_KEY = 'test-key-fast';
    process.env.TEST_QUALITY_KEY = 'test-key-quality';
  });

  afterEach(() => {
    delete process.env.TEST_FAST_KEY;
    delete process.env.TEST_QUALITY_KEY;
  });

  it('throws if no providers have API keys', () => {
    delete process.env.TEST_FAST_KEY;
    delete process.env.TEST_QUALITY_KEY;
    expect(() => new LlmRouter({ providers: TEST_PROVIDERS })).toThrow('No providers available');
  });

  it('routes a simple prompt to a provider', async () => {
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: MOCK_RESPONSE,
      latencyMs: 150,
      status: 200,
    });

    const router = new LlmRouter({ providers: TEST_PROVIDERS });
    const result = await router.chat({
      messages: [{ role: 'user', content: 'Hello' }],
    });

    expect(result.response.choices[0]!.message.content).toBe('Hello!');
    expect(result.attempts).toBe(1);
    expect(result.provider).toBe('test-fast');
    expect(result.model).toBe('fast-model');
  });

  it('retries on 429 with fallback provider', async () => {
    // First call: rate limited
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: null as unknown as ChatCompletionResponse,
      latencyMs: 50,
      status: 429,
    });
    // Second call: success
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: MOCK_RESPONSE,
      latencyMs: 200,
      status: 200,
    });

    const router = new LlmRouter({ providers: TEST_PROVIDERS });
    const result = await router.chat({
      messages: [{ role: 'user', content: 'Hello' }],
    });

    expect(result.attempts).toBe(2);
    expect(result.response.choices[0]!.message.content).toBe('Hello!');
  });

  it('retries on error with fallback provider', async () => {
    // First call: error
    vi.mocked(client.sendChatCompletion).mockRejectedValueOnce(new Error('Network error'));
    // Second call: success
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: MOCK_RESPONSE,
      latencyMs: 200,
      status: 200,
    });

    const router = new LlmRouter({ providers: TEST_PROVIDERS });
    const result = await router.chat({
      messages: [{ role: 'user', content: 'Hello' }],
    });

    expect(result.attempts).toBe(2);
  });

  it('throws after exhausting all retries', async () => {
    vi.mocked(client.sendChatCompletion).mockRejectedValue(new Error('All down'));

    const router = new LlmRouter({ providers: TEST_PROVIDERS, maxRetries: 2 });
    await expect(router.chat({ messages: [{ role: 'user', content: 'Hello' }] })).rejects.toThrow(
      'All providers exhausted'
    );
  });

  it('routes code prompts to code-capable models', async () => {
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: MOCK_RESPONSE,
      latencyMs: 200,
      status: 200,
    });

    const router = new LlmRouter({ providers: TEST_PROVIDERS });
    await router.chat({
      messages: [{ role: 'user', content: 'Write a typescript function to sort an array' }],
    });

    // Should have been called with quality-model (has 'code' strength)
    const callArgs = vi.mocked(client.sendChatCompletion).mock.calls[0]!;
    expect(callArgs[1]).toBe('quality-model');
  });

  it('fires telemetry callback on success', async () => {
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: MOCK_RESPONSE,
      latencyMs: 150,
      status: 200,
    });

    const telemetry = vi.fn();
    const router = new LlmRouter({ providers: TEST_PROVIDERS, onTelemetry: telemetry });
    await router.chat({ messages: [{ role: 'user', content: 'Hello' }] });

    expect(telemetry).toHaveBeenCalledWith(
      expect.objectContaining({ event: 'success', attempt: 1 })
    );
  });

  it('fires telemetry callback on rate limit', async () => {
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: null as unknown as ChatCompletionResponse,
      latencyMs: 50,
      status: 429,
    });
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: MOCK_RESPONSE,
      latencyMs: 200,
      status: 200,
    });

    const telemetry = vi.fn();
    const router = new LlmRouter({ providers: TEST_PROVIDERS, onTelemetry: telemetry });
    await router.chat({ messages: [{ role: 'user', content: 'Hello' }] });

    expect(telemetry).toHaveBeenCalledWith(expect.objectContaining({ event: 'rate_limit' }));
  });

  it('handles explicit provider:model routing', async () => {
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: MOCK_RESPONSE,
      latencyMs: 100,
      status: 200,
    });

    const router = new LlmRouter({ providers: TEST_PROVIDERS });
    const result = await router.chat({
      messages: [{ role: 'user', content: 'Hello' }],
      model: 'test-fast:fast-model',
    });

    expect(result.provider).toBe('test-fast');
    expect(result.model).toBe('fast-model');
  });

  it('throws for unknown explicit provider', async () => {
    const router = new LlmRouter({ providers: TEST_PROVIDERS });
    await expect(
      router.chat({ messages: [{ role: 'user', content: 'Hello' }], model: 'unknown:model' })
    ).rejects.toThrow('Provider "unknown" not found');
  });

  it('returns health snapshots', async () => {
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: MOCK_RESPONSE,
      latencyMs: 150,
      status: 200,
    });

    const router = new LlmRouter({ providers: TEST_PROVIDERS });
    await router.chat({ messages: [{ role: 'user', content: 'Hello' }] });

    const health = router.getHealth();
    expect(health).toContainEqual(
      expect.objectContaining({
        provider: 'test-fast',
        model: 'fast-model',
        successes: 1,
      })
    );
  });

  it('lists available providers', () => {
    const router = new LlmRouter({ providers: TEST_PROVIDERS });
    expect(router.getProviders()).toEqual(['test-fast', 'test-quality']);
  });

  it('plans the best provider/model without executing a request', () => {
    const router = new LlmRouter({ providers: TEST_PROVIDERS });
    const plan = router.plan({
      messages: [{ role: 'user', content: 'Write a TypeScript endpoint with validation' }],
    });

    expect(plan.provider.name).toBe('test-quality');
    expect(plan.model.id).toBe('quality-model');
    expect(plan.category).toBe('code');
    expect(plan.explicit).toBe(false);
  });

  it('plans local ollama models without requiring an API key', () => {
    const router = new LlmRouter({
      providers: [createLocalOllamaProvider(['qwen2.5-coder:7b', 'llama3.1:8b'])],
    });
    const plan = router.plan({
      messages: [{ role: 'user', content: 'Refactor this TypeScript function' }],
    });

    expect(plan.provider.name).toBe('local-ollama');
    expect(plan.model.id).toBe('qwen2.5-coder:7b');
  });

  it('fires telemetry for explicit model routing', async () => {
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: MOCK_RESPONSE,
      latencyMs: 100,
      status: 200,
    });

    const telemetry = vi.fn();
    const router = new LlmRouter({ providers: TEST_PROVIDERS, onTelemetry: telemetry });
    await router.chat({
      messages: [{ role: 'user', content: 'Hello' }],
      model: 'test-fast:fast-model',
    });

    expect(telemetry).toHaveBeenCalledWith(
      expect.objectContaining({
        event: 'success',
        provider: 'test-fast',
        model: 'fast-model',
        category: 'explicit',
      })
    );
  });

  it('records health on explicit model 429', async () => {
    vi.mocked(client.sendChatCompletion).mockResolvedValueOnce({
      response: null as unknown as ChatCompletionResponse,
      latencyMs: 50,
      status: 429,
    });

    const telemetry = vi.fn();
    const router = new LlmRouter({ providers: TEST_PROVIDERS, onTelemetry: telemetry });
    await expect(
      router.chat({ messages: [{ role: 'user', content: 'Hello' }], model: 'test-fast:fast-model' })
    ).rejects.toThrow('Rate limited');

    expect(telemetry).toHaveBeenCalledWith(
      expect.objectContaining({ event: 'rate_limit', provider: 'test-fast' })
    );

    // Health should have recorded the rate limit
    const health = router.getHealth();
    expect(health).toHaveLength(1);
    expect(health[0]!.rateLimits).toBe(1);
  });
});
