import type { ToolChoice } from "@oh-my-pi/pi-ai";

// ── Callback types ──────────────────────────────────────────────────────────

export interface ResolveInfo {
	/** The ToolChoice that was served to the LLM. */
	choice: ToolChoice;
}

export interface RejectInfo {
	/** The ToolChoice that was yielded but never (or unsuccessfully) served. */
	choice: ToolChoice;
	reason: "aborted" | "error" | "cleared" | "removed" | "unavailable" | "not_invoked";
}

/** "requeue" replays the lost yield next turn; "drop" (or void/undefined) discards it. */
export type RejectOutcome = "requeue" | "drop";

export interface DirectiveCallbacks {
	/** Fires when the yield completed; onInvoked directives require the requested tool to run first. */
	onResolved?: (info: ResolveInfo) => void;
	/**
	 * Fires when the yield is being discarded. Return "requeue" to replay the
	 * same value at the head of the queue for the next turn. Default: "drop".
	 */
	onRejected?: (info: RejectInfo) => RejectOutcome | undefined;
	/**
	 * Handler invoked when the model actually calls the forced tool. The queue
	 * directive carries the real execution logic; the tool's own execute() is
	 * bypassed. Returns the tool result directly.
	 */
	onInvoked?: (input: unknown) => Promise<unknown> | unknown;
}

// ── Directive ───────────────────────────────────────────────────────────────

export interface ToolChoiceDirective {
	generator: Iterator<ToolChoice>;
	/** Stable label for targeted removal and debugging (e.g. "user-force"). */
	label: string;
	callbacks: DirectiveCallbacks;
}

export interface PushOptions {
	/** Prepend to head instead of appending to tail. Default: false. */
	now?: boolean;
	label?: string;
	/** Lifecycle callbacks for this directive. */
	onResolved?: DirectiveCallbacks["onResolved"];
	onRejected?: DirectiveCallbacks["onRejected"];
	onInvoked?: DirectiveCallbacks["onInvoked"];
}

// ── Generators ──────────────────────────────────────────────────────────────

export function* onceGen(choice: ToolChoice): Generator<ToolChoice, void, unknown> {
	yield choice;
}

// ── In-flight state ─────────────────────────────────────────────────────────

interface InFlight {
	directive: ToolChoiceDirective;
	yielded: ToolChoice;
	invoked: boolean;
}

/**
 * A non-forcing pending preview invoker. Registered by `queueResolveHandler`
 * (resolve previews) so the `resolve` tool can dispatch to a staged action
 * WITHOUT this queue forcing `tool_choice`. The agent-loop's
 * SoftToolRequirement lifecycle (remind-then-escalate) owns any forcing.
 */
interface PendingInvoker {
	/** Unique id for this staged preview; never reused (never clobbered by label). */
	id: string;
	/** Source tool that staged the preview (e.g. "ast_edit"), for the reminder. */
	sourceToolName: string;
	onInvoked: (input: unknown) => Promise<unknown> | unknown;
}

// ── Queue ───────────────────────────────────────────────────────────────────

export class ToolChoiceQueue {
	#queue: ToolChoiceDirective[] = [];
	#inFlight: InFlight | undefined;
	/**
	 * Label of the directive whose last yield was resolved this turn.
	 * Consumers (e.g. todo reminder suppression) read via consumeLastServedLabel().
	 */
	#lastResolvedLabel: string | undefined;
	/**
	 * Non-forcing pending preview invokers, stacked by UNIQUE id. The `resolve`
	 * tool dispatches to the head; the agent-loop's soft-tool-requirement
	 * lifecycle drives resolution without this queue forcing `tool_choice`.
	 */
	#pendingInvokers: PendingInvoker[] = [];

	// ── Push ──────────────────────────────────────────────────────────────

	pushOnce(choice: ToolChoice, options?: PushOptions): void {
		this.push(onceGen(choice), options);
	}

	pushSequence(choices: ToolChoice[], options?: PushOptions): void {
		this.push(choices, options);
	}

	push(generator: Iterable<ToolChoice>, options?: PushOptions): void {
		const directive: ToolChoiceDirective = {
			generator: generator[Symbol.iterator](),
			label: options?.label ?? "anonymous",
			callbacks: {
				onResolved: options?.onResolved,
				onRejected: options?.onRejected,
				onInvoked: options?.onInvoked,
			},
		};
		if (options?.now) {
			this.#queue.unshift(directive);
		} else {
			this.#queue.push(directive);
		}
	}

	// ── Consume ───────────────────────────────────────────────────────────

	/**
	 * Advance the head directive and return its next yield. Records the value
	 * as in-flight until resolve() or reject() is called.
	 */
	nextToolChoice(): ToolChoice | undefined {
		while (this.#queue.length > 0) {
			const head = this.#queue[0]!;
			const result = head.generator.next();
			if (result.done) {
				this.#queue.shift();
				continue;
			}
			this.#inFlight = { directive: head, yielded: result.value, invoked: false };
			return result.value;
		}
		return undefined;
	}

	// ── Lifecycle ─────────────────────────────────────────────────────────

