import { afterAll, describe, expect, it } from 'vitest';
import { createOptimizationPresetRules } from './rule-service';
import { frenchPublicHolidaysForYear, isWorkingPlanningDate } from './declared-non-working-days';
import { OPTIMIZATION_CATALOG } from './planning-objectives';
import {
	calculatePlanningAsyncInternal,
	calculatePlanningInternal,
	type PlanningArtifacts,
	type PlanningOutcome,
} from './leave-planner.core';
import {
	calculatePlanningAsyncInternal as frozenCalculateAsync,
	calculatePlanningInternal as frozenCalculate,
} from './frozen-engine/leave-planner.core.frozen.mjs';

/** Rebuild instructions for the frozen engine: docs/benchmark/README.md. */
const BASE_SEED = 280_000;
const CASES = 5000;
const CHUNK = 200;

function mulberry32(seed: number): () => number {
	let a = seed >>> 0;
	return () => {
		a = (a + 0x6d2b79f5) >>> 0;
		let t = a;
		t = Math.imul(t ^ (t >>> 15), t | 1);
		t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
		return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
	};
}

const DAY = 86_400_000;
const weekday = (ms: number) => ![0, 6].includes(new Date(ms).getUTCDay());
const iso = (ms: number) => new Date(ms).toISOString().slice(0, 10);

type Req = Record<string, unknown> & { rules: Record<string, unknown>[] };

function generateRequest(seed: number): Req {
	const rnd = mulberry32(seed);
	const int = (min: number, max: number) => min + Math.floor(rnd() * (max - min + 1));
	const pick = <T>(items: readonly T[]): T => items[int(0, items.length - 1)];
	const chance = (p: number) => rnd() < p;
	const shuffle = <T>(items: T[]): T[] => {
		for (let i = items.length - 1; i > 0; i--) {
			const j = int(0, i);
			[items[i], items[j]] = [items[j], items[i]];
		}
		return items;
	};

	const long = chance(0.1);
	const mode = rnd();
	let startMs = Date.UTC(2025, 0, 1) + int(0, 1094) * DAY;
	if (mode < 0.15) startMs = Date.UTC(pick([2025, 2026]), 11, int(8, 31));
	else if (mode < 0.3) {
		const year = pick([2025, 2026, 2027]);
		const holidays = [...frenchPublicHolidaysForYear(year)].filter((d) =>
			weekday(Date.parse(d)),
		);
		startMs = Date.parse(pick(holidays)) - int(1, 10) * DAY;
	}
	const days = long ? int(29, 70) : int(7, 28);
	const start = iso(startMs);
	const date = (offset: number) => iso(startMs + offset * DAY);
	const types = ['cp', 'rtt', 'other'] as const;

	const balances = Array.from({ length: int(1, 3) }, (_, i) => ({
		id: `b${i}`,
		type: pick(types),
		quantity: int(0, long ? 4 : 3),
		...(chance(0.3) ? { availableFrom: date(int(1, days)) } : {}),
		expiryDate: date(int(-3, days + 40)),
	})).map((b) =>
		b.availableFrom && b.availableFrom > b.expiryDate
			? { ...b, expiryDate: b.availableFrom }
			: b,
	);

	const declaredDays: { date: string; type: string }[] = [];
	for (let n = int(0, 4); n > 0; n--) {
		const ms = startMs + int(0, days) * DAY;
		const d = iso(ms);
		if (isWorkingPlanningDate(d) && !declaredDays.some((x) => x.date === d))
			declaredDays.push({ date: d, type: pick(types) });
	}

	const selectors = ['cp', 'rtt', 'other', 'any'] as const;
	const rules: Record<string, unknown>[] = [];
	for (let n = long ? int(0, 1) : int(0, 2); n > 0; n--) {
		const from = int(0, days - 1);
		rules.push({
			id: `tw${rules.length}`,
			type: chance(0.5) ? 'time_window' : 'TIME_WINDOW',
			force: pick(['obligation', 'preference']),
			startDate: date(from),
			endDate: date(from + int(0, 10)),
			placeDays: int(1, 4),
			placeType: pick(selectors),
		});
	}
	for (let n = long && rules.length ? 0 : long ? int(0, 1) : int(0, 2); n > 0; n--) {
		rules.push({
			id: `it${rules.length}`,
			type: chance(0.5) ? 'if_then' : 'IF_THEN',
			force: pick(['obligation', 'preference']),
			dayOfWeek: int(1, 5),
			placeDays: int(1, 3),
			placeType: pick(selectors),
			placement: pick(['before', 'after', 'around', 'auto']),
			rewardDays: int(0, 2),
			rewardType: pick(types),
		});
	}
	const presets: Record<string, unknown>[] = shuffle(
		createOptimizationPresetRules() as unknown as Record<string, unknown>[],
	).slice(0, long ? int(1, 4) : int(0, 3));
	rules.push(...presets);
	if (!long && chance(0.3)) {
		const row = pick(OPTIMIZATION_CATALOG);
		const conditions = (row.conditions as readonly string[]).filter(() => chance(0.5));
		rules.push({
			id: 'composed',
			type: 'optimization',
			force: 'preference',
			configuration: {
				elements: row.elements,
				conditions,
				measure: row.measure,
				aggregation: pick(row.aggregations),
				direction: pick(['minimize', 'maximize']),
				...(row.measure === 'work-sequence-deficit' ? { thresholdDays: int(1, 8) } : {}),
			},
		});
	}
	for (const rule of shuffle(rules)) if (chance(0.08)) rule['enabled'] = false;

	const request: Req = {
		calculationDate: start,
		horizon: { start, end: date(days) },
		balances,
		declaredDays,
		rules,
	};
	if (chance(0.04)) {
		const breakIt = int(0, 3);
		if (breakIt === 0) balances[0].quantity = -1;
		else if (breakIt === 1) request['horizon'] = { start: date(days), end: start };
		else if (breakIt === 2) rules.push({ id: 'bad', type: 'unknown_type' });
		else declaredDays.push({ date: start, type: 'nope' });
	}
	return request;
}

