import type { ChatMessage, Chunk } from "../contracts/conversation.js";
import type { Logger, Span } from "../contracts/logging.js";
import type { ProviderContract, ProviderEvent, Usage } from "../contracts/provider.js";
import type { EventEmitter, RunTurnInput, RunTurnResult } from "../contracts/runtime.js";
import type { ToolCall, ToolContract } from "../contracts/tool.js";
import { createStepDispatcher, type StepDispatcher } from "./dispatch.js";
import {
	errorEvent,
	reasoningDeltaEvent,
	textDeltaEvent,
	toolCallEvent,
	toolResultEvent,
	usageEvent,
} from "./events.js";

export const MAX_STEPS = 50;

function zeroUsage(): Usage {
	return { inputTokens: 0, outputTokens: 0 };
}

function addUsage(a: Usage, b: Usage): Usage {
	const inputTokens = a.inputTokens + b.inputTokens;
	const outputTokens = a.outputTokens + b.outputTokens;

	if (a.cacheReadTokens !== undefined || b.cacheReadTokens !== undefined) {
		const cacheReadTokens = (a.cacheReadTokens ?? 0) + (b.cacheReadTokens ?? 0);
		if (a.cacheWriteTokens !== undefined || b.cacheWriteTokens !== undefined) {
			return {
				inputTokens,
				outputTokens,
				cacheReadTokens,
				cacheWriteTokens: (a.cacheWriteTokens ?? 0) + (b.cacheWriteTokens ?? 0),
			};
		}
		return { inputTokens, outputTokens, cacheReadTokens };
	}

	if (a.cacheWriteTokens !== undefined || b.cacheWriteTokens !== undefined) {
		return {
			inputTokens,
			outputTokens,
			cacheWriteTokens: (a.cacheWriteTokens ?? 0) + (b.cacheWriteTokens ?? 0),
		};
	}

	return { inputTokens, outputTokens };
}

function appendTextDelta(chunks: Chunk[], delta: string): void {
	const lastIdx = chunks.length - 1;
	const last = chunks[lastIdx];
	if (last !== undefined && last.type === "text") {
		chunks[lastIdx] = { type: "text", text: last.text + delta };
	} else {
		chunks.push({ type: "text", text: delta });
	}
}

function appendThinkingDelta(chunks: Chunk[], delta: string): void {
	const lastIdx = chunks.length - 1;
	const last = chunks[lastIdx];
	if (last !== undefined && last.type === "thinking") {
		chunks[lastIdx] = { type: "thinking", text: last.text + delta };
	} else {
		chunks.push({ type: "thinking", text: delta });
	}
}

interface StepContext {
	readonly provider: ProviderContract;
	readonly messages: ChatMessage[];
	readonly tools: readonly ToolContract[];
	readonly toolMap: Map<string, ToolContract>;
	readonly dispatch: RunTurnInput["dispatch"];
	readonly emit: EventEmitter;
	readonly signal: AbortSignal;
	readonly conversationId: string;
	readonly turnId: string;
	readonly logger: Logger;
	readonly toolSpans: Map<string, Span>;
}

interface StepResult {
	readonly assistantMessage: ChatMessage | undefined;
	readonly toolCalls: ToolCall[];
	readonly toolMessages: ChatMessage[];
	readonly usage: Usage;
	readonly finishReason: string;
}