	/**
	 * The in-flight yield completed normally. Directives with onInvoked are only
	 * consumed after their requested tool ran; a normal text turn or a different
	 * tool call requeues/rejects the directive instead.
	 */
	resolve(): void {
		const inFlight = this.#inFlight;
		if (!inFlight) return;
		if (inFlight.directive.callbacks.onInvoked && !inFlight.invoked) {
			this.reject("not_invoked");
			return;
		}
		this.#inFlight = undefined;

		this.#lastResolvedLabel = inFlight.directive.label;
		inFlight.directive.callbacks.onResolved?.({ choice: inFlight.yielded });
	}

	/**
	 * The in-flight yield was not served, or the turn aborted/errored.
	 * Fires onRejected to let the caller decide: "requeue" replays the exact
	 * lost value at the head of the queue; anything else drops it.
	 */
	reject(reason: RejectInfo["reason"]): void {
		const inFlight = this.#inFlight;
		this.#inFlight = undefined;
		if (!inFlight) return;

		const outcome = inFlight.directive.callbacks.onRejected?.({
			choice: inFlight.yielded,
			reason,
		});

		if (outcome === "requeue") {
			// Re-queue only the lost yield, not the rest of the sequence. Carry forward
			// callbacks so the replayed yield still executes and finalizes correctly,
			// and can requeue itself again if the next turn also aborts or skips it.
			this.#queue.unshift({
				generator: onceGen(inFlight.yielded),
				label: `${inFlight.directive.label}-requeued`,
				callbacks: {
					onResolved: inFlight.directive.callbacks.onResolved,
					onInvoked: inFlight.directive.callbacks.onInvoked,
					onRejected: inFlight.directive.callbacks.onRejected,
				},
			});
		}
	}

	/** True if there is an in-flight yield that hasn't been resolved or rejected. */
	get hasInFlight(): boolean {
		return this.#inFlight !== undefined;
	}

	/** Return the in-flight directive's onInvoked handler and mark it when called. */
	peekInFlightInvoker(): ((input: unknown) => Promise<unknown> | unknown) | undefined {
		const inFlight = this.#inFlight;
		const onInvoked = inFlight?.directive.callbacks.onInvoked;
		if (!inFlight || !onInvoked) return undefined;
		return (input: unknown): Promise<unknown> | unknown => {
			inFlight.invoked = true;
			return onInvoked(input);
		};
	}

	// ── Non-forcing pending invokers ──────────────────────────────────────
	// Preview producers (queueResolveHandler) register here so `resolve` can
	// dispatch to a staged action WITHOUT a forced tool_choice (no messages-cache
	// bust). Stacked by UNIQUE id: a re-register replaces only the same id, so
	// concurrent/sequential previews each survive and resolve independently.

	/** Register (or replace by exact id) a non-forcing pending preview invoker. */
	registerPendingInvoker(
		id: string,
		sourceToolName: string,
		onInvoked: (input: unknown) => Promise<unknown> | unknown,
	): void {
		this.removePendingInvoker(id);
		this.#pendingInvokers.push({ id, sourceToolName, onInvoked });
	}

	/** Drop the pending invoker with this id (e.g. after it resolves). */
	removePendingInvoker(id: string): void {
		this.#pendingInvokers = this.#pendingInvokers.filter(p => p.id !== id);
	}

	/** Drop every pending preview invoker without touching hard tool-choice directives. */
	clearPendingInvokers(): void {
		if (this.#pendingInvokers.length === 0) return;
		this.#pendingInvokers = [];
	}

	/** True when at least one non-forcing pending preview is registered. */
	get hasPendingInvoker(): boolean {
		return this.#pendingInvokers.length > 0;
	}

	/** The head (most-recently registered) pending invoker's handler, for resolve dispatch. */
	peekPendingInvoker(): ((input: unknown) => Promise<unknown> | unknown) | undefined {
		return this.#pendingInvokers.at(-1)?.onInvoked;
	}

	/** The head pending preview's stable id + source tool, for building the agent-level
	 *  SoftToolRequirement (the id drives reminder re-injection when the head changes). */
	peekPendingHead(): { id: string; sourceToolName: string } | undefined {
		const head = this.#pendingInvokers.at(-1);
		return head ? { id: head.id, sourceToolName: head.sourceToolName } : undefined;
	}

	// ── Cleanup ───────────────────────────────────────────────────────────

	/** Remove all directives with the given label. Rejects in-flight if it matches. */
	removeByLabel(label: string): void {
		if (this.#inFlight?.directive.label === label) {
			this.reject("removed");
		}
		this.#queue = this.#queue.filter(d => d.label !== label);
	}

	/** Empty the queue and reject any in-flight yield. */
	clear(): void {
		if (this.#inFlight) {
			this.reject("cleared");
		}
		this.#queue = [];
		this.#pendingInvokers = [];
		this.#lastResolvedLabel = undefined;
	}

	// ── Observation ───────────────────────────────────────────────────────

	/** Return the label of the most recently resolved directive, then clear it. */
	consumeLastServedLabel(): string | undefined {
		const label = this.#lastResolvedLabel;
		this.#lastResolvedLabel = undefined;
		return label;
	}

	/** For tests/debug: labels of currently queued directives in order. */
	inspect(): readonly string[] {
		return this.#queue.map(d => d.label);
	}
}