#!/usr/bin/env bun
import { Database } from "bun:sqlite";
/**
 * Title-generation benchmark harness.
 *
 * Samples random first-of-session messages from the local history DB, renders
 * the shipped `title-system.md` prompt, and runs every message against a matrix
 * of title models — the on-device ONNX models (LFM2 350M/700M, Gemma 270M) via
 * the tiny-title worker, plus a remote Ollama model (Llama 3.2 3B by default).
 * Each model lane runs concurrently; within a lane requests are sequential
 * because the local worker serializes generation on one pipeline.
 *
 * Results (per-sample titles + latency, plus per-model summaries) are written
 * to a timestamped JSON file so runs can be compared later.
 *
 * Usage:
 *   bun scripts/bench-title-models.ts
 *   bun scripts/bench-title-models.ts --count 30 --seed 42
 *   bun scripts/bench-title-models.ts --models lfm2-350m,gemma-270m
 *   bun scripts/bench-title-models.ts --ollama-url http://spark.internal:11434 --ollama-models llama3.2:3b,lfm2:2.6b
 *   bun scripts/bench-title-models.ts --db ~/.omp/agent/history.db --out bench.json
 */
import * as os from "node:os";
import * as path from "node:path";
import { prompt } from "@oh-my-pi/pi-utils";
import titleSystemPrompt from "../src/prompts/system/title-system.md" with { type: "text" };
import { preprocessTinyMessage } from "../src/tiny/message-preproc";
import { isTinyTitleLocalModelKey } from "../src/tiny/models";
import { normalizeGeneratedTitle } from "../src/tiny/text";
import { shutdownTinyTitleClient, tinyTitleClient } from "../src/tiny/title-client";

/** A sampled prompt with the cleaned text actually fed to the models. */
interface PreparedPrompt {
	id: number;
	raw: string;
	input: string;
}

/** One title produced for one input by one model, with wall-clock latency. */
interface BenchSample {
	id: number;
	input: string;
	title: string | null;
	ms: number;
}

/** All samples for one model plus the aggregate quality/latency summary. */
interface BenchLane {
	model: string;
	transport: "local" | "ollama";
	samples: BenchSample[];
	summary: BenchSummary;
}

/** Aggregate stats for a lane; latency percentiles skip the cold first call. */
interface BenchSummary {
	count: number;
	nulls: number;
	coldMs: number;
	warmMeanMs: number;
	warmMedianMs: number;
	warmP95Ms: number;
	lengthCompliant: string;
	punctuationFree: string;
}

interface BenchConfig {
	dbPath: string;
	count: number;
	seed: number;
	localModels: string[];
	ollamaUrl: string | null;
	ollamaModels: string[];
	outPath: string;
}

const DEFAULT_LOCAL_MODELS = ["lfm2-350m", "lfm2-700m", "gemma-270m"];
const DEFAULT_OLLAMA_URL = "http://spark.internal:11434";
const DEFAULT_OLLAMA_MODELS = ["llama3.2:3b", "lfm2:2.6b"];
const MIN_INPUT_CHARS = 10;
const MAX_INPUT_CHARS = 800;

/** System prompt with examples (used for the capable Ollama model). */
const TITLE_PROMPT_WITH_EXAMPLES = prompt.render(titleSystemPrompt, { includeExamples: true });
/** Example-free prompt matching what the on-device worker ships to tiny models. */
const TITLE_PROMPT_NO_EXAMPLES = prompt.render(titleSystemPrompt, { includeExamples: false });

/** Deterministic mulberry32 PRNG so `--seed` reproduces a sample set. */
function createRng(seed: number): () => number {
	let state = seed >>> 0;
	return () => {
		state |= 0;
		state = (state + 0x6d2b79f5) | 0;
		let t = Math.imul(state ^ (state >>> 15), 1 | state);
		t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
		return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
	};
}

/** Pick `count` distinct random first-of-session prompts within the size band. */
function sampleHistoryPrompts(dbPath: string, count: number, rng: () => number): { id: number; prompt: string }[] {
	const db = new Database(dbPath, { readonly: true });
	try {
		const rows = db
			.query(
				`WITH firsts AS (
					SELECT session_id, MIN(id) AS id FROM history
					WHERE session_id IS NOT NULL
					GROUP BY session_id
				)
				SELECT h.id AS id, h.prompt AS prompt
				FROM history h JOIN firsts ON firsts.id = h.id
				WHERE length(trim(h.prompt)) BETWEEN ? AND ?`,
			)
			.all(MIN_INPUT_CHARS, MAX_INPUT_CHARS) as { id: number; prompt: string }[];
		const seen = new Set<string>();
		const unique: { id: number; prompt: string }[] = [];
		for (const row of rows) {
			const key = row.prompt.trim();
			if (seen.has(key)) continue;
			seen.add(key);
			unique.push(row);
		}
		// Fisher–Yates with the seeded RNG, then take the first `count`.
		for (let i = unique.length - 1; i > 0; i--) {
			const j = Math.floor(rng() * (i + 1));
			[unique[i], unique[j]] = [unique[j], unique[i]];
		}
		return unique.slice(0, Math.min(count, unique.length));
	} finally {
		db.close();
	}
}

