/**
 * llmindex.io — eval batch runner: perturbed items → OpenRouter → graded responses
 * Author:  Simon-Pierre Boucher
 * Contact: contact@spboucher.ai
 * License: Proprietary — © Simon-Pierre Boucher, all rights reserved
 */
import { createHash } from 'node:crypto';
import { Resvg } from '@resvg/resvg-js';
import { prisma } from '@llmindex/db';
import { extractAnswer, extractBlockAnswer, generateBatch, gradeAnswer } from '@llmindex/items';
import { OpenRouterClient, type ChatMessage, type Pricing } from '@llmindex/openrouter';
import { INDEX_VERSION, type Domain } from '@llmindex/scoring';
import { maxRunCostUsd } from './env';

const CONCURRENCY = 4;
/** Batch is flagged degraded when more than 2% of calls fail (§6). */
export const DEGRADED_FAILURE_RATE = 0.02;
/** Conservative per-call token estimate for the pre-flight cost gate. */
export const EST_TOKENS = { prompt: 1200, completion: 1500 };
/** Completion budget: reasoning models need headroom or answers truncate (measured in the eval-harness literature). */
export const MAX_COMPLETION_TOKENS = 16384;

export interface EvalBatchOptions {
  modelSlug: string;
  domain: Domain;
  n: number;
  /** k>1 adds consistency samples at the model's default temperature. */
  kSamples?: number;
  seed?: string;
  /** Live progress callback (completed calls, total calls) — best-effort. */
  onProgress?: (done: number, total: number) => void | Promise<void>;
}

export function estimateBatchCostUsd(
  n: number,
  kSamples: number,
  pricing: Pricing,
): number {
  const calls = n + n * Math.max(0, kSamples - 1);
  return (
    (calls * (EST_TOKENS.prompt * pricing.promptPerM + EST_TOKENS.completion * pricing.completionPerM)) /
    1_000_000
  );
}

export class CostCapError extends Error {}

async function mapLimit<T, R>(
  items: T[],
  limit: number,
  fn: (item: T, index: number) => Promise<R>,
): Promise<R[]> {
  const results: R[] = new Array(items.length);
  let next = 0;
  const lanes = Array.from({ length: Math.min(limit, items.length) }, async () => {
    for (;;) {
      const i = next++;
      if (i >= items.length) return;
      results[i] = await fn(items[i]!, i);
    }
  });
  await Promise.all(lanes);
  return results;
}

