/**
 * Bench engine: speed benchmark runner.
 *
 * Suite = grid of (prompt_tokens x gen_tokens x concurrency) x repetitions.
 * TTFT measured client-side at first stream delta.
 * llama.cpp timings parsed from final stream chunk.
 * Bounded fan-out via Promise.allSettled at suite-declared concurrency.
 * Warmup excluded from results.
 */

import type { Sql } from '../db.js';
import type { DeltaEmitter } from './delta-emitter.js';
import { publishJob } from './publish-job.js';
import { jsonbObject } from './jsonb.js';

// ─── types ──────────────────────────────────────────────────────────────────

export interface BenchSuite {
  id: string;
  name: string;
  providerId: string;
  model: string;
  promptTokens: number[];
  genTokens: number[];
  concurrency: number[];
  repetitions: number;
  temperature?: number;
  topP?: number;
  metadata?: Record<string, unknown>;
}

export interface BenchRunParams {
  suite: BenchSuite;
  baseUrl: string;
  temperature?: number;
  topP?: number;
}

export interface BenchTimings {
  promptPerSecond: number;
  predictedPerSecond: number;
  cacheN: number;
}

export interface BenchSample {
  promptTokens: number;
  genTokens: number;
  concurrency: number;
  repetition: number;
  ttftMs: number | null;
  totalMs: number | null;
  promptTps: number | null;
  genTps: number | null;
  cacheN: number | null;
  error: string | null;
}

// ─── stream parser ──────────────────────────────────────────────────────────

/**
 * Parse llama.cpp timings from the final chunk of a streaming response.
 * llama.cpp returns timings in the last chunk's usage or as a separate field:
 *   { "timings": { "prompt_per_second": N, "predicted_per_second": N, "cache_n": N } }
 * or in the usage object.
 */
export function parseLlamaTimings(chunk: string): BenchTimings | null {
  try {
    // Strip "data: " prefix if present
    const jsonStr = chunk.startsWith('data: ') ? chunk.slice(6) : chunk;
    if (jsonStr.trim() === '[DONE]') return null;

    const parsed = JSON.parse(jsonStr) as Record<string, unknown>;

    // Try the timings object first (llama.cpp standard)
    const timings = parsed.timings as {
      prompt_per_second?: number;
      predicted_per_second?: number;
      cache_n?: number;
    } | undefined;
    if (timings) {
      return {
        promptPerSecond: timings.prompt_per_second ?? 0,
        predictedPerSecond: timings.predicted_per_second ?? 0,
        cacheN: timings.cache_n ?? 0,
      };
    }

    // Fallback: check usage.completion_tokens_details or completion_tokens
    const usage = parsed.usage as {
      prompt_tokens?: number;
      completion_tokens?: number;
    } | undefined;
    if (usage) {
      return {
        promptPerSecond: 0,
        predictedPerSecond: 0,
        cacheN: 0,
      };
    }

    return null;
  } catch {
    return null;
  }
}

// ─── single request runner ──────────────────────────────────────────────────

/**
 * Run a single bench request: stream completion, capture TTFT, parse timings.
 * Returns a BenchSample.
 */
