import type { Clock, IdSource, TraceEnvelope, TraceEventType } from "./types.ts";

interface ActiveToolSpan {
	spanId: string;
	parentSpanId?: string;
}

function extractToolCallIds(message: unknown): string[] {
	if (typeof message !== "object" || message === null) return [];
	const record = message as Record<string, unknown>;
	if (record.role !== "assistant" || !Array.isArray(record.content)) return [];

	const ids: string[] = [];
	for (const block of record.content) {
		if (typeof block !== "object" || block === null) continue;
		const content = block as Record<string, unknown>;
		if (content.type === "toolCall" && typeof content.id === "string") {
			ids.push(content.id);
		}
	}
	return ids;
}

export class TraceCorrelator {
	private readonly sessionId: string;
	private readonly runId: string;
	private readonly clock: Clock;
	private readonly ids: IdSource;
	private sequence = 0;
	private activeInteractionId: string | undefined;
	private currentTurnId: string | undefined;
	private activeLlmSpanId: string | undefined;
	private readonly toolParents = new Map<string, string>();
	private readonly activeTools = new Map<string, ActiveToolSpan>();

	constructor(sessionId: string, runId: string, clock: Clock, ids: IdSource) {
		this.sessionId = sessionId;
		this.runId = runId;
		this.clock = clock;
		this.ids = ids;
	}

	startRun(reason: string): TraceEnvelope {
		return this.createEvent("run.start", { reason });
	}

	endRun(reason: string): TraceEnvelope {
		return this.createEvent("run.end", { reason });
	}

	startInteraction(input: unknown): TraceEnvelope {
		this.activeInteractionId = this.ids.next();
		return this.createEvent("interaction.start", { input });
	}

	endInteraction(reason: string): TraceEnvelope {
		const event = this.createEvent("interaction.end", { reason });
		this.activeInteractionId = undefined;
		return event;
	}

	startTurn(turnIndex: number, timestamp: number): TraceEnvelope {
		this.currentTurnId = this.ids.next();
		this.activeLlmSpanId = undefined;
		this.toolParents.clear();
		this.activeTools.clear();
		return this.createEvent(
			"turn.start",
			{ turnIndex, timestamp },
			{ turnId: this.currentTurnId, spanId: this.currentTurnId },
		);
	}

	endTurn(turnIndex: number, message: unknown, toolResults: unknown): TraceEnvelope {
		const event = this.createEvent(
			"turn.end",
			{ turnIndex, message, toolResults },
			{ turnId: this.currentTurnId, spanId: this.currentTurnId },
		);
		this.currentTurnId = undefined;
		this.activeLlmSpanId = undefined;
		this.toolParents.clear();
		this.activeTools.clear();
		return event;
	}

	startLlm(payload: unknown): TraceEnvelope {
		this.activeLlmSpanId = this.ids.next();
		return this.createEvent("llm.request", payload, {
			turnId: this.currentTurnId,
			spanId: this.activeLlmSpanId,
			parentSpanId: this.currentTurnId,
		});
	}

	finishLlm(message: unknown): TraceEnvelope {
		const spanId = this.activeLlmSpanId;
		if (spanId !== undefined) {
			for (const toolCallId of extractToolCallIds(message)) {
				this.toolParents.set(toolCallId, spanId);
			}
		}
		const event = this.createEvent("llm.response", message, {
			turnId: this.currentTurnId,
			spanId,
			parentSpanId: this.currentTurnId,
		});
		this.activeLlmSpanId = undefined;
		return event;
	}

	startTool(toolCallId: string, toolName: string, args: unknown): TraceEnvelope {
		const spanId = this.ids.next();
		const parentSpanId = this.toolParents.get(toolCallId) ?? this.currentTurnId;
		this.activeTools.set(toolCallId, { spanId, parentSpanId });
		return this.createEvent("tool.request", { toolCallId, toolName, args }, {
			turnId: this.currentTurnId,
			spanId,
			parentSpanId,
		});
	}

	finishTool(toolCallId: string, toolName: string, result: unknown, isError: boolean): TraceEnvelope {
		const active = this.activeTools.get(toolCallId);
		const event = this.createEvent("tool.response", { toolCallId, toolName, result, isError }, {
			turnId: this.currentTurnId,
			spanId: active?.spanId,
			parentSpanId: active?.parentSpanId ?? this.currentTurnId,
		});
		this.activeTools.delete(toolCallId);
		return event;
	}

	warning(payload: unknown): TraceEnvelope {
		return this.createEvent("trace.warning", payload, { turnId: this.currentTurnId });
	}

	private createEvent(
		type: TraceEventType,
		payload: unknown,
		context: { interactionId?: string; turnId?: string; spanId?: string; parentSpanId?: string } = {},
	): TraceEnvelope {
		const event: TraceEnvelope = {
			schemaVersion: 1,
			seq: ++this.sequence,
			eventId: this.ids.next(),
			timestamp: this.clock.now().toISOString(),
			sessionId: this.sessionId,
			runId: this.runId,
			type,
			payload,
		};
		const interactionId = context.interactionId ?? this.activeInteractionId;
		if (interactionId !== undefined) event.interactionId = interactionId;
		if (context.turnId !== undefined) event.turnId = context.turnId;
		if (context.spanId !== undefined) event.spanId = context.spanId;
		if (context.parentSpanId !== undefined) event.parentSpanId = context.parentSpanId;
		return event;
	}
}