/** Run one local ONNX model over every prompt (sequential; worker is single-lane). */
async function runLocalLane(model: string, prompts: PreparedPrompt[]): Promise<BenchSample[]> {
	const samples: BenchSample[] = [];
	for (const item of prompts) {
		const started = performance.now();
		const title = await tinyTitleClient.generate(model, item.input, { systemPrompt: TITLE_PROMPT_NO_EXAMPLES });
		samples.push({ id: item.id, input: item.input, title, ms: performance.now() - started });
	}
	return samples;
}

/** Extract the `<title>` payload from a free-form chat completion. */
function parseChatTitle(text: string, sourceText: string): string | null {
	if (!text || /<title\s*\/>/i.test(text)) return null;
	const closed = /<title>([\s\S]*?)<\/title>/i.exec(text);
	const open = closed ? null : /<title>([\s\S]*)/i.exec(text);
	return normalizeGeneratedTitle(closed?.[1] ?? open?.[1] ?? text, sourceText);
}

/** Run one Ollama chat model over every prompt via the /api/chat endpoint. */
async function runOllamaLane(baseUrl: string, model: string, prompts: PreparedPrompt[]): Promise<BenchSample[]> {
	const samples: BenchSample[] = [];
	for (const item of prompts) {
		const started = performance.now();
		const response = await fetch(new URL("/api/chat", baseUrl), {
			method: "POST",
			headers: { "content-type": "application/json" },
			body: JSON.stringify({
				model,
				stream: false,
				keep_alive: "10m",
				messages: [
					{ role: "system", content: TITLE_PROMPT_WITH_EXAMPLES },
					{ role: "user", content: `<user>\n${item.input}\n</user>` },
				],
				options: { temperature: 0, num_predict: 32 },
			}),
		});
		if (!response.ok) throw new Error(`Ollama ${response.status}: ${await response.text()}`);
		const payload = (await response.json()) as { message?: { content?: string } };
		const raw = payload.message?.content ?? "";
		samples.push({
			id: item.id,
			input: item.input,
			title: parseChatTitle(raw, item.input),
			ms: performance.now() - started,
		});
	}
	return samples;
}

/** Fold a lane's samples into latency percentiles and title-quality ratios. */
function summarize(samples: BenchSample[]): BenchSummary {
	const warm = samples
		.slice(1)
		.map(sample => sample.ms)
		.sort((a, b) => a - b);
	const outputs = samples.filter(sample => sample.title !== null);
	const percentile = (sorted: number[], q: number): number =>
		sorted.length === 0 ? 0 : sorted[Math.min(sorted.length - 1, Math.max(0, Math.ceil(sorted.length * q) - 1))];
	const wordCompliant = outputs.filter(sample => {
		const words = sample.title!.trim().split(/\s+/).length;
		return words >= 3 && words <= 7;
	}).length;
	const punctuationFree = outputs.filter(sample => !/\p{P}/u.test(sample.title!)).length;
	return {
		count: samples.length,
		nulls: samples.length - outputs.length,
		coldMs: Number((samples[0]?.ms ?? 0).toFixed(1)),
		warmMeanMs: Number((warm.reduce((sum, value) => sum + value, 0) / (warm.length || 1)).toFixed(1)),
		warmMedianMs: Number(percentile(warm, 0.5).toFixed(1)),
		warmP95Ms: Number(percentile(warm, 0.95).toFixed(1)),
		lengthCompliant: `${wordCompliant}/${outputs.length}`,
		punctuationFree: `${punctuationFree}/${outputs.length}`,
	};
}