export async function runSingleBenchRequest(
  baseUrl: string,
  model: string,
  promptTokens: number,
  genTokens: number,
  repetition: number,
  temperature: number = 0.7,
  topP: number = 0.9,
): Promise<BenchSample> {
  const sample: BenchSample = {
    promptTokens,
    genTokens,
    concurrency: 1, // set by the fan-out caller
    repetition,
    ttftMs: null,
    totalMs: null,
    promptTps: null,
    genTps: null,
    cacheN: null,
    error: null,
  };

  // Generate a deterministic prompt of the target length.
  const prompt = generatePrompt(promptTokens);

  const startTime = Date.now();
  let firstDeltaTime: number | null = null;
  let timings: BenchTimings | null = null;

  try {
    const res = await fetch(`${baseUrl}/v1/chat/completions`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        model,
        messages: [{ role: 'user', content: prompt }],
        temperature,
        top_p: topP,
        max_tokens: genTokens,
        stream: true,
      }),
      signal: AbortSignal.timeout(120_000),
    });

    if (!res.ok) {
      const errBody = await res.text().catch(() => '');
      throw new Error(`bench request failed: ${res.status} ${errBody.slice(0, 200)}`);
    }

    const reader = res.body?.getReader();
    if (!reader) {
      throw new Error('no response body');
    }

    const decoder = new TextDecoder();
    let buffer = '';

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;

      buffer += decoder.decode(value, { stream: true });
      const lines = buffer.split('\n');
      buffer = lines.pop() ?? '';

      for (const line of lines) {
        const trimmed = line.trim();
        if (!trimmed || trimmed === 'data: [DONE]') continue;

        // TTFT: capture at first delta
        if (firstDeltaTime === null) {
          firstDeltaTime = Date.now();
        }

        // Parse timings from the final chunk
        const t = parseLlamaTimings(trimmed);
        if (t) {
          timings = t;
        }
      }
    }

    sample.ttftMs = firstDeltaTime !== null ? firstDeltaTime - startTime : null;
    sample.totalMs = Date.now() - startTime;

    if (timings) {
      sample.promptTps = timings.promptPerSecond;
      sample.genTps = timings.predictedPerSecond;
      sample.cacheN = timings.cacheN;
    }
  } catch (err) {
    sample.error = (err as Error).message ?? String(err);
  }

  return sample;
}

/**
 * Generate a deterministic prompt with approximately the target token count.
 * Uses a repeating pattern that averages ~1.3 chars per token for GPT-style tokenizers.
 */
function generatePrompt(targetTokens: number): string {
  // Simple pattern: repeat a sentence that tokenizes predictably.
  // ~1.3 chars/token is a rough average for English text.
  const charsPerToken = 4;
  const targetChars = targetTokens * charsPerToken;
  const base = 'The quick brown fox jumps over the lazy dog. ';
  let result = '';
  while (result.length < targetChars) {
    result += base;
  }
  return result.slice(0, targetChars);
}

// ─── bench runner ───────────────────────────────────────────────────────────

export interface BenchRunProgress {
  jobId: string;
  totalSamples: number;
  completedSamples: number;
  currentPromptTokens: number;
  currentGenTokens: number;
  currentConcurrency: number;
  currentRepetition: number;
}

/**
 * Run a full bench suite: grid of all combinations.
 * Bounded fan-out via Promise.allSettled at suite-declared concurrency.
 * Warmup excluded from results (1 warmup request per unique grid cell, discarded).
 */
