import { TraceCorrelator } from "./correlator.ts";
import { serializeTraceEvent } from "./serialize.ts";
import type { Clock, IdSource, TraceConfig, TraceEnvelope } from "./types.ts";

export interface TraceEventSink {
	write(event: TraceEnvelope, line: string): void;
	start?(): void;
	flush?(timeoutMs: number): Promise<boolean>;
	shutdown?(timeoutMs: number): Promise<void>;
	open?(): void;
}

export interface TraceRuntimeOptions {
	sessionId: string;
	runId: string;
	config: TraceConfig;
	clock: Clock;
	ids: IdSource;
	persistence?: TraceEventSink;
	viewer?: TraceEventSink;
	onWarning: (warning: string) => void;
}

export class TraceRuntime {
	private readonly correlator: TraceCorrelator;
	private readonly persistence: TraceEventSink | undefined;
	private readonly viewer: TraceEventSink | undefined;
	private readonly onWarning: (warning: string) => void;
	private readonly warnedSinks = new Set<string>();
	private enabled: boolean;
	private ended = false;
	private interactionActive = false;
	private ignoreNextUserMessage = false;
	private pendingTurn: { turnIndex: number; timestamp: number } | undefined;

	constructor(options: TraceRuntimeOptions) {
		this.correlator = new TraceCorrelator(options.sessionId, options.runId, options.clock, options.ids);
		this.persistence = options.config.persistence ? options.persistence : undefined;
		this.viewer = options.viewer;
		this.onWarning = options.onWarning;
		this.enabled = options.config.enabled;
	}

	start(reason: string): void {
		if (this.enabled) this.startViewer();
		this.capture(() => this.correlator.startRun(reason));
	}
	startAgentInput(prompt: string, images: unknown): void {
		this.startInteraction({ role: "user", content: prompt, images });
		this.ignoreNextUserMessage = this.interactionActive;
	}
	observeMessage(message: unknown): void {
		if (typeof message !== "object" || message === null || !("role" in message) || message.role !== "user") return;
		if (this.ignoreNextUserMessage) {
			this.ignoreNextUserMessage = false;
			return;
		}
		this.startInteraction(message);
	}
	endAgent(): void {
		this.flushPendingTurn();
		this.closeInteraction("agent-end");
	}
	startTurn(turnIndex: number, timestamp: number): void {
		this.flushPendingTurn();
		this.pendingTurn = { turnIndex, timestamp };
	}
	endTurn(turnIndex: number, message: unknown, toolResults: unknown): void {
		this.flushPendingTurn();
		this.capture(() => this.correlator.endTurn(turnIndex, message, toolResults));
	}
	startLlm(payload: unknown): void {
		this.flushPendingTurn();
		this.capture(() => this.correlator.startLlm(payload));
	}
	finishLlm(message: unknown): void { this.capture(() => this.correlator.finishLlm(message)); }
	startTool(toolCallId: string, toolName: string, args: unknown): void {
		this.flushPendingTurn();
		this.capture(() => this.correlator.startTool(toolCallId, toolName, args));
	}
	finishTool(toolCallId: string, toolName: string, result: unknown, isError: boolean): void {
		this.capture(() => this.correlator.finishTool(toolCallId, toolName, result, isError));
	}

	setEnabled(enabled: boolean): void {
		const wasEnabled = this.enabled;
		this.enabled = enabled;
		if (enabled && !wasEnabled) this.startViewer();
	}
	isEnabled(): boolean { return this.enabled; }
	open(): void {
		try { this.viewer?.open?.(); } catch (error) { this.warnOnce("viewer", error); }
	}

	async shutdown(reason: string, timeoutMs: number): Promise<void> {
		if (this.ended) return;
		this.flushPendingTurn();
		this.closeInteraction("session-shutdown", true);
		this.ended = true;
		this.capture(() => this.correlator.endRun(reason), true);
		await Promise.allSettled([
			this.persistence?.flush?.(timeoutMs),
			this.viewer?.shutdown?.(timeoutMs),
		]);
	}

	private startInteraction(input: unknown): void {
		this.closeInteraction("new-user-message");
		this.interactionActive = this.capture(() => this.correlator.startInteraction(input)) !== undefined;
	}

	private closeInteraction(reason: string, force = false): void {
		if (!this.interactionActive) return;
		this.capture(() => this.correlator.endInteraction(reason), force);
		this.interactionActive = false;
		this.ignoreNextUserMessage = false;
	}

	private flushPendingTurn(): void {
		const pending = this.pendingTurn;
		if (pending === undefined) return;
		this.pendingTurn = undefined;
		this.capture(() => this.correlator.startTurn(pending.turnIndex, pending.timestamp));
	}

	private capture(createEvent: () => TraceEnvelope, force = false): TraceEnvelope | undefined {
		if ((!this.enabled && !force) || this.ended && !force) return undefined;
		try {
			const event = createEvent();
			const line = serializeTraceEvent(event);
			this.writeSink("persistence", this.persistence, event, line);
			this.writeSink("viewer", this.viewer, event, line);
			return event;
		} catch (error) {
			this.warnOnce("runtime", error);
			return undefined;
		}
	}

	private writeSink(name: string, sink: TraceEventSink | undefined, event: TraceEnvelope, line: string): void {
		if (sink === undefined) return;
		try { sink.write(event, line); } catch (error) { this.warnOnce(name, error); }
	}

	private startViewer(): void {
		try { this.viewer?.start?.(); } catch (error) { this.warnOnce("viewer", error); }
	}

	private warnOnce(subsystem: string, error: unknown): void {
		if (this.warnedSinks.has(subsystem)) return;
		this.warnedSinks.add(subsystem);
		try {
			this.onWarning(`Trace ${subsystem} failed: ${error instanceof Error ? error.message : String(error)}`);
		} catch {
			// Notifications are observational.
		}
	}
}