// Live elapsed-time labels for folded thinking rows.
//
// Pi renders one static label ("thoughts") per hidden thinking run. This
// module upgrades that label while the fold itself stays pi's: while the model
// is thinking, the row shows a spinner and the running elapsed time; when the
// run ends, the row settles to the actual duration ("thoughts · 8.4s").
//
// Label targeting: pi's ctx.ui.setHiddenThinkingLabel is global — it rewrites
// every assistant message component. A final duration belongs to one message,
// so it is written per component instance instead: the streaming component
// stays mounted in the chat container after message_end (pi only clears its
// reference), so a duck-typed walk — the same pattern thinking-fold.ts uses
// for actionHandlers — can address exactly the row that just finished
// thinking. When the walk finds nothing (non-TUI mode, pi refactor), the
// global label API is the degraded fallback: the animation still runs, and
// the final duration lingers briefly before the base label returns.

import type { AssistantMessageEvent } from "@earendil-works/pi-ai";
import type { TUI } from "@earendil-works/pi-tui";
import type { IconGlyphs } from "./icons.ts";
import { tuiT } from "./tui-i18n.ts";

/**
 * Duck-typed AssistantMessageComponent: the two public setters pi gives it.
 * No other component in pi's chat tree exposes this pair.
 */
export interface ThinkingLabelTarget {
	setHiddenThinkingLabel(label: string): void;
	setHideThinkingBlock(hide: boolean): void;
}

function isThinkingLabelTarget(node: object): node is ThinkingLabelTarget {
	const candidate = node as Partial<ThinkingLabelTarget>;
	return (
		typeof candidate.setHiddenThinkingLabel === "function"
		&& typeof candidate.setHideThinkingBlock === "function"
	);
}

/**
 * Collect label targets in document order. The newest assistant message is
 * always the last target: pi appends the streaming component to the chat
 * container on message_start, and later siblings (tool executions, the next
 * message) are not targets.
 */
export function findThinkingLabelTargets(tui: TUI): ThinkingLabelTarget[] {
	const targets: ThinkingLabelTarget[] = [];
	const seen = new Set<unknown>();
	const walk = (node: unknown): void => {
		if (typeof node !== "object" || node === null || seen.has(node)) return;
		seen.add(node);
		if (isThinkingLabelTarget(node)) targets.push(node);
		const children = (node as { children?: unknown }).children;
		if (Array.isArray(children)) {
			for (const child of children) walk(child);
		}
	};
	walk(tui);
	return targets;
}

export function formatThinkingDuration(ms: number): string {
	if (!Number.isFinite(ms) || ms <= 0) return "0.0s";
	if (ms < 10_000) return `${(ms / 1000).toFixed(1)}s`;
	const totalSeconds = Math.round(ms / 1000);
	if (totalSeconds < 60) return `${totalSeconds}s`;
	const minutes = Math.floor(totalSeconds / 60);
	return `${minutes}m${String(totalSeconds % 60).padStart(2, "0")}s`;
}

export interface ThinkingTimerOptions {
	getTui: () => TUI | undefined;
	/** Instance label writes do not schedule a paint; the owner must. */
	requestRender: () => void;
	/** Label folded rows currently show; undefined means pi's own default. */
	getBaseLabel: () => string | undefined;
	getGlyphs: () => IconGlyphs;
	/** Effective hideThinkingBlock — labels only render while thinking is folded. */
	isThinkingHidden: () => boolean;
	isEnabled: () => boolean;
	/** Static mode: keep a stable label (no live ticker); duration settles at end. */
	isStatic?: () => boolean;
	/** Global label fallback for when no component instance is reachable. */
	setGlobalLabel: (label: string | undefined) => void;
	now?: () => number;
	/** How long the fallback final label lingers before the base returns. */
	settleMs?: number;
}

const DEFAULT_SETTLE_MS = 5_000;
export const THINKING_TICK_PERIOD_MS = 1_000;

export class ThinkingFoldTimer {
	readonly #options: ThinkingTimerOptions;
	#startedAt: number | undefined;
	#ticker: ReturnType<typeof setInterval> | undefined;
	#settleTimer: ReturnType<typeof setTimeout> | undefined;
	#target: ThinkingLabelTarget | undefined;
	#lastElapsedSecond: number | undefined;
	// Invalidates pending settle restores across stop/reset boundaries so a
	// late timer can never rewrite a label a newer run already owns.
	#generation = 0;

	constructor(options: ThinkingTimerOptions) {
		this.#options = options;
	}

	onAssistantMessageStart(): void {
		this.#rebuildTarget();
	}

	onAssistantMessageEvent(event: AssistantMessageEvent | undefined): void {
		if (!event) return;
		if (!this.#options.isEnabled()) {
			this.stop();
			return;
		}
		if (event.type === "thinking_start") this.#begin();
		else if (event.type === "thinking_end") this.#finish();
	}

