import { GENAI } from '../constants/telemetry/index.js'
import { taskFailed } from '../tools/coordinator/outcome.js'
import type { AgentInput } from '../types/agent/base.js'
import type { AgentManagerContract } from '../types/agent/manager.js'
import type {
	CreateTaskOptions,
	SiblingFailurePolicy,
	TaskHandle,
	TaskScheduler,
} from '../types/agent/scheduler.js'
import type { AgentTaskContext } from '../types/agent/task.js'
import type { TaskId } from '../types/ids/index.js'
import { createUserMessage } from '../types/message/index.js'
import type { CancelCause } from '../types/session/cancel-cause.js'
import type { ChildSessionLifecycleEvent, SessionEventListener } from '../types/session/events.js'
import { toErrorMessage } from '../utils/error.js'
import { SCOPE_ATTRIBUTE } from '../utils/log/types.js'
import { type Logger, resolveLogger } from '../utils/logger.js'

/**
 * How many launched tasks a gateway remembers.
 *
 * High enough that no realistic single turn reaches it — a fan-out is eight,
 * a long supervisory turn is dozens — so the listing a supervisor reads at the
 * end of its turn is always complete. It exists for the host that reuses one
 * gateway across turns, where the alternative is a set and a map that grow for
 * the life of the process.
 */
const GATEWAY_TASK_LEDGER_CAP = 1_000

interface ObserverDelivery {
	busy: boolean
	readonly queue: Array<{
		readonly event: Parameters<SessionEventListener>[0]
		readonly observer: 'scheduler' | 'task'
	}>
}

export class LocalTaskScheduler implements TaskScheduler {
	private agentManager: AgentManagerContract
	private taskContext: AgentTaskContext
	private listener: SessionEventListener | undefined
	private trackedTaskIds: Set<TaskId> = new Set()

	private parentInput?: Pick<AgentInput, 'taskStore' | 'runtimeToolOverrides' | 'runtimeContext'>

	private completionListeners: Set<(handle: TaskHandle) => void> = new Set()

	/**
	 * The settled summary of each task, kept past the manager's eviction.
	 *
	 * Terminal tasks leave the manager 30 seconds after they finish, and
	 * `listTasks` rebuilt itself by looking every tracked id back up — so a
	 * task that finished a minute ago simply vanished from the tool whose
	 * whole job is the end-of-turn check. A supervisor could not tell an
	 * evicted task from one that never launched; both read as absence.
	 *
	 * Eviction is there to release the heavy state — messages, controllers,
	 * spawn records — not the fact that the task ran. This is a handful of
	 * fields per task, bounded by the number the gateway itself launched.
	 */
	private settledHandles: Map<TaskId, TaskHandle> = new Map()

	private siblingFailurePolicy: SiblingFailurePolicy = 'continue'
	/** See {@link onTaskProgress}. */
	private readonly progressListeners = new Set<(taskId: TaskId) => void>()
	private readonly childSessionListeners = new Set<(event: ChildSessionLifecycleEvent) => void>()
	/**
	 * One non-blocking delivery chain per observer.
	 *
	 * Child streaming never awaits these promises, but a listener still sees
	 * its own events in source order. Keeping the chains separate also means a
	 * slow exporter cannot hold up a task-specific screen (or another task's
	 * screen) that uses a different callback.
	 */
	private readonly observerDeliveries = new WeakMap<SessionEventListener, ObserverDelivery>()
	/** Raw, unresolved — kept as the caller handed it so each of the two log
	 * sites below resolves it independently via `resolveLogger`, rather than
	 * this constructor baking in ONE `.child()` binding both would then share
	 * (which would also change how many `component:` bindings this file has,
	 * a different ratchet than the one this task moves). */
	private readonly log?: Logger

	constructor(
		agentManager: AgentManagerContract,
		taskContext: AgentTaskContext,
		listener?: SessionEventListener,
		parentInput?: Pick<AgentInput, 'taskStore' | 'runtimeToolOverrides' | 'runtimeContext'>,
		options?: { siblingFailurePolicy?: SiblingFailurePolicy; log?: Logger },
	) {
		this.agentManager = agentManager
		this.taskContext = taskContext
		this.listener = listener
		this.parentInput = parentInput
		this.log = options?.log
		if (options?.siblingFailurePolicy) {
			this.siblingFailurePolicy = options.siblingFailurePolicy
		}
	}

	get budget() {
		return this.taskContext.budget
	}