export async function runBenchSuite(
  params: BenchRunParams,
  sql: Sql,
  emitter: DeltaEmitter,
  seq: number,
  onProgress: (progress: BenchRunProgress) => void,
): Promise<void> {
  const { suite, baseUrl } = params;

  // A4: suite-defined sampling params with fallback defaults.
  const temperature = suite.temperature ?? params.temperature ?? 0.7;
  const topP = suite.topP ?? params.topP ?? 0.9;
  const jobId = suite.id;

  // Build the full grid of combinations.
  const grid: Array<{
    promptTokens: number;
    genTokens: number;
    concurrency: number;
    repetition: number;
  }> = [];

  for (const pt of suite.promptTokens) {
    for (const gt of suite.genTokens) {
      for (const conc of suite.concurrency) {
        for (let rep = 0; rep < suite.repetitions; rep++) {
          grid.push({ promptTokens: pt, genTokens: gt, concurrency: conc, repetition: rep });
        }
      }
    }
  }

  const totalSamples = grid.length;

  // Persist the run record with jobType (A2) and sampling params (A4).
  const runId = `${jobId}_${Date.now()}`;
  await sql`
    INSERT INTO bench_runs (id, suite_id, job_type, status, started_at, total_samples, temperature, top_p)
    VALUES (${runId}, ${suite.id}, 'bench', 'running', clock_timestamp(), ${totalSamples}, ${temperature}, ${topP})
  `;

  publishJob(emitter, {
    seq,
    jobType: 'bench',
    jobId: runId,
    status: 'running',
    detail: {
      suiteId: suite.id,
      providerId: suite.providerId,
      model: suite.model,
      totalSamples,
    },
  });

  // A5: Warmup pass — 1 request per unique (promptTokens, genTokens) cell, discarded.
  const uniqueCells = new Set<string>();
  for (const item of grid) {
    const cellKey = `${item.promptTokens}_${item.genTokens}`;
    if (!uniqueCells.has(cellKey)) {
      uniqueCells.add(cellKey);
    }
  }
  const warmupPromises = Array.from(uniqueCells).map(async (cellKey) => {
    const parts = cellKey.split('_').map(Number);
    const pt = parts[0] ?? 0;
    const gt = parts[1] ?? 0;
    return runSingleBenchRequest(baseUrl, suite.model, pt, gt, 0, temperature, topP);
  });
  await Promise.allSettled(warmupPromises);

  let completed = 0;
  const samples: BenchSample[] = [];

  // Group by (promptTokens, genTokens, concurrency) for fan-out; each group
  // runs 'repetitions' requests concurrently.
  const groups = new Map<string, typeof grid>();
  for (const item of grid) {
    const key = `${item.promptTokens}_${item.genTokens}_${item.concurrency}`;
    if (!groups.has(key)) {
      groups.set(key, []);
    }
    groups.get(key)!.push(item);
  }

  for (const [_key, group] of groups) {
    const concurrency = group[0]!.concurrency;
    const batchSize = Math.min(concurrency, group.length);

    // Process in batches of 'concurrency' size using Promise.allSettled.
    for (let batchStart = 0; batchStart < group.length; batchStart += batchSize) {
      const batch = group.slice(batchStart, batchStart + batchSize);

      const promises = batch.map(async (item) => {
        const sample = await runSingleBenchRequest(
          baseUrl,
          suite.model,
          item.promptTokens,
          item.genTokens,
          item.repetition,
          temperature,
          topP,
        );
        sample.concurrency = item.concurrency;
        return sample;
      });

      const results = await Promise.allSettled(promises);
      for (const result of results) {
        if (result.status === 'fulfilled') {
          samples.push(result.value);
        }
        completed++;

        // Progress callback
        const current = batch[0]!;
        onProgress({
          jobId: runId,
          totalSamples,
          completedSamples: completed,
          currentPromptTokens: current.promptTokens,
          currentGenTokens: current.genTokens,
          currentConcurrency: current.concurrency,
          currentRepetition: current.repetition,
        });

        publishJob(emitter, {
          seq,
          jobType: 'bench',
          jobId: runId,
          status: 'running',
          detail: {
            completedSamples: completed,
            totalSamples,
            percent: Math.round((completed / totalSamples) * 100),
          },
        });
      }
    }
  }

  // Persist all samples.
  for (const s of samples) {
    await sql`
      INSERT INTO bench_samples (run_id, prompt_tokens, gen_tokens, concurrency, repetition, ttft_ms, total_ms, prompt_tps, gen_tps, cache_n, error)
      VALUES (${runId}, ${s.promptTokens}, ${s.genTokens}, ${s.concurrency}, ${s.repetition}, ${s.ttftMs ?? null}, ${s.totalMs ?? null}, ${s.promptTps ?? null}, ${s.genTps ?? null}, ${s.cacheN ?? null}, ${s.error ?? null})
    `;
  }

  // Compute aggregates.
  const validSamples = samples.filter((s) => !s.error && s.genTps != null);
  const aggregate = computeAggregates(validSamples);

  // A1: Baseline persistence + regression flag.
  // Compare against existing baseline; first run seeds it.
  const baselineRows = await sql<{ aggregate: string }[]>`
    SELECT aggregate FROM bench_baselines
    WHERE provider_id = ${suite.providerId} AND model = ${suite.model}
  `;

  const regressionFlag = computeRegressionFlag(aggregate, baselineRows[0]?.aggregate);

  // Upsert baseline.
  await sql`
    INSERT INTO bench_baselines (provider_id, model, aggregate, run_id)
    VALUES (${suite.providerId}, ${suite.model}, ${sql.json(aggregate as never)}, ${runId})
    ON CONFLICT (provider_id, model) DO UPDATE SET
      aggregate = EXCLUDED.aggregate,
      run_id = EXCLUDED.run_id,
      created_at = clock_timestamp()
  `;

  // Update run record with regression flag.
  await sql`
    UPDATE bench_runs
    SET status = 'completed', finished_at = clock_timestamp(), completed_samples = ${completed},
        aggregate = ${sql.json(aggregate as never)}, regression_flag = ${regressionFlag}
    WHERE id = ${runId}
  `;

  publishJob(emitter, {
    seq,
    jobType: 'bench',
    jobId: runId,
    status: 'completed',
    detail: { ...aggregate, regressionFlag },
  });
}

