import { BaseCallbackHandler } from '@langchain/core/callbacks/base';
import type { SerializedFields } from '@langchain/core/dist/load/map_keys';
import { getModelNameForTiktoken } from '@langchain/core/language_models/base';
import type {
	Serialized,
	SerializedNotImplemented,
	SerializedSecret,
} from '@langchain/core/load/serializable';
import type { BaseMessage } from '@langchain/core/messages';
import type { LLMResult } from '@langchain/core/outputs';
import pick from 'lodash/pick';
import type { IDataObject, ISupplyDataFunctions, JsonObject } from 'n8n-workflow';
import { NodeConnectionTypes, NodeError, NodeOperationError } from 'n8n-workflow';

import { logAiEvent } from './log-ai-event';
import { estimateTokensFromStringList } from './tokenizer/token-estimator';

/** Normalized token usage returned by TokensUsageParser. */
type TokenUsageResult = {
	completionTokens: number;
	promptTokens: number;
	totalTokens: number;
	/** Cost may be undefined when the provider returns token counts but no pricing fields. */
	cost?: number;
};

/** Raw provider tokenUsage payload. Some providers report `totalCost` instead of `cost`. */
type ProviderTokenUsageResult = TokenUsageResult & {
	totalCost?: number;
};

type TokensUsageParser = (result: LLMResult) => TokenUsageResult;

type RunDetail = {
	index: number;
	messages: BaseMessage[] | string[] | string;
	options: SerializedSecret | SerializedNotImplemented | SerializedFields;
};

const TIKTOKEN_ESTIMATE_MODEL = 'gpt-4o';

type TracingWriter = {
	setMetadata: (metadata: { tracing: LlmTokenTracingMetadata }) => void;
};

/** Keys written by `applyTracingTokenMetadata` into execution tracing metadata. */
type LlmTokenTracingMetadata = {
	'llm.tokens.in': number;
	'llm.tokens.out': number;
	'llm.tokens.total': number;
	'llm.tokens.estimated': boolean;
	'llm.cost.total'?: number;
};

function canWriteTracingMetadata(context: unknown): context is TracingWriter {
	return (
		typeof context === 'object' &&
		context !== null &&
		'setMetadata' in context &&
		typeof context.setMetadata === 'function'
	);
}

export class N8nLlmTracing extends BaseCallbackHandler {
	name = 'N8nLlmTracing';

	// This flag makes sure that LangChain will wait for the handlers to finish before continuing
	// This is crucial for the handleLLMError handler to work correctly (it should be called before the error is propagated to the root node)
	awaitHandlers = true;

	connectionType = NodeConnectionTypes.AiLanguageModel;

	promptTokensEstimate = 0;

	completionTokensEstimate = 0;

	#parentRunIndex?: number;

	/**
	 * A map to associate LLM run IDs to run details.
	 * Key: Unique identifier for each LLM run (run ID)
	 * Value: RunDetails object
	 *
	 */
	runsMap: Record<string, RunDetail> = {};

	options: {
		tokensUsageParser: TokensUsageParser;
		errorDescriptionMapper: (error: NodeError) => string | null | undefined;
	} = {
		// Default(OpenAI format) parser
		tokensUsageParser: (result: LLMResult) => {
			const tokenUsage = result?.llmOutput?.tokenUsage as
				| Partial<ProviderTokenUsageResult>
				| undefined;
			const completionTokens = tokenUsage?.completionTokens ?? 0;
			const promptTokens = tokenUsage?.promptTokens ?? 0;
			const cost = tokenUsage?.cost ?? tokenUsage?.totalCost;

			return {
				completionTokens,
				promptTokens,
				totalTokens: completionTokens + promptTokens,
				cost,
			};
		},
		errorDescriptionMapper: (error: NodeError) => error.description,
	};

	constructor(
		private executionFunctions: ISupplyDataFunctions,
		options?: {
			tokensUsageParser?: TokensUsageParser;
			errorDescriptionMapper?: (error: NodeError) => string;
		},
	) {
		super();
		this.options = { ...this.options, ...options };
	}

	async estimateTokensFromGeneration(generations: LLMResult['generations']) {
		const messages = generations.flatMap((gen) => gen.map((g) => g.text));
		return await this.estimateTokensFromStringList(messages);
	}

	async estimateTokensFromStringList(list: string[]) {
		const embeddingModel = getModelNameForTiktoken(TIKTOKEN_ESTIMATE_MODEL);
		return await estimateTokensFromStringList(list, embeddingModel);
	}

