import { createLogger } from '@wispplace/observability'
import type { RetainedUploadLease } from './request-body-admission'
import { failUploadJob } from './upload-jobs'

const logger = createLogger('main-app')
const UPLOAD_SHUTDOWN_MESSAGE = 'Upload failed'

type UploadWork = (signal: AbortSignal) => Promise<void>

interface TrackedUpload {
	jobId: string
	controller: AbortController
	bodyLease?: RetainedUploadLease
	promise: Promise<void>
}

export class PublicUploadLifecycle {
	private accepting = true
	private readonly tasks = new Set<TrackedUpload>()
	private stopPromise: Promise<void> | undefined

	isAccepting(): boolean {
		return this.accepting
	}

	track(jobId: string, work: UploadWork, bodyLease?: RetainedUploadLease): boolean {
		if (!this.accepting) return false
		const controller = new AbortController()
		let tracked: TrackedUpload
		const promise = Promise.resolve()
			.then(async () => await work(controller.signal))
			.catch(() => {
				logger.error('Public upload background task failed', { errorKind: 'upload_background_failed' })
				failUploadJob(jobId, UPLOAD_SHUTDOWN_MESSAGE)
			})
			.finally(() => {
				bodyLease?.release()
				this.tasks.delete(tracked)
			})
		tracked = { jobId, controller, bodyLease, promise }
		this.tasks.add(tracked)
		return true
	}

	async stopAndDrain(graceMs: number): Promise<void> {
		this.stopPromise ??= this.stopAndDrainOnce(graceMs)
		return await this.stopPromise
	}

	stats(): { accepting: boolean; tasks: number; retainedBodies: number } {
		return {
			accepting: this.accepting,
			tasks: this.tasks.size,
			retainedBodies: Array.from(this.tasks).filter((task) => task.bodyLease).length,
		}
	}

	private async stopAndDrainOnce(graceMs: number): Promise<void> {
		this.accepting = false
		const pending = Promise.allSettled(Array.from(this.tasks, (task) => task.promise))
		if (this.tasks.size === 0 || (await settlesWithin(pending, graceMs))) return

		// Abort cooperatively before marking terminal. Pipeline checkpoints prevent
		// a timed-out task from committing a root manifest after shutdown begins.
		for (const task of this.tasks) {
			task.controller.abort()
			failUploadJob(task.jobId, UPLOAD_SHUTDOWN_MESSAGE)
			// The server is no longer accepting work. Release the budget even if a
			// non-cooperative lock/PDS call needs process shutdown to interrupt it.
			task.bodyLease?.release()
			task.bodyLease = undefined
		}
	}
}

function settlesWithin(pending: Promise<unknown>, graceMs: number): Promise<boolean> {
	return new Promise((resolve) => {
		const timer = setTimeout(() => resolve(false), Math.max(0, graceMs))
		timer.unref?.()
		void pending.then(() => {
			clearTimeout(timer)
			resolve(true)
		})
	})
}

export const publicUploadLifecycle = new PublicUploadLifecycle()

export function getPublicUploadLifecycleStats(): { accepting: boolean; tasks: number; retainedBodies: number } {
	return publicUploadLifecycle.stats()
}

export function stopAndDrainPublicUploads(graceMs: number): Promise<void> {
	return publicUploadLifecycle.stopAndDrain(graceMs)
}