	async createTask(options: CreateTaskOptions): Promise<TaskHandle> {
		// Filled once the spawn resolves. A box rather than a bare binding
		// because the assignment happens AFTER the `await` that the reader is
		// passed into — see the progress tee below for why it cannot simply
		// read the `task` const it is declared beside.
		const launched: { id?: TaskId } = {}

		const task = await this.agentManager.sendMessage(
			{
				agentId: options.agentId,
				beforeStart: options.beforeStart,
				...(options.planId ? { planId: options.planId } : {}),
				...(options.planStepId ? { planStepId: options.planStepId } : {}),
				// Display grouping travels with the spawn so the manager can put it
				// on `agent_pending`. Spread conditionally, like the plan edge
				// above: a host that groups nothing must not be made to look like
				// one that grouped everything under an empty label.
				...(options.workflow ? { workflow: options.workflow } : {}),
				...(options.phase ? { phase: options.phase } : {}),
				...(options.phaseDetail ? { phaseDetail: options.phaseDetail } : {}),
				...(options.phaseOrder !== undefined ? { phaseOrder: options.phaseOrder } : {}),
				input: {
					messages: [createUserMessage(options.prompt)],
					workingDirectory: options.workingDirectory,
					taskStore: this.parentInput?.taskStore,
					runtimeToolOverrides: this.parentInput?.runtimeToolOverrides,
					runtimeContext: options.runtimeContext ?? this.parentInput?.runtimeContext,
				},
				// Phase 6: spawn scope propagates from the gateway's task context.
				// The caller built it at SupervisorAgent boundary (§12.1).
				parentSessionId: this.taskContext.sessionId,
				tenantId: this.taskContext.tenantId,
				projectId: this.taskContext.projectId,
				parentActor: this.taskContext.parentActor,
				// The caller's overrides, plus the span the caller supplied so a
				// delegated session joins the trace it belongs to instead of
				// starting its own root.
				//
				// `options.configOverrides` used to be dropped here: this built
				// a fresh object from `parentSpan` and never looked at the
				// field, so a caller pinning a child to a cheaper model got the
				// agent's default and no sign anything had been ignored. The
				// dedicated `parentSpan` option is applied last because it is
				// the specific field for that job — a caller who sets both is
				// saying the same thing twice, and the named one is the answer.
				...(options.configOverrides || options.parentSpan
					? {
							configOverrides: {
								...options.configOverrides,
								...(options.parentSpan ? { parentSpan: options.parentSpan } : {}),
							},
						}
					: {}),
			},
			// Every sibling reserves from this same authority before provisioning.
			this.taskContext,
			// The host's listener still sees everything it always did; this
			// only tees off the fact that SOMETHING happened, which is what an
			// idle bound measures. The event itself is not forwarded — a
			// progress signal that carried the child's output would be a
			// second, undocumented way to read a worker's work.
			//
			// The id comes from `launched.id`, NOT from the `task` const
			// below. This callback is handed to the very `await` that assigns
			// `task`, so a child that emits anything before `sendMessage`
			// resolves reached it inside the temporal dead zone and threw
			// `Cannot access 'task' before initialization` — killing the launch
			// outright.
			//
			// It survived because a single sequential launch usually resolves
			// before the child says anything. A concurrent fan-out does not:
			// with four `create_task` calls from one turn — the shape this
			// tool's own description tells the model to use — the event loop
			// interleaves and three of the four died. Found by running one.
			(event) => {
				// A scheduler-wide observer and this task's observer are independent.
				// Either may throw or reject: observation is not authority over the
				// child, and one broken screen/export must neither stop the turn nor
				// suppress the other observer. Async listeners are deliberately not
				// awaited, so a slow renderer cannot backpressure model streaming.
				this.deliverEvent(this.listener, event, 'scheduler')
				if (options.onEvent !== this.listener) {
					this.deliverEvent(options.onEvent, event, 'task')
				}
				// The parent turn's own record of its children. Synchronous and
				// before anything else can settle, so `child_session_ended` is
				// queued ahead of the parent's next record.
				if (
					event.type === 'child_session_spawned' ||
					event.type === 'child_session_messaged' ||
					event.type === 'child_session_idled'
				) {
					for (const notify of this.childSessionListeners) {
						try {
							notify(event)
						} catch (err) {
							resolveLogger(this.log)
								.child({ [SCOPE_ATTRIBUTE]: 'scheduler/local' })
								.warn('Child session observer failed', {
									'exception.message': toErrorMessage(err),
								})
						}
					}
				}
				// No id yet means nothing is waiting on this task: the caller
				// does not hold the handle, so an idle bound cannot be running
				// against it. There is no progress to report to anyone.
				if (launched.id === undefined) return
				for (const notify of this.progressListeners) notify(launched.id)
			},
		)

		launched.id = task.taskId
		this.trackedTaskIds.add(task.taskId)
		this.forgetOldestBeyondCap()

		this.agentManager
			.waitForCompletion(task.taskId)
			.then(() => {
				const completed = this.agentManager.getInstance(task.taskId)
				if (completed) {
					const handle = toHandle(completed)
					// Snapshot now, while the manager still holds it — in 30
					// seconds eviction takes the record away.
					this.settledHandles.set(task.taskId, handle)
					this.applySiblingPolicy(handle)
					for (const cb of this.completionListeners) {
						cb(handle)
					}
				}
			})
			.catch((err) => {
				resolveLogger(this.log)
					.child({ [SCOPE_ATTRIBUTE]: 'scheduler/local' })
					.error('Task completion tracking failed', {
						'namzu.task.id': task.taskId,
						'exception.message': toErrorMessage(err),
					})
			})

		return toHandle(task)
	}