	/**
	 * Settle an interrupted run: message_end fires for aborted and failed
	 * messages too (agent_end only removes the row afterwards), so the
	 * duration written here is still accurate and lands on a mounted row.
	 */
	onAssistantMessageEnd(): void {
		this.#rebuildTarget();
		if (this.#startedAt !== undefined) this.#finish();
		this.#target = undefined;
	}

	/** One animation frame. The internal interval drives it; tests call it directly. */
	tick(): void {
		if (this.#startedAt === undefined) return;
		// Static mode keeps a stable label; the live elapsed only ticks when not static.
		if (this.#options.isStatic?.()) return;
		const now = this.#now();
		const elapsed = now - this.#startedAt;
		const elapsedSecond = Math.max(0, Math.floor(elapsed / 1000));
		if (elapsedSecond === this.#lastElapsedSecond) return;
		this.#lastElapsedSecond = elapsedSecond;
		// No spinner prefix: the thinking row shows a live elapsed without any
		// animated indicator (the host working line is the only spinner left).
		this.#paint(`${tuiT("thinking.thinking")} ${formatThinkingDuration(elapsed)}`);
	}

	/** Re-evaluate the label policy after a static-mode toggle mid-run. */
	syncMode(): void {
		if (this.#startedAt === undefined) return;
		if (this.#options.isStatic?.()) {
			this.#clearTimers();
			this.#lastElapsedSecond = undefined;
			this.#paint(this.#options.getBaseLabel() ?? tuiT("thinking.thinking"));
			return;
		}
		if (!this.#ticker) {
			this.#ticker = setInterval(() => this.tick(), THINKING_TICK_PERIOD_MS);
			this.#ticker.unref?.();
		}
		this.tick();
	}

	/** The agent stopped: drop the run without settling a duration. */
	stop(): void {
		this.#generation++;
		this.#clearTimers();
		this.#startedAt = undefined;
		this.#target = undefined;
		this.#lastElapsedSecond = undefined;
	}

	/** Session boundary: forget the run; pi rebuilds the chat from entries. */
	reset(): void {
		this.#generation++;
		this.#clearTimers();
		this.#startedAt = undefined;
		this.#target = undefined;
		this.#lastElapsedSecond = undefined;
	}

	#begin(): void {
		// The gate is per-run: when thinking is unfolded the label row is not
		// rendered, so there is nothing to animate and no duration to keep.
		if (!this.#options.isThinkingHidden()) return;
		this.#startedAt = this.#now();
		this.#lastElapsedSecond = undefined;
		if (this.#options.isStatic?.()) {
			// Static mode: a stable label, no live ticker; #finish settles the duration.
			this.#paint(this.#options.getBaseLabel() ?? tuiT("thinking.thinking"));
			return;
		}
		if (!this.#ticker) {
			this.#ticker = setInterval(() => this.tick(), THINKING_TICK_PERIOD_MS);
			this.#ticker.unref?.();
		}
		// Paint immediately so the row flips to the live label with the first
		// frame instead of waiting one second.
		this.tick();
	}

	#finish(): void {
		const startedAt = this.#startedAt;
		if (startedAt === undefined) return;
		this.#clearTimers();
		this.#startedAt = undefined;
		this.#lastElapsedSecond = undefined;
		const duration = formatThinkingDuration(this.#now() - startedAt);
		const base = this.#options.getBaseLabel() ?? tuiT("thinking.thoughts");
		this.#paint(`${base}${this.#options.getGlyphs().separator}${duration}`, true);
	}

	#paint(label: string, settle = false): void {
		const target = this.#target;
		if (target) {
			try {
				target.setHiddenThinkingLabel(label);
				this.#options.requestRender();
			} catch {
				// component may be mid-teardown; the global path recovers
				this.#target = undefined;
				this.#paintGlobal(label, settle);
			}
			return;
		}
		this.#paintGlobal(label, settle);
	}

	#rebuildTarget(): void {
		const tui = this.#options.getTui();
		this.#target = tui ? findThinkingLabelTargets(tui).at(-1) : undefined;
	}

	#paintGlobal(label: string, settle: boolean): void {
		try {
			this.#options.setGlobalLabel(label);
		} catch {
			// no UI surface at all — nothing to paint
		}
		if (!settle) return;
		// The global label is shared by every folded row, so the final
		// duration cannot stay: give it a moment to be read, then hand the
		// label back to whatever the owner currently wants.
		this.#generation++;
		const generation = this.#generation;
		if (this.#settleTimer) clearTimeout(this.#settleTimer);
		this.#settleTimer = setTimeout(() => {
			this.#settleTimer = undefined;
			if (generation !== this.#generation) return;
			try {
				this.#options.setGlobalLabel(this.#options.getBaseLabel());
			} catch {
				// best-effort restore
			}
		}, this.#options.settleMs ?? DEFAULT_SETTLE_MS);
		this.#settleTimer.unref?.();
	}

	#clearTimers(): void {
		if (this.#ticker) {
			clearInterval(this.#ticker);
			this.#ticker = undefined;
		}
		if (this.#settleTimer) {
			clearTimeout(this.#settleTimer);
			this.#settleTimer = undefined;
		}
	}

	#now(): number {
		return this.#options.now?.() ?? Date.now();
	}
}