import type { TaskAttempt, TaskRuntime, UsageStats, WorkerHistoryMessage } from "./types.ts";
import { emptyUsage } from "./types.ts";

function refreshTotal(usage: UsageStats): UsageStats {
	usage.totalTokens = usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
	return usage;
}

export function formatTokenCount(value: number): string {
	if (!Number.isFinite(value)) return "0";
	const absolute = Math.abs(value);
	if (absolute < 1_000) return String(Math.round(value));
	const units = ["k", "m", "b"];
	let scaled = value;
	let unitIndex = -1;
	while (Math.abs(scaled) >= 1_000 && unitIndex < units.length - 1) {
		scaled /= 1_000;
		unitIndex += 1;
	}
	const formatted = scaled.toFixed(2).replace(/\.?(0+)$/, "");
	return `${formatted}${units[unitIndex]}`;
}

export function cacheHitRate(usage: UsageStats): number {
	const promptTokens = usage.input + usage.cacheRead + usage.cacheWrite;
	return promptTokens > 0 ? usage.cacheRead / promptTokens : 0;
}

export function formatCacheHitRate(usage: UsageStats): string {
	return `${(cacheHitRate(usage) * 100).toFixed(1)}%`;
}

export function aggregateUsage(values: Iterable<UsageStats | undefined>): UsageStats {
	const result = emptyUsage();
	for (const usage of values) {
		if (!usage) continue;
		result.input += usage.input;
		result.output += usage.output;
		result.cacheRead += usage.cacheRead;
		result.cacheWrite += usage.cacheWrite;
		result.cost += usage.cost;
		result.contextTokens = Math.max(result.contextTokens, usage.contextTokens);
		result.turns += usage.turns;
	}
	return refreshTotal(result);
}

export function usageFromHistory(history: readonly WorkerHistoryMessage[]): UsageStats {
	const result = emptyUsage();
	for (const message of history) {
		const usage = message.usage;
		if (!usage) continue;
		result.input += usage.input ?? 0;
		result.output += usage.output ?? 0;
		result.cacheRead += usage.cacheRead ?? 0;
		result.cacheWrite += usage.cacheWrite ?? 0;
		result.cost += usage.cost?.total ?? 0;
		if (message.role === "assistant") {
			result.contextTokens = usage.totalTokens ?? (usage.input ?? 0) + (usage.output ?? 0) + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0);
			result.turns += 1;
		}
	}
	return refreshTotal(result);
}

export function attemptUsage(attempt: TaskAttempt): UsageStats {
	return aggregateUsage([attempt.worker?.usage, attempt.reviewer?.usage]);
}

export function taskUsage(task: TaskRuntime): UsageStats {
	return aggregateUsage(task.attempts.map(attemptUsage));
}

export function clusterUsage(tasks: readonly TaskRuntime[]): UsageStats {
	return aggregateUsage(tasks.map(taskUsage));
}