import { calculateCost, type ModelPricing } from './pricing.js'
import type { CostEntry, CostTrend, ProviderCostComparison, TaskCostSummary } from './types.js'

/**
 * In-memory cost tracker that accumulates CostEntry records.
 *
 * Provides summarization, provider comparison, and trending.
 * Serializable via toJSON/fromJSON for persistence.
 */
export class CostTracker {
	private entries: CostEntry[] = []
	private customPricing?: Record<string, ModelPricing>

	constructor(customPricing?: Record<string, ModelPricing>) {
		this.customPricing = customPricing
	}

	/**
	 * Record a cost entry. Cost is calculated automatically from token counts and model pricing.
	 */
	record(entry: Omit<CostEntry, 'costUsd'>): void {
		const { costUsd } = calculateCost(
			entry.inputTokens,
			entry.outputTokens,
			entry.modelId,
			this.customPricing,
		)

		this.entries.push({ ...entry, costUsd })
	}

	/**
	 * Summarize all cost entries for a given task.
	 */
	summarize(taskId: string): TaskCostSummary {
		const taskEntries = this.entries.filter((e) => e.taskId === taskId)

		const totalCostUsd = taskEntries.reduce((sum, e) => sum + e.costUsd, 0)
		const totalInputTokens = taskEntries.reduce((sum, e) => sum + e.inputTokens, 0)
		const totalOutputTokens = taskEntries.reduce((sum, e) => sum + e.outputTokens, 0)
		const totalThinkingTokens = taskEntries.reduce((sum, e) => sum + (e.thinkingTokens ?? 0), 0)
		const totalDurationMs = taskEntries.reduce((sum, e) => sum + (e.durationMs ?? 0), 0)
		const sessionCount = taskEntries.length
		const costPerSession = sessionCount > 0 ? totalCostUsd / sessionCount : 0

		return {
			taskId,
			totalCostUsd,
			totalInputTokens,
			totalOutputTokens,
			totalThinkingTokens,
			sessionCount,
			totalDurationMs,
			entries: taskEntries,
			costPerSession,
		}
	}

	/**
	 * Compare costs across providers/models for a given task.
	 */
	compare(taskId: string): ProviderCostComparison {
		const taskEntries = this.entries.filter((e) => e.taskId === taskId)

		// Group by providerId + modelId
		const groups = new Map<
			string,
			{
				providerId: string
				modelId: string
				totalCostUsd: number
				totalTokens: number
				sessionCount: number
			}
		>()

		for (const entry of taskEntries) {
			const key = `${entry.providerId}:${entry.modelId}`
			const existing = groups.get(key)
			if (existing) {
				existing.totalCostUsd += entry.costUsd
				existing.totalTokens += entry.inputTokens + entry.outputTokens
				existing.sessionCount += 1
			} else {
				groups.set(key, {
					providerId: entry.providerId,
					modelId: entry.modelId,
					totalCostUsd: entry.costUsd,
					totalTokens: entry.inputTokens + entry.outputTokens,
					sessionCount: 1,
				})
			}
		}

		return {
			taskId,
			byProvider: Array.from(groups.values()),
		}
	}

	/**
	 * Generate cost trend data points aggregated by date.
	 * @param periodDays Number of days to include (default 30)
	 */
	trend(periodDays = 30): CostTrend {
		// Group entries by date
		const byDate = new Map<
			string,
			{ tasks: Set<string>; totalCostUsd: number; tokenTotal: number }
		>()

		for (const entry of this.entries) {
			const date = entry.timestamp.slice(0, 10) // YYYY-MM-DD
			const existing = byDate.get(date)
			if (existing) {
				existing.tasks.add(entry.taskId)
				existing.totalCostUsd += entry.costUsd
				existing.tokenTotal += entry.inputTokens + entry.outputTokens
			} else {
				byDate.set(date, {
					tasks: new Set([entry.taskId]),
					totalCostUsd: entry.costUsd,
					tokenTotal: entry.inputTokens + entry.outputTokens,
				})
			}
		}

		// Sort by date and limit to periodDays
		const sortedDates = Array.from(byDate.keys()).sort()
		const limitedDates = sortedDates.slice(-periodDays)

		const dataPoints = limitedDates.map((date) => {
			const data = byDate.get(date)!
			const taskCount = data.tasks.size
			return {
				date,
				taskCount,
				totalCostUsd: data.totalCostUsd,
				avgCostPerTask: taskCount > 0 ? data.totalCostUsd / taskCount : 0,
				tokenTotal: data.tokenTotal,
			}
		})

		return { dataPoints, periodDays }
	}

	/**
	 * Serialize all entries for persistence.
	 */
	toJSON(): CostEntry[] {
		return [...this.entries]
	}

	/**
	 * Restore a CostTracker from persisted entries.
	 */
	static fromJSON(entries: CostEntry[]): CostTracker {
		const tracker = new CostTracker()
		tracker.entries = [...entries]
		return tracker
	}
}
