import { describe, it, expect, vi, beforeEach } from 'vitest';
import { parseLlamaTimings, computeAggregates, runSingleBenchRequest } from '../../index.js';
import { computeRegressionFlag } from '../bench-engine.js';
import { createFleetState, ensureHostState } from '../fleet-state.js';
import { createDeltaEmitter } from '../delta-emitter.js';
import type { Sql } from '../../db.js';
import type { Config } from '../../config.js';
import type { BenchSuite } from '../bench-engine.js';

// ─── parseLlamaTimings tests ────────────────────────────────────────────────

describe('parseLlamaTimings', () => {
  it('parses timings from a standard llama.cpp chunk', () => {
    const chunk = 'data: {"choices":[],"timings":{"prompt_per_second":150,"predicted_per_second":80,"cache_n":50}}';
    const result = parseLlamaTimings(chunk);
    expect(result).not.toBeNull();
    expect(result!.promptPerSecond).toBe(150);
    expect(result!.predictedPerSecond).toBe(80);
    expect(result!.cacheN).toBe(50);
  });

  it('parses timings without data: prefix', () => {
    const chunk = '{"timings":{"prompt_per_second":200,"predicted_per_second":100,"cache_n":0}}';
    const result = parseLlamaTimings(chunk);
    expect(result).not.toBeNull();
    expect(result!.promptPerSecond).toBe(200);
  });

  it('returns null for [DONE] chunk', () => {
    expect(parseLlamaTimings('data: [DONE]')).toBeNull();
  });

  it('returns null for chunk without timings', () => {
    const chunk = 'data: {"choices":[{"delta":{"content":"hello"}}]}';
    expect(parseLlamaTimings(chunk)).toBeNull();
  });

  it('returns null for malformed JSON', () => {
    expect(parseLlamaTimings('data: not-json')).toBeNull();
  });
});

// ─── computeAggregates tests ────────────────────────────────────────────────

describe('computeAggregates', () => {
  it('returns nulls for empty samples', () => {
    const result = computeAggregates([]);
    expect(result.totalSamples).toBe(0);
    expect(result.avgTtftMs).toBeNull();
    expect(result.avgGenTps).toBeNull();
  });

  it('computes averages correctly', () => {
    const samples = [
      { ttftMs: 100, genTps: 50, promptTps: 100, error: null } as any,
      { ttftMs: 200, genTps: 100, promptTps: 200, error: null } as any,
      { ttftMs: 300, genTps: 150, promptTps: 300, error: null } as any,
    ];
    const result = computeAggregates(samples);
    expect(result.avgTtftMs).toBe(200);
    expect(result.avgGenTps).toBe(100);
    expect(result.avgPromptTps).toBe(200);
    expect(result.totalSamples).toBe(3);
    expect(result.errorSamples).toBe(0);
  });

  it('computes median correctly for odd count', () => {
    const samples = [
      { ttftMs: 100, genTps: 50, promptTps: 100, error: null } as any,
      { ttftMs: 200, genTps: 100, promptTps: 200, error: null } as any,
      { ttftMs: 300, genTps: 150, promptTps: 300, error: null } as any,
    ];
    const result = computeAggregates(samples);
    expect(result.medianTtftMs).toBe(200);
    expect(result.medianGenTps).toBe(100);
  });

  it('computes median correctly for even count', () => {
    const samples = [
      { ttftMs: 100, genTps: 50, promptTps: 100, error: null } as any,
      { ttftMs: 200, genTps: 100, promptTps: 200, error: null } as any,
      { ttftMs: 300, genTps: 150, promptTps: 300, error: null } as any,
      { ttftMs: 400, genTps: 200, promptTps: 400, error: null } as any,
    ];
    const result = computeAggregates(samples);
    expect(result.medianTtftMs).toBe(250);
    expect(result.medianGenTps).toBe(125);
  });

  it('computes p95 TTFT', () => {
    const samples = Array.from({ length: 20 }, (_, i) => ({
      ttftMs: (i + 1) * 10,
      genTps: 50,
      promptTps: 100,
      error: null,
    })) as any[];
    const result = computeAggregates(samples);
    expect(result.p95TtftMs).toBeCloseTo(190, -1);
  });

  it('filters out null values', () => {
    const samples = [
      { ttftMs: 100, genTps: 50, promptTps: 100, error: null } as any,
      { ttftMs: null, genTps: null, promptTps: null, error: 'timeout' } as any,
    ];
    const result = computeAggregates(samples);
    expect(result.avgTtftMs).toBe(100);
    expect(result.errorSamples).toBe(1);
  });
});

// ─── bench runner pipeline test (mock fetch + real functions) ────────────────