interface Run {
	outcome: PlanningOutcome;
	artifacts?: Pick<PlanningArtifacts, 'explanation' | 'finalState' | 'explanationError'>;
}

function run(
	calculate: (
		request: unknown,
		c?: undefined,
		a?: (artifacts: PlanningArtifacts) => void,
	) => PlanningOutcome,
	request: Req,
): Run {
	let artifacts: Run['artifacts'];
	const outcome = calculate(structuredClone(request), undefined, (a) => {
		artifacts = {
			explanation: a.explanation,
			finalState: a.finalState,
			explanationError: a.explanationError,
		};
	});
	return { outcome, artifacts };
}

const kinds = new Map<string, number>();
const ruleTypes = new Set<string>();
const count = (key: string) => kinds.set(key, (kinds.get(key) ?? 0) + 1);

let slowest = { seed: 0, ms: 0 };

function crossesYearAndHoliday(request: Req): void {
	const { start, end } = request['horizon'] as { start: string; end: string };
	if (start.slice(0, 4) !== end.slice(0, 4)) count('crosses-year');
	for (let ms = Date.parse(start); ms <= Date.parse(end); ms += DAY)
		if (
			weekday(ms) &&
			frenchPublicHolidaysForYear(new Date(ms).getUTCFullYear()).has(iso(ms))
		) {
			count('holiday-week');
			return;
		}
}

function assertSameAsFrozen(seed: number): void {
	const request = generateRequest(seed);
	const t0 = performance.now();
	const current = run(calculatePlanningInternal, request);
	const frozen = run(frozenCalculate, request);
	const ms = performance.now() - t0;
	if (ms > slowest.ms) slowest = { seed, ms: Math.round(ms) };
	crossesYearAndHoliday(request);
	count(current.outcome.kind);
	for (const rule of request.rules) ruleTypes.add(String(rule['type']).toLowerCase());
	if ((request['declaredDays'] as unknown[]).length) count('declared');
	if (request.rules.some((r) => r['rewardDays'])) count('reward');
	try {
		expect(current).toStrictEqual(frozen);
	} catch (error) {
		throw new Error(
			`Engine differs from frozen engine. seed=${seed}\nrequest=${JSON.stringify(request)}\n${
				(error as Error).message
			}`,
			{ cause: error },
		);
	}
}

type AsyncEngine = typeof calculatePlanningAsyncInternal;