export async function runEvalBatch(opts: EvalBatchOptions): Promise<{ runId: string; status: string }> {
  const kSamples = opts.kSamples ?? 1;
  const model = await prisma.model.findUnique({ where: { slug: opts.modelSlug } });
  if (!model) throw new Error(`Model not in DB (run db:seed to sync): ${opts.modelSlug}`);
  if (model.promptPricePerM == null || model.completionPricePerM == null) {
    throw new Error(`Model has no pricing (cost tracking mandatory): ${opts.modelSlug}`);
  }
  const pricing: Pricing = {
    promptPerM: model.promptPricePerM,
    completionPerM: model.completionPricePerM,
  };

  const estimated = estimateBatchCostUsd(opts.n, kSamples, pricing);
  const cap = maxRunCostUsd();
  if (estimated > cap) {
    throw new CostCapError(
      `Estimated batch cost $${estimated.toFixed(2)} exceeds MAX_RUN_COST_USD=$${cap} — refusing to start`,
    );
  }

  const seed = opts.seed ?? `batch:${Date.now()}`;
  const items = generateBatch({ domain: opts.domain, n: opts.n, seed });
  const itemSetHash = createHash('sha256')
    .update(items.map((i) => i.perturbSeed).join('|'))
    .digest('hex');

  const run = await prisma.scoreRun.create({
    data: {
      indexVersion: INDEX_VERSION,
      kind: 'eval_batch',
      status: 'running',
      itemSetHash,
      modelSet: [model.slug],
      notes: `domain=${opts.domain} n=${opts.n} k=${kSamples} seed=${seed}`,
    },
  });

  const dbItems = await Promise.all(
    items.map((item) =>
      prisma.evalItem.create({
        data: {
          templateId: item.templateId,
          domain: item.domain,
          prompt: item.prompt,
          answerKey: item.answerKey,
          perturbSeed: item.perturbSeed,
          isAnchor: item.isAnchor,
        },
      }),
    ),
  );

  const client = new OpenRouterClient();
  let failures = 0;
  let completedCalls = 0;
  // Small-context models reject max_tokens near/above their window (HTTP 400):
  // cap the completion budget by the context length, keeping prompt headroom.
  const maxTokens = model.contextLength
    ? Math.min(MAX_COMPLETION_TOKENS, Math.max(2048, model.contextLength - 4096))
    : MAX_COMPLETION_TOKENS;

  interface Call {
    itemIdx: number;
    sampleIndex: number;
  }
  const calls: Call[] = [];
  for (let i = 0; i < items.length; i++) {
    for (let s = 0; s < kSamples; s++) calls.push({ itemIdx: i, sampleIndex: s });
  }

  await mapLimit(calls, CONCURRENCY, async ({ itemIdx, sampleIndex }) => {
    const item = items[itemIdx]!;
    const dbItem = dbItems[itemIdx]!;
    // Vision items: rasterize the generated SVG scene → PNG data URL.
    let messages: ChatMessage[];
    if (item.svg) {
      const png = new Resvg(item.svg, { fitTo: { mode: 'width', value: 1024 } }).render().asPng();
      messages = [
        {
          role: 'user',
          content: [
            { type: 'text', text: item.prompt },
            { type: 'image_url', image_url: { url: `data:image/png;base64,${png.toString('base64')}` } },
          ],
        },
      ];
    } else {
      messages = [{ role: 'user', content: item.prompt }];
    }
    // Scored sample: temperature 0. Consistency samples: model default temperature (§6).
    const requestParams =
      sampleIndex === 0
        ? { model: model.slug, messages, temperature: 0, max_tokens: maxTokens }
        : { model: model.slug, messages, max_tokens: maxTokens };
    // Audit copy: keep exact params but elide the base64 image payload (the
    // scene is reproducible from the item's stored SVG/perturb seed).
    const auditParams = item.svg
      ? {
          ...requestParams,
          messages: [{ role: 'user', content: [{ type: 'text', text: item.prompt }, { type: 'image_url', image_url: { url: '[generated png elided — reproducible from eval_item svg]' } }] }],
        }
      : requestParams;
    try {
      const result = await client.chat(requestParams, pricing);
      const blockMode = item.grading === 'json' || item.grading === 'lines';
      const { answer, confidence } = blockMode
        ? extractBlockAnswer(result.text)
        : extractAnswer(result.text);
      // Truncated completions are unscored (correct=null), never "wrong":
      // grading a cut-off chain of thought measures the budget, not the model.
      const truncated = result.raw.choices?.[0]?.finish_reason === 'length' && answer === null;
      const correct = truncated ? null : gradeAnswer(answer, item.answerKey, item.grading);
      await prisma.modelResponse.create({
        data: {
          runId: run.id,
          modelId: model.id,
          itemId: dbItem.id,
          sampleIndex,
          requestParams: auditParams as object,
          rawResponse: result.raw as unknown as object,
          answerExtracted: answer,
          correct,
          error: truncated ? 'truncated' : null,
          confidence,
          tokensIn: result.usage?.prompt_tokens ?? null,
          tokensOut: result.usage?.completion_tokens ?? null,
          latencyMs: result.latencyMs,
          costUsd: result.costUsd,
        },
      });
    } catch (err) {
      failures += 1;
      await prisma.modelResponse.create({
        data: {
          runId: run.id,
          modelId: model.id,
          itemId: dbItem.id,
          sampleIndex,
          requestParams: auditParams as object,
          rawResponse: {},
          error: String(err).slice(0, 2000),
        },
      });
    }
    completedCalls += 1;
    if (opts.onProgress && (completedCalls % 5 === 0 || completedCalls === calls.length)) {
      try {
        await opts.onProgress(completedCalls, calls.length);
      } catch {
        /* progress is best-effort */
      }
    }
  });

  const failureRate = failures / calls.length;
  const status = failureRate > DEGRADED_FAILURE_RATE ? 'degraded' : 'complete';
  await prisma.scoreRun.update({
    where: { id: run.id },
    data: {
      status,
      completedAt: new Date(),
      fitDiagnostics: { calls: calls.length, failures, failureRate },
    },
  });
  console.log(
    `[eval] run ${run.id} ${status}: model=${model.slug} domain=${opts.domain} calls=${calls.length} failures=${failures}`,
  );
  return { runId: run.id, status };
}