import type { BaseLanguageModel } from '@langchain/core/language_models/base';
import type { VectorStore } from '@langchain/core/vectorstores';
import { VectorDBQAChain } from '@langchain/classic/chains';
import { VectorStoreQATool } from '@langchain/classic/tools';
import type {
	IExecuteFunctions,
	INodeExecutionData,
	INodeType,
	INodeTypeDescription,
	ISupplyDataFunctions,
	SupplyData,
} from 'n8n-workflow';
import { NodeConnectionTypes, nodeNameToToolName } from 'n8n-workflow';

import { logWrapper, getConnectionHintNoticeField } from '@n8n/ai-utilities';

async function getTool(
	ctx: ISupplyDataFunctions | IExecuteFunctions,
	itemIndex: number,
): Promise<VectorStoreQATool> {
	const node = ctx.getNode();
	const { typeVersion } = node;
	const name =
		typeVersion <= 1
			? (ctx.getNodeParameter('name', itemIndex) as string)
			: nodeNameToToolName(node);
	const toolDescription = ctx.getNodeParameter('description', itemIndex) as string;
	const topK = ctx.getNodeParameter('topK', itemIndex, 4) as number;
	const description = VectorStoreQATool.getDescription(name, toolDescription);
	const vectorStore = (await ctx.getInputConnectionData(
		NodeConnectionTypes.AiVectorStore,
		itemIndex,
	)) as VectorStore;
	const llm = (await ctx.getInputConnectionData(
		NodeConnectionTypes.AiLanguageModel,
		itemIndex,
	)) as BaseLanguageModel;

	const vectorStoreTool = new VectorStoreQATool(name, description, {
		llm,
		vectorStore,
	});

	vectorStoreTool.chain = VectorDBQAChain.fromLLM(llm, vectorStore, {
		k: topK,
	});

	return vectorStoreTool;
}

export class ToolVectorStore implements INodeType {
	description: INodeTypeDescription = {
		displayName: 'Vector Store Question Answer Tool',
		name: 'toolVectorStore',
		icon: 'node:vector-store-question-answer-tool',
		iconColor: 'black',
		group: ['transform'],
		version: [1, 1.1],
		description: 'Answer questions with a vector store',
		defaults: {
			name: 'Answer questions with a vector store',
		},
		codex: {
			categories: ['AI'],
			subcategories: {
				AI: ['Tools'],
				Tools: ['Other Tools'],
			},
			resources: {
				primaryDocumentation: [
					{
						url: 'https://docs.n8n.io/integrations/builtin/cluster-nodes/sub-nodes/n8n-nodes-langchain.toolvectorstore/',
					},
				],
			},
		},

		inputs: [
			{
				displayName: 'Vector Store',
				maxConnections: 1,
				type: NodeConnectionTypes.AiVectorStore,
				required: true,
			},
			{
				displayName: 'Model',
				maxConnections: 1,
				type: NodeConnectionTypes.AiLanguageModel,
				required: true,
			},
		],

		outputs: [NodeConnectionTypes.AiTool],
		outputNames: ['Tool'],
		builderHint: {
			inputs: {
				ai_vectorStore: { required: true },
				ai_languageModel: { required: true },
			},
		},
		properties: [
			getConnectionHintNoticeField([NodeConnectionTypes.AiAgent]),
			{
				displayName: 'Data Name',
				name: 'name',
				type: 'string',
				default: '',
				placeholder: 'e.g. users_info',
				validateType: 'string-alphanumeric',
				description:
					'Name of the data in vector store. This will be used to fill this tool description: Useful for when you need to answer questions about [name]. Whenever you need information about [data description], you should ALWAYS use this. Input should be a fully formed question.',
				displayOptions: {
					show: {
						'@version': [1],
					},
				},
			},
			{
				displayName: 'Description of Data',
				name: 'description',
				type: 'string',
				default: '',
				placeholder: "[Describe your data here, e.g. a user's name, email, etc.]",
				description:
					'Describe the data in vector store. This will be used to fill this tool description: Useful for when you need to answer questions about [name]. Whenever you need information about [data description], you should ALWAYS use this. Input should be a fully formed question.',
				typeOptions: {
					rows: 3,
				},
			},
			{
				displayName: 'Limit',
				name: 'topK',
				type: 'number',
				default: 4,
				description: 'The maximum number of results to return',
			},
		],
	};

	async supplyData(this: ISupplyDataFunctions, itemIndex: number): Promise<SupplyData> {
		const vectorStoreTool = await getTool(this, itemIndex);

		return {
			response: logWrapper(vectorStoreTool, this),
		};
	}

	async execute(this: IExecuteFunctions): Promise<INodeExecutionData[][]> {
		const inputData = this.getInputData();
		const result: INodeExecutionData[] = [];
		for (let itemIndex = 0; itemIndex < inputData.length; itemIndex++) {
			const tool = await getTool(this, itemIndex);
			const outputData = await tool.invoke(inputData[itemIndex].json);
			result.push({
				json: {
					response: outputData,
				},
				pairedItem: {
					item: itemIndex,
				},
			});
		}

		return [result];
	}
}