/** Cancels after a seed-derived step count: during preparation, mid-search, or never. */
function cancelAfter(seed: number, total: number): number {
	const rnd = mulberry32(seed ^ 0x5bd1e995);
	const roll = rnd();
	if (roll < 0.1) return Math.floor(rnd() * total);
	if (roll < 0.25) return Number.POSITIVE_INFINITY;
	// Preparation is most of the steps; spend the rest of the budget on the search.
	return Math.floor(total * (0.9 + rnd() * 0.1));
}

async function runCancelled(engine: AsyncEngine, request: Req, limit: number) {
	const calls: unknown[] = [];
	const artifacts: unknown[] = [];
	const explanations: unknown[] = [];
	const outcomes: unknown[] = [];
	let polls = 0;
	const outcome = await engine(
		structuredClone(request),
		(plan, state) => calls.push({ plan, state }),
		({ timings: _timings, ...rest }) => artifacts.push(rest),
		(o) => outcomes.push(o),
		(...args) => explanations.push(args),
		() => ++polls > limit,
	);
	return { outcome, calls, artifacts, explanations, outcomes, polls };
}

const cancelledKinds = new Map<string, number>();
let incumbents = 0;

async function assertCancelledSameAsFrozen(seed: number): Promise<void> {
	const request = generateRequest(seed);
	const total = (await runCancelled(calculatePlanningAsyncInternal, request, Infinity)).polls;
	const limit = cancelAfter(seed, total);
	const current = await runCancelled(calculatePlanningAsyncInternal, request, limit);
	const frozen = await runCancelled(frozenCalculateAsync, request, limit);
	const kind = current.outcome.kind;
	cancelledKinds.set(kind, (cancelledKinds.get(kind) ?? 0) + 1);
	if (kind === 'PlanningIncomplete' && current.calls.length) incumbents++;
	try {
		expect(current).toStrictEqual(frozen);
	} catch (error) {
		throw new Error(
			`Cancelled run differs from frozen engine. seed=${seed} limit=${limit}\nrequest=${JSON.stringify(request)}\n${
				(error as Error).message
			}`,
			{ cause: error },
		);
	}
}

const ONLY_SEED = process.env['FROZEN_ENGINE_SEED'];

describe('frozen engine differential check', () => {
	if (ONLY_SEED) {
		it(`matches the frozen engine for seed ${ONLY_SEED}`, async () => {
			assertSameAsFrozen(Number(ONLY_SEED));
			await assertCancelledSameAsFrozen(Number(ONLY_SEED));
		}, 120_000);
		return;
	}
	for (let from = 0; from < CASES; from += CHUNK) {
		it(`matches the frozen engine for seeds ${BASE_SEED + from}..${BASE_SEED + from + CHUNK - 1}`, () => {
			for (let i = from; i < from + CHUNK; i++) assertSameAsFrozen(BASE_SEED + i);
		}, 120_000);
	}

	for (let from = 0; from < CASES; from += CHUNK * 2) {
		it(`matches the frozen engine when cancelled, seeds ${BASE_SEED + from}.. every 10th`, async () => {
			for (let i = from; i < from + CHUNK * 2; i += 10)
				await assertCancelledSameAsFrozen(BASE_SEED + i);
		}, 300_000);
	}

	it('retains incumbents of cancelled calculations', () => {
		console.info('cancelled coverage', Object.fromEntries(cancelledKinds), { incumbents });
		expect(cancelledKinds.get('PlanningIncomplete') ?? 0).toBeGreaterThan(20);
		expect(incumbents).toBeGreaterThan(10);
	});

	afterAll(() => {
		console.info(
			'frozen-engine coverage of generated cases',
			Object.fromEntries(kinds),
			[...ruleTypes],
			'slowest case',
			slowest,
		);
	});

	it('generates every outcome kind and rule type', () => {
		for (const kind of [
			'Plan',
			'PlanningConflict',
			'InvalidInput',
			'declared',
			'reward',
			'crosses-year',
			'holiday-week',
		])
			expect(kinds.get(kind) ?? 0, kind).toBeGreaterThan(20);
		expect(['if_then', 'optimization', 'time_window'].filter((t) => !ruleTypes.has(t))).toEqual(
			[],
		);
	});
});