describe('bench runner pipeline', () => {
  let mockSql: Sql;
  let executedQueries: Array<{ query: string; values: unknown[] }>;

  beforeEach(() => {
    executedQueries = [];
    mockSql = Object.assign(
      (strings: TemplateStringsArray, ...values: unknown[]) => {
        const query = strings.reduce((acc: string, s: string, i: number) => acc + s + (values[i] ?? ''), '');
        executedQueries.push({ query, values });
        return Promise.resolve([]);
      },
      {
        json: (v: unknown) => v,
        unsafe: async (q: string) => { executedQueries.push({ query: q, values: [] }); return []; },
      },
    ) as unknown as Sql;
  });

  it('runSingleBenchRequest captures TTFT and timings on successful stream', async () => {
    const fakeStream = createFakeStreamResponse([
      'data: {"choices":[{"delta":{"content":"H"}}]}',
      'data: {"choices":[{"delta":{"content":"ello"}}]}',
      'data: {"choices":[],"timings":{"prompt_per_second":150,"predicted_per_second":80,"cache_n":10}}',
      'data: [DONE]',
    ]);

    vi.spyOn(global, 'fetch').mockResolvedValueOnce(fakeStream);

    const sample = await runSingleBenchRequest(
      'http://localhost:8401',
      'test-model',
      10,
      20,
      0,
      0.7,
      0.9,
    );

    expect(sample.error).toBeNull();
    expect(sample.ttftMs).toBeGreaterThanOrEqual(0);
    expect(sample.ttftMs).toBeLessThan(5000);
    expect(sample.totalMs).toBeGreaterThanOrEqual(0);
    expect(sample.promptTps).toBe(150);
    expect(sample.genTps).toBe(80);
    expect(sample.cacheN).toBe(10);
    expect(sample.promptTokens).toBe(10);
    expect(sample.genTokens).toBe(20);
    expect(sample.repetition).toBe(0);

    vi.restoreAllMocks();
  });

  it('runSingleBenchRequest captures error on HTTP failure', async () => {
    vi.spyOn(global, 'fetch').mockResolvedValueOnce({
      ok: false,
      status: 500,
      text: async () => 'Internal Server Error',
    } as Response);

    const sample = await runSingleBenchRequest(
      'http://localhost:8401',
      'test-model',
      10,
      20,
      0,
    );

    expect(sample.error).toContain('500');
    expect(sample.ttftMs).toBeNull();

    vi.restoreAllMocks();
  });

  it('runSingleBenchRequest captures error on fetch exception', async () => {
    vi.spyOn(global, 'fetch').mockRejectedValueOnce(new Error('ECONNREFUSED'));

    const sample = await runSingleBenchRequest(
      'http://localhost:8401',
      'test-model',
      10,
      20,
      0,
    );

    expect(sample.error).toContain('ECONNREFUSED');

    vi.restoreAllMocks();
  });
});

// ─── helper: create a fake streaming Response ────────────────────────────────

function createFakeStreamResponse(lines: string[]): Response {
  const encoder = new TextEncoder();
  let position = 0;

  const stream = new ReadableStream({
    async pull(controller) {
      if (position >= lines.length) {
        controller.close();
        return;
      }
      const line = lines[position]! + '\n\n';
      controller.enqueue(encoder.encode(line));
      position++;
      // Small delay to simulate network latency for TTFT measurement
      await new Promise((r) => setTimeout(r, 5));
    },
  });

  return new Response(stream, {
    status: 200,
    headers: { 'Content-Type': 'text/event-stream' },
  });
}

// ─── computeRegressionFlag tests (A1) ────────────────────────────────────────

describe('computeRegressionFlag', () => {
  it('returns baseline for first run (no baseline)', () => {
    const current = computeAggregates([
      { ttftMs: 100, genTps: 80, promptTps: 150, error: null } as any,
    ]);
    expect(computeRegressionFlag(current, undefined)).toBe('baseline');
  });

  it('returns regression when gen tok/s drops below -10%', () => {
    const current = computeAggregates([
      { ttftMs: 200, genTps: 70, promptTps: 100, error: null } as any,
    ]);
    const baseline = JSON.stringify({
      avgGenTps: 100,
      avgTtftMs: 100,
      totalSamples: 1,
    });
    expect(computeRegressionFlag(current, baseline)).toBe('regression');
  });

  it('returns improvement when gen tok/s rises above +5%', () => {
    const current = computeAggregates([
      { ttftMs: 80, genTps: 120, promptTps: 200, error: null } as any,
    ]);
    const baseline = JSON.stringify({
      avgGenTps: 100,
      avgTtftMs: 100,
      totalSamples: 1,
    });
    expect(computeRegressionFlag(current, baseline)).toBe('improvement');
  });

  it('returns baseline when within threshold', () => {
    const current = computeAggregates([
      { ttftMs: 100, genTps: 98, promptTps: 150, error: null } as any,
    ]);
    const baseline = JSON.stringify({
      avgGenTps: 100,
      avgTtftMs: 100,
      totalSamples: 1,
    });
    expect(computeRegressionFlag(current, baseline)).toBe('baseline');
  });

  it('returns null for divide-by-zero (N5: baseline avgGenTps is 0)', () => {
    const current = computeAggregates([
      { ttftMs: 100, genTps: 50, promptTps: 100, error: null } as any,
    ]);
    const baseline = JSON.stringify({
      avgGenTps: 0,
      avgTtftMs: 100,
      totalSamples: 1,
    });
    expect(computeRegressionFlag(current, baseline)).toBeNull();
  });

  it('returns null for null current avgGenTps', () => {
    const current = computeAggregates([]);
    expect(computeRegressionFlag(current, JSON.stringify({ avgGenTps: 100 }))).toBeNull();
  });

  it('returns null for malformed baseline JSON', () => {
    const current = computeAggregates([
      { ttftMs: 100, genTps: 80, promptTps: 150, error: null } as any,
    ]);
    expect(computeRegressionFlag(current, 'not-json')).toBeNull();
  });
});
