import {
	type Api,
	type AssistantMessage,
	type AssistantMessageEvent,
	type AssistantMessageEventStream,
	calculateCost,
	type Model,
	type StopReason,
} from "@earendil-works/pi-ai";
import type { AcpPromptMeta } from "./types.js";

function cloneMessage(message: AssistantMessage): AssistantMessage {
	return {
		...message,
		content: message.content.map((block) => ({ ...block })),
		usage: { ...message.usage, cost: { ...message.usage.cost } },
	};
}

export function createOutput(model: Model<Api>): AssistantMessage {
	return {
		role: "assistant",
		content: [],
		api: model.api,
		provider: model.provider,
		model: model.id,
		usage: {
			input: 0,
			output: 0,
			cacheRead: 0,
			cacheWrite: 0,
			totalTokens: 0,
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
		},
		stopReason: "pending",
		timestamp: Date.now(),
	};
}

export function applyUsage(model: Model<Api>, output: AssistantMessage, meta?: AcpPromptMeta): void {
	if (!meta) return;
	output.usage.input = meta.inputTokens ?? 0;
	output.usage.output = meta.outputTokens ?? 0;
	output.usage.cacheRead = meta.cachedReadTokens ?? 0;
	output.usage.cacheWrite = 0;
	output.usage.totalTokens =
		meta.totalTokens ?? output.usage.input + output.usage.output + output.usage.cacheRead;
	calculateCost(model, output.usage);
	const responseId = meta.requestId ?? meta.promptId;
	if (responseId) output.responseId = responseId;
}

export function mapStopReason(reason: string | undefined): StopReason {
	const normalized = reason?.toLowerCase();
	if (normalized === "max_tokens" || normalized === "length") return "length";
	if (normalized === "aborted" || normalized === "cancelled" || normalized === "canceled") return "aborted";
	if (normalized === "error") return "error";
	return "stop";
}

export class ContentEmitter {
	private started = false;
	private thinkingIndex = -1;
	private textIndex = -1;
	private readonly stream: AssistantMessageEventStream;
	private readonly output: AssistantMessage;

	constructor(stream: AssistantMessageEventStream, output: AssistantMessage) {
		this.stream = stream;
		this.output = output;
	}

	start(): void {
		if (this.started) return;
		this.started = true;
		this.push({ type: "start", partial: this.output });
	}

	appendThinking(delta: string): void {
		if (!delta) return;
		this.start();
		this.closeText();
		if (this.thinkingIndex < 0) {
			this.thinkingIndex = this.output.content.length;
			this.output.content.push({ type: "thinking", thinking: "" });
			this.push({ type: "thinking_start", contentIndex: this.thinkingIndex, partial: this.output });
		}
		const block = this.output.content[this.thinkingIndex];
		if (block?.type !== "thinking") return;
		block.thinking += delta;
		this.push({ type: "thinking_delta", contentIndex: this.thinkingIndex, delta, partial: this.output });
	}

	appendText(delta: string): void {
		if (!delta) return;
		this.start();
		this.closeThinking();
		if (this.textIndex < 0) {
			this.textIndex = this.output.content.length;
			this.output.content.push({ type: "text", text: "" });
			this.push({ type: "text_start", contentIndex: this.textIndex, partial: this.output });
		}
		const block = this.output.content[this.textIndex];
		if (block?.type !== "text") return;
		block.text += delta;
		this.push({ type: "text_delta", contentIndex: this.textIndex, delta, partial: this.output });
	}

	noteActivity(text: string): void {
		this.appendThinking(`${text}\n`);
		this.closeThinking();
	}

	done(reason: StopReason): void {
		this.start();
		this.closeThinking();
		this.closeText();
		this.output.stopReason = reason;
		const doneReason = reason === "length" ? "length" : "stop";
		this.push({ type: "done", reason: doneReason, message: this.output });
		this.stream.end(cloneMessage(this.output));
	}

	error(message: string, aborted: boolean): void {
		this.start();
		this.closeThinking();
		this.closeText();
		this.output.stopReason = aborted ? "aborted" : "error";
		this.output.errorMessage = message;
		if (this.output.content.length === 0) this.output.content.push({ type: "text", text: message });
		this.push({ type: "error", reason: this.output.stopReason, error: this.output });
		this.stream.end(cloneMessage(this.output));
	}

	private closeThinking(): void {
		if (this.thinkingIndex < 0) return;
		const index = this.thinkingIndex;
		this.thinkingIndex = -1;
		const block = this.output.content[index];
		if (block?.type === "thinking") {
			this.push({ type: "thinking_end", contentIndex: index, content: block.thinking, partial: this.output });
		}
	}

	private closeText(): void {
		if (this.textIndex < 0) return;
		const index = this.textIndex;
		this.textIndex = -1;
		const block = this.output.content[index];
		if (block?.type === "text") {
			this.push({ type: "text_end", contentIndex: index, content: block.text, partial: this.output });
		}
	}

	private push(event: AssistantMessageEvent): void {
		if ("partial" in event) this.stream.push({ ...event, partial: cloneMessage(this.output) });
		else if (event.type === "done") this.stream.push({ ...event, message: cloneMessage(this.output) });
		else if (event.type === "error") this.stream.push({ ...event, error: cloneMessage(this.output) });
		else this.stream.push(event);
	}
}