/**
 * A1: Compute regression flag against baseline.
 * Threshold: gen tok/s -10% = regression, +5% = improvement.
 * N5: guards against divide-by-zero.
 */
export function computeRegressionFlag(
  current: BenchAggregate,
  // Accepts the raw bench_baselines.aggregate value: porsager returns jsonb
  // already-parsed (object), while tests pass a JSON string. jsonbObject handles
  // both. undefined => no baseline row yet => seed.
  baselineJson: unknown,
): 'baseline' | 'regression' | 'improvement' | null {
  if (!current.avgGenTps) return null;
  if (!baselineJson) return 'baseline';

  const baseline = jsonbObject(baselineJson) as BenchAggregate | null;
  if (!baseline) return null;

  if (!baseline.avgGenTps || baseline.avgGenTps === 0) return null;

  const delta = (current.avgGenTps - baseline.avgGenTps) / baseline.avgGenTps;
  if (delta < -0.1) return 'regression';
  if (delta > 0.05) return 'improvement';
  return 'baseline';
}

export interface BenchAggregate {
  avgTtftMs: number | null;
  medianTtftMs: number | null;
  avgGenTps: number | null;
  medianGenTps: number | null;
  avgPromptTps: number | null;
  medianPromptTps: number | null;
  totalSamples: number;
  errorSamples: number;
  p95TtftMs: number | null;
}

export function computeAggregates(samples: BenchSample[]): BenchAggregate {
  if (samples.length === 0) {
    return {
      avgTtftMs: null,
      medianTtftMs: null,
      avgGenTps: null,
      medianGenTps: null,
      avgPromptTps: null,
      medianPromptTps: null,
      totalSamples: 0,
      errorSamples: 0,
      p95TtftMs: null,
    };
  }

  const ttfts = samples.map((s) => s.ttftMs).filter((v): v is number => v != null).sort((a, b) => a - b);
  const genTps = samples.map((s) => s.genTps).filter((v): v is number => v != null).sort((a, b) => a - b);
  const promptTps = samples.map((s) => s.promptTps).filter((v): v is number => v != null).sort((a, b) => a - b);

  const avg = (arr: number[]) => arr.length ? arr.reduce((a, b) => a + b, 0) / arr.length : null;
  const median = (arr: number[]) => {
    if (arr.length === 0) return null;
    const mid = Math.floor(arr.length / 2);
    return arr.length % 2 ? arr[mid]! : (arr[mid - 1]! + arr[mid]!) / 2;
  };
  const p95 = (arr: number[]) => {
    if (arr.length === 0) return null;
    const idx = Math.ceil(arr.length * 0.95) - 1;
    return arr[Math.max(0, idx)] ?? null;
  };

  return {
    avgTtftMs: avg(ttfts),
    medianTtftMs: median(ttfts),
    avgGenTps: avg(genTps),
    medianGenTps: median(genTps),
    avgPromptTps: avg(promptTps),
    medianPromptTps: median(promptTps),
    totalSamples: samples.length,
    errorSamples: samples.filter((s) => s.error).length,
    p95TtftMs: p95(ttfts),
  };
}