function processEvent(
	event: ProviderEvent,
	chunks: Chunk[],
	toolCalls: ToolCall[],
	dispatcher: StepDispatcher,
	ctx: StepContext,
): void {
	switch (event.type) {
		case "text-delta":
			appendTextDelta(chunks, event.delta);
			ctx.emit(textDeltaEvent(ctx.conversationId, ctx.turnId, event.delta));
			break;
		case "reasoning-delta":
			appendThinkingDelta(chunks, event.delta);
			ctx.emit(reasoningDeltaEvent(ctx.conversationId, ctx.turnId, event.delta));
			break;
		case "tool-call": {
			const call: ToolCall = {
				id: event.toolCallId,
				name: event.toolName,
				input: event.input,
			};
			toolCalls.push(call);
			chunks.push({
				type: "tool-call",
				toolCallId: event.toolCallId,
				toolName: event.toolName,
				input: event.input,
			});
			ctx.emit(
				toolCallEvent(
					ctx.conversationId,
					ctx.turnId,
					event.toolCallId,
					event.toolName,
					event.input,
				),
			);

			// Open a tool-call span (attrs: name, toolCallId)
			try {
				const tcSpan = ctx.logger.span("tool-call", {
					name: event.toolName,
					toolCallId: event.toolCallId,
				});
				ctx.toolSpans.set(event.toolCallId, tcSpan);
			} catch {
				// Swallow — D7: logging never breaks the turn.
			}

			if (ctx.dispatch.eager) {
				dispatcher.submit(call);
			}
			break;
		}
		case "usage":
			ctx.emit(usageEvent(ctx.conversationId, ctx.turnId, event.usage));
			break;
		case "finish":
			break;
		case "error":
			if (event.code !== undefined) {
				chunks.push({ type: "error", message: event.message, code: event.code });
			} else {
				chunks.push({ type: "error", message: event.message });
			}
			ctx.emit(errorEvent(ctx.conversationId, ctx.turnId, event.message, event.code));
			break;
	}
}

async function executeStep(ctx: StepContext): Promise<StepResult> {
	const chunks: Chunk[] = [];
	const toolCalls: ToolCall[] = [];
	let stepUsage = zeroUsage();
	let finishReason = "stop";

	// Open a step span with the verbatim pre-mutation prompt in its body (BEFORE capture).
	let stepSpan: Span | undefined;
	try {
		stepSpan = ctx.logger.span("step");
		// Emit the verbatim pre-mutation prompt as a log record on the step span's logger.
		// This is the "BEFORE" capture — the messages + tools as handed to provider.stream.
		stepSpan.log.info("prompt:before", {
			"prompt.messages": JSON.stringify(ctx.messages),
			"prompt.tools": JSON.stringify(
				ctx.tools.map((t) => ({
					name: t.name,
					description: t.description,
					parameters: t.parameters,
				})),
			),
		});
	} catch {
		// Swallow — D7.
	}

	const dispatcher = createStepDispatcher(
		ctx.toolMap,
		ctx.dispatch,
		ctx.signal,
		ctx.emit,
		ctx.conversationId,
		ctx.turnId,
		ctx.toolSpans,
	);

	try {
		const stream = ctx.provider.stream(ctx.messages, ctx.tools);
		for await (const event of stream) {
			if (ctx.signal.aborted) break;
			processEvent(event, chunks, toolCalls, dispatcher, ctx);
			if (event.type === "usage") {
				stepUsage = addUsage(stepUsage, event.usage);
			}
			if (event.type === "finish") {
				finishReason = event.reason;
			}
		}
	} catch (err) {
		const message = err instanceof Error ? err.message : String(err);
		chunks.push({ type: "error", message });
		ctx.emit(errorEvent(ctx.conversationId, ctx.turnId, message));
		finishReason = "error";
		// Close step span with error
		try {
			stepSpan?.end({ err });
		} catch {
			// Swallow — D7.
		}
		stepSpan = undefined;
	}

	if (!ctx.dispatch.eager) {
		for (const call of toolCalls) {
			dispatcher.submit(call);
		}
	}

	const results = await dispatcher.drain();

	// Close remaining tool-call spans
	for (const call of toolCalls) {
		const tcSpan = ctx.toolSpans.get(call.id);
		if (tcSpan !== undefined) {
			const result = results.get(call.id);
			try {
				tcSpan.end({
					attrs: {
						isError: result?.isError ?? false,
						contentLength: result?.content.length ?? 0,
					},
				});
			} catch {
				// Swallow — D7.
			}
			ctx.toolSpans.delete(call.id);
		}
	}

	const toolMessages: ChatMessage[] = [];
	for (const call of toolCalls) {
		const result = results.get(call.id);
		if (result !== undefined) {
			const isError = result.isError ?? false;
			ctx.emit(
				toolResultEvent(
					ctx.conversationId,
					ctx.turnId,
					call.id,
					call.name,
					result.content,
					isError,
				),
			);
			toolMessages.push({
				role: "tool",
				chunks: [
					{
						type: "tool-result",
						toolCallId: call.id,
						toolName: call.name,
						content: result.content,
						isError,
					},
				],
			});
		}
	}

	// Close step span (if not already closed by error)
	if (stepSpan !== undefined) {
		try {
			stepSpan.end({
				attrs: {
					finishReason,
					usage_inputTokens: stepUsage.inputTokens,
					usage_outputTokens: stepUsage.outputTokens,
				},
			});
		} catch {
			// Swallow — D7.
		}
	}

	const assistantMessage: ChatMessage | undefined =
		chunks.length > 0 ? { role: "assistant", chunks } : undefined;

	return { assistantMessage, toolCalls, toolMessages, usage: stepUsage, finishReason };
}

