import type {
	IExecuteFunctions,
	INodeExecutionData,
	INodeType,
	INodeTypeDescription,
	Logger,
} from 'n8n-workflow';
import { NodeConnectionType, NodeOperationError } from 'n8n-workflow';
import { pipeline } from '@huggingface/transformers';

export class NamedEntityRecognition implements INodeType {
	description: INodeTypeDescription = {
		displayName: 'Named Entity Recognition',
		name: 'namedEntityRecognition',
		icon: 'fa:tags',
		group: ['ai'],
		version: 1,
		description: 'Extracts named entities (like persons, locations, organizations) from text using Transformers.js.',
		defaults: {
			name: 'Named Entity Recognition',
		},
		inputs: [NodeConnectionType.Main],
		outputs: [NodeConnectionType.Main],
		usableAsTool: true,
		properties: [
			{
				displayName: 'Input Text',
				name: 'inputText',
				type: 'string',
				typeOptions: {
					rows: 5,
				},
				default: '',
				required: true,
				description: 'The text from which to extract named entities.',
				placeholder: 'e.g., N8N is a workflow automation tool based in Berlin, developed by Johannes and a great team.',
			},
			{
				displayName: 'Aggregation Strategy',
				name: 'aggregationStrategy',
				type: 'options',
				options: [
					{ name: 'None', value: 'none' },
					{ name: 'Simple (Recommended)', value: 'simple' },
					{ name: 'First', value: 'first' },
					{ name: 'Average', value: 'average' },
					{ name: 'Max', value: 'max' },
				],
				default: 'simple',
				description: "Strategy to group token parts (e.g. B-PER, I-PER) into single entities. 'Simple' is often a good default.",
			},
			{
				displayName: 'Output Field Name',
				name: 'outputFieldName',
				type: 'string',
				default: 'entities',
				required: true,
				description: 'The field name where the array of extracted entities will be stored.',
			},
		],
	};

	private static pipelineCache: Map<string, any> = new Map();

	async execute(this: IExecuteFunctions): Promise<INodeExecutionData[][]> {
		const itemIndex = 0; // We will only ever process one item.

		// Define the model to use
		const model = 'Xenova/bert-base-NER';

		// Get node parameters
		const outputFieldName = this.getNodeParameter('outputFieldName', itemIndex, 'entities') as string;
		const aggregationStrategy = this.getNodeParameter('aggregationStrategy', itemIndex, 'simple') as string;
		const inputText = this.getNodeParameter('inputText', itemIndex, '') as string;

		if (!inputText || !inputText.trim()) {
			throw new NodeOperationError(this.getNode(), 'Input Text parameter is required. Please provide a static value or an expression.');
		}

		// Initialize the pipeline (cache it to avoid reloading)
		let pipe = NamedEntityRecognition.pipelineCache.get(model);
		if (!pipe) {
			try {
				this.logger.info(`Loading NER model: ${model}`);
				pipe = await pipeline('token-classification', model);
				NamedEntityRecognition.pipelineCache.set(model, pipe);
				this.logger.info(`NER Model ${model} loaded successfully`);
			} catch (error) {
				throw new NodeOperationError(
					this.getNode(),
					`Failed to load NER model: ${(error as Error).message}`,
				);
			}
		}

		// Process the text
		try {
			this.logger.info(`Extracting entities from text...`);

			const pipelineOptions: any = {};
			if (aggregationStrategy && aggregationStrategy !== 'none') {
				pipelineOptions.aggregation_strategy = aggregationStrategy;
			}

			const processedOutput = await pipe(inputText, pipelineOptions);
			this.logger.info(`Raw NER output: ${JSON.stringify(processedOutput)}`);

			if (!Array.isArray(processedOutput)) {
				throw new NodeOperationError(this.getNode(), 'NER output is not in the expected array format.');
			}

			const resultJson: any = {};
			resultJson[outputFieldName] = processedOutput;

			return [this.helpers.returnJsonArray([{ json: resultJson }])];

		} catch (error) {
			if (this.continueOnFail()) {
				const errorItem = {
					json: {
						error: (error as Error).message,
					},
				};
				return [this.helpers.returnJsonArray([errorItem])];
			} else {
				throw new NodeOperationError(this.getNode(), (error as Error).message || String(error));
			}
		}
	}
}