	async handleLLMEnd(output: LLMResult, runId: string) {
		// The fallback should never happen since handleLLMStart should always set the run details
		// but just in case, we set the index to the length of the runsMap
		const runDetails = this.runsMap[runId] ?? { index: Object.keys(this.runsMap).length };

		output.generations = output.generations.map((gen) =>
			gen.map((g) => pick(g, ['text', 'generationInfo'])),
		);

		const tokenUsageEstimate = {
			completionTokens: 0,
			promptTokens: 0,
			totalTokens: 0,
		};
		const tokenUsage = this.options.tokensUsageParser(output);

		if (output.generations.length > 0) {
			tokenUsageEstimate.completionTokens = await this.estimateTokensFromGeneration(
				output.generations,
			);

			tokenUsageEstimate.promptTokens = this.promptTokensEstimate;
			tokenUsageEstimate.totalTokens =
				tokenUsageEstimate.completionTokens + this.promptTokensEstimate;
		}
		const response: {
			response: { generations: LLMResult['generations'] };
			tokenUsageEstimate?: typeof tokenUsageEstimate;
			tokenUsage?: typeof tokenUsage;
		} = {
			response: { generations: output.generations },
		};

		// If the LLM response contains actual tokens usage, otherwise fallback to the estimate
		if (tokenUsage.completionTokens > 0) {
			response.tokenUsage = tokenUsage;
			this.applyTracingTokenMetadata({
				promptTokens: tokenUsage.promptTokens,
				completionTokens: tokenUsage.completionTokens,
				totalTokens: tokenUsage.totalTokens,
				isEstimated: false,
				cost: tokenUsage.cost,
			});
		} else {
			response.tokenUsageEstimate = tokenUsageEstimate;
			this.applyTracingTokenMetadata({
				promptTokens: tokenUsageEstimate.promptTokens,
				completionTokens: tokenUsageEstimate.completionTokens,
				totalTokens: tokenUsageEstimate.totalTokens,
				isEstimated: true,
			});
		}

		const parsedMessages =
			typeof runDetails.messages === 'string'
				? runDetails.messages
				: runDetails.messages.map((message) => {
						if (typeof message === 'string') return message;
						if (typeof message?.toJSON === 'function') return message.toJSON();

						return message;
					});

		const sourceNodeRunIndex =
			this.#parentRunIndex !== undefined ? this.#parentRunIndex + runDetails.index : undefined;

		this.executionFunctions.addOutputData(
			this.connectionType,
			runDetails.index,
			[[{ json: { ...response } }]],
			undefined,
			sourceNodeRunIndex,
		);

		logAiEvent(this.executionFunctions, 'ai-llm-generated-output', {
			messages: parsedMessages,
			options: runDetails.options,
			response,
		});
	}

	async handleLLMStart(llm: Serialized, prompts: string[], runId: string) {
		const estimatedTokens = await this.estimateTokensFromStringList(prompts);
		const sourceNodeRunIndex =
			this.#parentRunIndex !== undefined
				? this.#parentRunIndex + this.executionFunctions.getNextRunIndex()
				: undefined;

		const options = llm.type === 'constructor' ? llm.kwargs : llm;
		const { index } = this.executionFunctions.addInputData(
			this.connectionType,
			[
				[
					{
						json: {
							messages: prompts,
							estimatedTokens,
							options,
						},
					},
				],
			],
			sourceNodeRunIndex,
		);

		// Save the run details for later use when processing `handleLLMEnd` event
		this.runsMap[runId] = {
			index,
			options,
			messages: prompts,
		};
		this.promptTokensEstimate = estimatedTokens;
	}

	async handleLLMError(error: IDataObject | Error, runId: string, parentRunId?: string) {
		const runDetails = this.runsMap[runId] ?? { index: Object.keys(this.runsMap).length };

		// Filter out non-x- headers to avoid leaking sensitive information in logs
		// eslint-disable-next-line no-prototype-builtins
		if (typeof error === 'object' && error?.hasOwnProperty('headers')) {
			const errorWithHeaders = error as { headers: Record<string, unknown> };

			Object.keys(errorWithHeaders.headers).forEach((key) => {
				if (!key.startsWith('x-')) {
					delete errorWithHeaders.headers[key];
				}
			});
		}

		if (error instanceof NodeError) {
			if (this.options.errorDescriptionMapper) {
				error.description = this.options.errorDescriptionMapper(error);
			}

			this.executionFunctions.addOutputData(this.connectionType, runDetails.index, error);
		} else {
			// If the error is not a NodeError, we wrap it in a NodeOperationError
			this.executionFunctions.addOutputData(
				this.connectionType,
				runDetails.index,
				new NodeOperationError(this.executionFunctions.getNode(), error as JsonObject, {
					functionality: 'configuration-node',
				}),
			);
		}

		logAiEvent(this.executionFunctions, 'ai-llm-errored', {
			// eslint-disable-next-line @typescript-eslint/no-base-to-string
			error: Object.keys(error).length === 0 ? error.toString() : error,
			runId,
			parentRunId,
		});
	}

	// Used to associate subsequent runs with the correct parent run in subnodes of subnodes
	setParentRunIndex(runIndex: number) {
		this.#parentRunIndex = runIndex;
	}

	private applyTracingTokenMetadata(params: {
		promptTokens: number;
		completionTokens: number;
		totalTokens: number;
		isEstimated: boolean;
		cost?: number;
	}) {
		if (!canWriteTracingMetadata(this.executionFunctions)) return;

		const tracing: LlmTokenTracingMetadata = {
			'llm.tokens.in': params.promptTokens,
			'llm.tokens.out': params.completionTokens,
			'llm.tokens.total': params.totalTokens,
			'llm.tokens.estimated': params.isEstimated,
		};
		if (typeof params.cost === 'number' && Number.isFinite(params.cost)) {
			tracing['llm.cost.total'] = params.cost;
		}

		this.executionFunctions.setMetadata({ tracing });
	}
}