import {
	Elysia,
	type MapResponse,
	type Context,
	type TraceEvent,
	type TraceProcess
} from 'elysia'

type MaybePromise<T> = T | Promise<T>

class TimeLogger {
	private startTimes: Map<string, number>
	private durations: Map<string, number>
	private totalStartTime: number | null

	constructor() {
		this.startTimes = new Map()
		this.durations = new Map()
		this.totalStartTime = null
	}

	startTime(name: string) {
		this.startTimes.set(name, performance.now())
	}

	endTime(name: string) {
		const startTime = this.startTimes.get(name)
		if (startTime !== undefined) {
			const duration = performance.now() - startTime
			this.durations.set(name, duration)
			this.startTimes.delete(name)
		}
	}

	getCompletedDurations() {
		return Array.from(this.durations.entries()).map(([name, duration]) => ({
			name,
			duration
		}))
	}

	startTotal(): void {
		this.totalStartTime = performance.now()
	}

	endTotal(): number | null {
		if (this.totalStartTime === null) return null
		return performance.now() - this.totalStartTime
	}
}

export interface ServerTimingOptions {
	/**
	 * Should Elysia report data back to client via 'Server-Sent-Event'
	 */
	report?: boolean
	/**
	 * Allow Server Timing to log specified life-cycle events
	 */
	trace?: {
		/**
		 * Capture duration from request
		 *
		 * @default true
		 */
		request?: boolean
		/**
		 * Capture duration from parse
		 *
		 * @default true
		 */
		parse?: boolean
		/**
		 * Capture duration from transform
		 *
		 * @default true
		 */
		transform?: boolean
		/**
		 * Capture duration from beforeHandle
		 *
		 * @default true
		 */
		beforeHandle?: boolean
		/**
		 * Capture duration from handle
		 *
		 * @default true
		 */
		handle?: boolean
		/**
		 * Capture duration from afterHandle
		 *
		 * @default true
		 */
		afterHandle?: boolean
		/**
		 * Capture duration from mapResponse
		 *
		 * @default true
		 */
		error?: boolean
		/**
		 * Capture duration from mapResponse
		 *
		 * @default true
		 */
		mapResponse?: boolean
		/**
		 * Capture total duration from start to finish
		 *
		 * @default true
		 */
		total?: boolean
	}
	/**
	 * Determine whether or not Server Timing should be enabled
	 *
	 * @default NODE_ENV !== 'production'
	 */
	enabled?: boolean
	/**
	 * A condition whether server timing should be log
	 *
	 * @default undefined
	 */
	allow?:
		| MaybePromise<boolean>
		| ((context: Omit<Context, 'path'>) => MaybePromise<boolean>)
	/**
	 * A custom mapResponse provided by user
	 *
	 * @default undefined
	 */
	mapResponse?: MapResponse
}

const getLabel = (
	event: TraceEvent,
	listener: (
		callback: (process: TraceProcess<'begin', true>) => unknown
	) => unknown,
	write: (value: string) => void
) => {
	listener(async ({ onStop, onEvent, total }) => {
		let label = ''

		if (total === 0) return

		onEvent(({ name, index, onStop }) => {
			onStop(({ elapsed }) => {
				label += `${event}.${index}.${name || 'anon'};dur=${elapsed},`
			})
		})

		onStop(({ elapsed }) => {
			label += `${event};dur=${elapsed},`

			write(label)
		})
	})
}

export const serverTiming = ({
	allow,
	enabled = process.env.NODE_ENV !== 'production',
	trace: {
		request: traceRequest = true,
		parse: traceParse = true,
		transform: traceTransform = true,
		beforeHandle: traceBeforeHandle = true,
		handle: traceHandle = true,
		afterHandle: traceAfterHandle = true,
		error: traceError = true,
		mapResponse: traceMapResponse = true,
		total: traceTotal = true
	} = {},
	mapResponse
}: ServerTimingOptions = {}) => {
	const app = new Elysia().decorate('timeLog', new TimeLogger()).trace(
		{ as: 'global' },
		async ({
			onRequest,
			onParse,
			onTransform,
			onBeforeHandle,
			onHandle,
			onAfterHandle,
			onMapResponse,
			onError,
			set,
			context,
			response,
			context: {
				request: { method }
			}
		}) => {
			
			if (!enabled) return
			let label = ''

			const write = (nextValue: string) => {
				label += nextValue
			}

			let start: number

			onRequest(({ begin }) => {
				context.timeLog.startTotal()
				start = begin
			})

			if (traceRequest) getLabel('request', onRequest, write)
			if (traceParse) getLabel('parse', onParse, write)
			if (traceTransform) getLabel('transform', onTransform, write)
			if (traceBeforeHandle)
				getLabel('beforeHandle', onBeforeHandle, write)
			if (traceAfterHandle) getLabel('afterHandle', onAfterHandle, write)
			if (traceError) getLabel('error', onError, write)
			if (traceMapResponse) getLabel('mapResponse', onMapResponse, write)

			if (traceHandle)
				onHandle(({ name, onStop }) => {
					onStop(({ elapsed }) => {
						label += `handle.${name};dur=${elapsed},`
					})
				})

			onMapResponse(({ onStop }) => {
				onStop(async ({ end }) => {
					const completedDurations =
						context.timeLog.getCompletedDurations()
					if (completedDurations.length > 0) {
						label +=
							completedDurations
								.map(
									({ name, duration }) =>
										`${name};dur=${duration}`
								)
								.join(', ') + ','
					}
					const elapsed = context.timeLog.endTotal();

					let allowed = allow
					if (allowed instanceof Promise) allowed = await allowed

					if (traceTotal) label += `total;dur=${elapsed}`
					else label = label.slice(0, -1)

					// ? Must wait until request is reported
					switch (typeof allowed) {
						case 'boolean':
							if (allowed === false)
								delete set.headers['Server-Timing']

							set.headers['Server-Timing'] = label

							break

						case 'function':
							if ((await allowed(context)) === false)
								delete set.headers['Server-Timing']

							set.headers['Server-Timing'] = label

							break

						default:
							set.headers['Server-Timing'] = label
					}
				})
			})
		}
	)
	return app
}

export default serverTiming