	private deliverEvent(
		listener: SessionEventListener | undefined,
		event: Parameters<SessionEventListener>[0],
		observer: 'scheduler' | 'task',
	): void {
		if (!listener) return

		// Each observer receives its own value graph. SessionEvent is readonly at the
		// type boundary, but a JavaScript consumer can still mutate an object it
		// was handed; that must not forge what the next independent observer sees.
		const snapshot = structuredClone(event)
		let delivery = this.observerDeliveries.get(listener)
		if (!delivery) {
			delivery = { busy: false, queue: [] }
			this.observerDeliveries.set(listener, delivery)
		}
		delivery.queue.push({ event: snapshot, observer })
		this.drainObserver(listener, delivery)
	}

	private drainObserver(listener: SessionEventListener, delivery: ObserverDelivery): void {
		if (delivery.busy) return
		while (delivery.queue.length > 0) {
			const next = delivery.queue.shift()
			if (!next) return
			try {
				const pending = listener(next.event)
				if (!pending) continue
				delivery.busy = true
				void pending.then(
					() => {
						delivery.busy = false
						this.drainObserver(listener, delivery)
					},
					(error) => {
						this.reportObserverFailure(next.event.type, next.observer, error)
						delivery.busy = false
						this.drainObserver(listener, delivery)
					},
				)
				return
			} catch (error) {
				this.reportObserverFailure(next.event.type, next.observer, error)
			}
		}
	}

	private reportObserverFailure(
		eventType: Parameters<SessionEventListener>[0]['type'],
		observer: 'scheduler' | 'task',
		error: unknown,
	): void {
		resolveLogger(this.log)
			.child({ [SCOPE_ATTRIBUTE]: 'scheduler/local' })
			.warn('Task event observer failed', {
				'namzu.event.type': eventType,
				'namzu.scheduler.observer': observer,
				'exception.message': toErrorMessage(error),
			})
	}

	/**
	 * Decide what a failed child means for the ones still running.
	 *
	 * The primitive to stop them already existed — every child holds an
	 * abort controller chained to the parent's, and `AgentManager.cancel`
	 * uses it — but nothing connected a failure to it. So a supervisor that
	 * fanned out five tasks and watched one die had no way to say the other
	 * four were now pointless: they ran to completion spending budget on
	 * work whose premise had gone.
	 *
	 * `'continue'` stays the default, and deliberately. Partial results are
	 * usually worth having, and a policy that tore down healthy siblings on
	 * any failure would make one flaky child able to waste four good ones.
	 * The point is that the choice is now expressible, not that the answer
	 * changed.
	 */
	private applySiblingPolicy(finished: TaskHandle): void {
		if (this.siblingFailurePolicy !== 'cancel-siblings') return
		if (!taskFailed(finished)) return

		const cancelled: TaskId[] = []
		for (const taskId of this.trackedTaskIds) {
			if (taskId === finished.taskId) continue
			const sibling = this.agentManager.getInstance(taskId)
			// `cancel` is already a no-op on a terminal task, but checking
			// here keeps the log honest about what was actually stopped.
			if (!sibling || sibling.state === 'completed' || sibling.state === 'failed') continue
			this.agentManager.cancel(taskId)
			cancelled.push(taskId)
		}

		if (cancelled.length > 0) {
			resolveLogger(this.log)
				.child({ [SCOPE_ATTRIBUTE]: 'scheduler/local' })
				.info('Cancelled siblings after a child failed', {
					'namzu.scheduler.failed': finished.taskId,
					[GENAI.AGENT_ID]: finished.agentId,
					'namzu.scheduler.cancelled': cancelled,
				})
		}
	}