function parseArgs(argv: string[]): BenchConfig {
	const get = (flag: string): string | undefined => {
		const index = argv.indexOf(flag);
		return index >= 0 ? argv[index + 1] : undefined;
	};
	const has = (flag: string): boolean => argv.includes(flag);
	const modelsArg = get("--models");
	const ollamaModelsArg = get("--ollama-models");
	const ollamaUrlArg = get("--ollama-url");
	const stamp = new Date().toISOString().replace(/[:.]/g, "-");
	return {
		dbPath: (get("--db") ?? path.join(os.homedir(), ".omp/agent/history.db")).replace(/^~/, os.homedir()),
		count: Number(get("--count") ?? 20),
		seed: Number(get("--seed") ?? Date.now() & 0xffffffff),
		localModels: modelsArg
			? modelsArg
					.split(",")
					.map(model => model.trim())
					.filter(Boolean)
			: DEFAULT_LOCAL_MODELS,
		ollamaUrl: has("--no-ollama") ? null : (ollamaUrlArg ?? DEFAULT_OLLAMA_URL),
		ollamaModels: ollamaModelsArg
			? ollamaModelsArg
					.split(",")
					.map(model => model.trim())
					.filter(Boolean)
			: DEFAULT_OLLAMA_MODELS,
		outPath: get("--out") ?? path.join(os.tmpdir(), `title-bench-${stamp}.json`),
	};
}

async function main(): Promise<void> {
	const config = parseArgs(Bun.argv.slice(2));
	const rng = createRng(config.seed);
	const rows = sampleHistoryPrompts(config.dbPath, config.count, rng);
	if (rows.length === 0) throw new Error(`No history prompts found in ${config.dbPath}`);
	const prepared: PreparedPrompt[] = rows.map(row => ({
		id: row.id,
		raw: row.prompt,
		input: preprocessTinyMessage(row.prompt),
	}));

	const invalidLocal = config.localModels.filter(model => !isTinyTitleLocalModelKey(model));
	if (invalidLocal.length > 0) throw new Error(`Unknown local title model(s): ${invalidLocal.join(", ")}`);

	console.info(`Benchmarking ${rows.length} prompts (seed ${config.seed}) from ${config.dbPath}`);

	// Each model is its own concurrent lane; the local worker still serializes
	// its own lanes internally, but the Ollama lane genuinely runs in parallel.
	const laneTasks: Promise<BenchLane>[] = [
		...config.localModels.map(async (model): Promise<BenchLane> => {
			const samples = await runLocalLane(model, prepared);
			return { model, transport: "local", samples, summary: summarize(samples) };
		}),
	];
	if (config.ollamaUrl) {
		const url = config.ollamaUrl;
		for (const model of config.ollamaModels) {
			laneTasks.push(
				(async (): Promise<BenchLane> => {
					const samples = await runOllamaLane(url, model, prepared);
					return { model: `${model}@ollama`, transport: "ollama", samples, summary: summarize(samples) };
				})(),
			);
		}
	}

	const settled = await Promise.allSettled(laneTasks);
	const lanes: BenchLane[] = [];
	for (const result of settled) {
		if (result.status === "fulfilled") lanes.push(result.value);
		else
			console.error(
				`Lane failed: ${result.reason instanceof Error ? result.reason.message : String(result.reason)}`,
			);
	}

	await shutdownTinyTitleClient();

	// Prompt-centric view: each row is one input with every model's title beside it.
	const matrix = prepared.map(item => {
		const titles: Record<string, string> = {};
		for (const lane of lanes) titles[lane.model] = lane.samples.find(sample => sample.id === item.id)?.title ?? "∅";
		return { id: item.id, raw: item.raw, input: item.input, titles };
	});

	const report = {
		generatedAt: new Date().toISOString(),
		config: { ...config, prompts: prepared },
		matrix,
		lanes,
	};
	await Bun.write(config.outPath, JSON.stringify(report, null, 2));

	for (const entry of matrix) {
		console.info(`\n[#${entry.id}] ${entry.raw.replace(/\s+/g, " ").slice(0, 140)}`);
		if (entry.input !== entry.raw.trim())
			console.info(`  (cleaned) ${entry.input.replace(/\s+/g, " ").slice(0, 140)}`);
		console.table(Object.fromEntries(lanes.map(lane => [lane.model, { output: entry.titles[lane.model] }])));
	}

	console.info("\nSummary:");
	console.table(
		Object.fromEntries(
			lanes.map(lane => [
				lane.model,
				{
					cold: lane.summary.coldMs,
					warmMean: lane.summary.warmMeanMs,
					warmP95: lane.summary.warmP95Ms,
					nulls: lane.summary.nulls,
					len3to7: lane.summary.lengthCompliant,
					punctFree: lane.summary.punctuationFree,
				},
			]),
		),
	);
	console.info(`\nWrote ${config.outPath}`);
}

await main();