export async function runTurn(input: RunTurnInput): Promise<RunTurnResult> {
	const messages: ChatMessage[] = [...input.messages];
	const resultMessages: ChatMessage[] = [];
	let totalUsage = zeroUsage();
	let finishReason = "stop";

	const toolMap = new Map<string, ToolContract>();
	for (const tool of input.tools) {
		toolMap.set(tool.name, tool);
	}

	const conversationId = input.conversationId;
	const turnId = input.turnId;
	const signal = input.signal ?? new AbortController().signal;
	const logger = input.logger;

	// Open a turn span (attrs: conversationId, turnId, model)
	let turnSpan: Span | undefined;
	if (logger !== undefined) {
		try {
			turnSpan = logger.span("turn", {
				conversationId,
				turnId,
				model: input.providerOpts?.model ?? input.provider.id,
			});
		} catch {
			// Swallow — D7.
		}
	}

	// Track open tool-call spans across steps so we can close them on abort
	const toolSpans = new Map<string, Span>();

	try {
		for (let step = 0; step < MAX_STEPS; step++) {
			if (signal.aborted) {
				finishReason = "aborted";
				break;
			}

			const stepResult = await executeStep({
				provider: input.provider,
				messages,
				tools: input.tools,
				toolMap,
				dispatch: input.dispatch,
				emit: input.emit,
				signal,
				conversationId,
				turnId,
				logger: turnSpan?.log ?? logger ?? createNoopLogger(),
				toolSpans,
			});

			totalUsage = addUsage(totalUsage, stepResult.usage);

			if (stepResult.assistantMessage !== undefined) {
				messages.push(stepResult.assistantMessage);
				resultMessages.push(stepResult.assistantMessage);
			}

			for (const msg of stepResult.toolMessages) {
				messages.push(msg);
				resultMessages.push(msg);
			}

			if (signal.aborted) {
				finishReason = "aborted";
				break;
			}

			if (stepResult.toolCalls.length === 0) {
				finishReason = stepResult.finishReason;
				break;
			}

			if (step === MAX_STEPS - 1) {
				finishReason = "max-steps";
			}
		}
	} finally {
		// Close any orphaned tool-call spans (e.g. abort mid-tool)
		for (const [id, tcSpan] of toolSpans) {
			try {
				tcSpan.end({ attrs: { orphaned: true } });
			} catch {
				// Swallow — D7.
			}
			toolSpans.delete(id);
		}

		// Close the turn span
		if (turnSpan !== undefined) {
			try {
				turnSpan.end({
					attrs: {
						finishReason,
						usage_inputTokens: totalUsage.inputTokens,
						usage_outputTokens: totalUsage.outputTokens,
					},
				});
			} catch {
				// Swallow — D7.
			}
		}
	}

	return { messages: resultMessages, usage: totalUsage, finishReason };
}

function createNoopLogger(): Logger {
	return {
		debug() {},
		info() {},
		warn() {},
		error() {},
		child() {
			return createNoopLogger();
		},
		span() {
			return {
				id: "noop",
				log: createNoopLogger(),
				setAttributes() {},
				addLink() {},
				child() {
					return this;
				},
				end() {},
			};
		},
	};
}