	async waitForTask(taskId: TaskId): Promise<TaskHandle> {
		await this.agentManager.waitForCompletion(taskId)
		const task = this.agentManager.getInstance(taskId)
		if (!task) {
			throw new Error(`Task ${taskId} not found after completion`)
		}
		return toHandle(task)
	}

	async continueTask(taskId: TaskId, message: string): Promise<void> {
		await this.agentManager.continueTask(taskId, message)
	}

	cancelTask(taskId: TaskId, cause?: CancelCause): void {
		this.agentManager.cancel(taskId, cause)
	}

	getTask(taskId: TaskId): TaskHandle | undefined {
		const task = this.agentManager.getInstance(taskId)
		return task ? toHandle(task) : undefined
	}

	/**
	 * Snapshots a task's terminal state so it survives eviction.
	 *
	 * Called when the task settles, while the manager still holds it.
	 */
	rememberSettled(taskId: TaskId): void {
		const task = this.agentManager.getInstance(taskId)
		if (task) this.settledHandles.set(taskId, toHandle(task))
	}

	/**
	 * Drop the oldest tasks once the ledger passes {@link GATEWAY_TASK_LEDGER_CAP}.
	 *
	 * A gateway constructed per turn is bounded by that turn and this never
	 * fires. But `SupervisorAgentConfig.gateway` lets a host supply its own,
	 * and a long-lived host reusing one accumulates an id and a settled handle
	 * per task it ever launched, for the life of the process — the doc above
	 * says "bounded by the number the gateway itself launched", which is true
	 * and is not a bound when the gateway outlives the turn.
	 *
	 * Both collections are evicted **together and in insertion order**. Losing
	 * a tracked id while keeping its handle, or the reverse, would make a task
	 * that ran read as one that never launched — which is the exact defect the
	 * settled-handle map was added to fix, reintroduced by its own cleanup.
	 */
	private forgetOldestBeyondCap(): void {
		while (this.trackedTaskIds.size > GATEWAY_TASK_LEDGER_CAP) {
			const oldest = this.trackedTaskIds.values().next().value
			if (oldest === undefined) return
			this.trackedTaskIds.delete(oldest)
			this.settledHandles.delete(oldest)
		}
	}

	listTasks(): TaskHandle[] {
		const handles: TaskHandle[] = []
		for (const taskId of this.trackedTaskIds) {
			// The live task wins: a remembered snapshot must never shadow
			// state that is still being updated.
			const task = this.agentManager.getInstance(taskId)
			if (task) {
				handles.push(toHandle(task))
				continue
			}
			const settled = this.settledHandles.get(taskId)
			if (settled) handles.push(settled)
		}
		return handles
	}

	/**
	 * Every event a child emits, reduced to "this one is still alive".
	 *
	 * Deliberately just the id. A caller that wanted the event itself has
	 * the turn listener; what an idle clock needs is the fact, and passing
	 * the payload here would make this a second way to read a worker's
	 * output — one nobody documented and nothing frames as untrusted.
	 */
	onTaskProgress(callback: (taskId: TaskId) => void): () => void {
		this.progressListeners.add(callback)
		return () => {
			this.progressListeners.delete(callback)
		}
	}

	onTaskCompleted(callback: (handle: TaskHandle) => void): () => void {
		this.completionListeners.add(callback)
		return () => {
			this.completionListeners.delete(callback)
		}
	}

	onChildSessionEvent(callback: (event: ChildSessionLifecycleEvent) => void): () => void {
		this.childSessionListeners.add(callback)
		return () => {
			this.childSessionListeners.delete(callback)
		}
	}
}

/**
 * Did this child fail?
 *
 * Two answers have to agree. `state` is `'failed'` only when the spawn
 * machinery itself threw; a child whose agent RAN and returned
 * `status: 'failed'` lands in `markCompleted` regardless, carrying the
 * failure in its result rather than its state. Reading only the state
 * would therefore miss the ordinary case — an agent that tried and could
 * not — and catch only the exceptional one.
 */
function toHandle(task: import('../types/agent/task.js').AgentTask): TaskHandle {
	return {
		taskId: task.taskId,
		agentId: task.agentId,
		state: task.state,
		result: task.result,
		createdAt: task.createdAt,
		completedAt: task.completedAt,
	}
}