import {
	Capability,
	hasCapability,
	languageModels,
} from "@giselles-ai/language-model";
import {
	type Connection,
	isAppEntryNode,
	isContentGenerationNode,
	isDataQueryNode,
	isDataStoreNode,
	isEndNode,
	isFileNode,
	isImageGenerationNode,
	isQueryNode,
	isTextGenerationNode,
	isVectorStoreNode,
	type NodeLike,
} from "@giselles-ai/protocol";

export type ConnectionValidationResult =
	| { canConnect: true }
	| { canConnect: false; message: string };

export interface ConnectionValidationOptions {
	existingConnections?: Connection[];
}

export function isSupportedConnection(
	outputNode: NodeLike,
	inputNode: NodeLike,
	options?: ConnectionValidationOptions,
): ConnectionValidationResult {
	// prevent self-loop
	if (outputNode.id === inputNode.id) {
		return {
			canConnect: false,
			message: "Connecting to the same node is not allowed",
		};
	}

	// only operation node can receive inputs
	if (inputNode.type !== "operation") {
		return {
			canConnect: false,
			message: "This node does not receive inputs",
		};
	}

	// v2container: AppEntryNode and EndNode cannot be connected directly.
	if (
		(isAppEntryNode(outputNode) && isEndNode(inputNode)) ||
		(isEndNode(outputNode) && isAppEntryNode(inputNode))
	) {
		return {
			canConnect: false,
			message: "Connecting Start and End nodes is not allowed",
		};
	}

	if (inputNode.content.type === "end") {
		return { canConnect: true };
	}

	// prevent unsupported inputs for image generation node
	if (isImageGenerationNode(inputNode)) {
		if (outputNode.content.type === "github") {
			return {
				canConnect: false,
				message: "GitHub node is not supported as an input for this node",
			};
		}
	}

	// trigger and action node can be connected to any node (except ImageGenerationNode, handled above)
	if (
		outputNode.content.type === "trigger" ||
		outputNode.content.type === "action"
	) {
		return {
			canConnect: true,
		};
	}

	// image generation can be connected to generation node if the model have a capability to handle generated image input
	if (isImageGenerationNode(outputNode)) {
		if (!isTextGenerationNode(inputNode) && !isImageGenerationNode(inputNode)) {
			return {
				canConnect: false,
				message:
					"Image generation node can only be connected to text generation or image generation",
			};
		}

		const inputNodeLLMId = inputNode.content.llm.id;
		const inputNodeLanguageModel = languageModels.find(
			(languageModel) => languageModel.id === inputNodeLLMId,
		);

		if (inputNodeLanguageModel === undefined) {
			return {
				canConnect: false,
				message: "This node is not supported as an input for Image generation",
			};
		}
		if (
			hasCapability(inputNodeLanguageModel, Capability.ImageGenerationInput)
		) {
			return {
				canConnect: true,
			};
		}

		return {
			canConnect: false,
			message:
				"Image generation node is not supported as an input for this node",
		};
	}
	if (outputNode.content.type === "github") {
		return {
			canConnect: false,
			message: "GitHub node is not supported as an output",
		};
	}

	// file can be connected to generation node if the model have a capability to handle file input
	if (isFileNode(outputNode)) {
		if (!isTextGenerationNode(inputNode) && !isImageGenerationNode(inputNode)) {
			return {
				canConnect: false,
				message: "File node is not supported as an input for this node",
			};
		}

		const inputNodeLLMId = inputNode.content.llm.id;
		const inputNodeLanguageModel = languageModels.find(
			(languageModel) => languageModel.id === inputNodeLLMId,
		);

		if (inputNodeLanguageModel === undefined) {
			return {
				canConnect: false,
				message: "This node is not supported as an input for File",
			};
		}
		if (hasCapability(inputNodeLanguageModel, Capability.GenericFileInput)) {
			return {
				canConnect: true,
			};
		}
		if (outputNode.content.category === "text") {
			return {
				canConnect: true,
			};
		}

		if (
			outputNode.content.category === "image" &&
			hasCapability(inputNodeLanguageModel, Capability.ImageFileInput)
		) {
			return {
				canConnect: true,
			};
		}

		if (
			outputNode.content.category === "pdf" &&
			hasCapability(inputNodeLanguageModel, Capability.PdfFileInput)
		) {
			return {
				canConnect: true,
			};
		}

		return {
			canConnect: false,
			message: "File node is not supported as an input for this node",
		};
	}

	// Vector store node can only be connected to query node
	if (isVectorStoreNode(outputNode)) {
		if (isQueryNode(inputNode)) {
			return {
				canConnect: true,
			};
		}
		return {
			canConnect: false,
			message: "Vector store node can only be connected to query node",
		};
	}

	// Data store node can only be connected to data query, text generation, image generation, or content generation node
	if (isDataStoreNode(outputNode)) {
		if (
			isDataQueryNode(inputNode) ||
			isTextGenerationNode(inputNode) ||
			isImageGenerationNode(inputNode) ||
			isContentGenerationNode(inputNode)
		) {
			// Data Query can only have one Data Store connected
			if (isDataQueryNode(inputNode) && options?.existingConnections) {
				const hasExistingDataStore = options.existingConnections.some(
					(conn) =>
						conn.inputNode.id === inputNode.id &&
						conn.outputNode.content.type === "dataStore",
				);
				if (hasExistingDataStore) {
					return {
						canConnect: false,
						message: "Data Query node can only have one Data Store connected",
					};
				}
			}
			return {
				canConnect: true,
			};
		}
		return {
			canConnect: false,
			message:
				"Data store node can only be connected to data query or generation nodes",
		};
	}

	// query can only be connected to text generation or image generation
	if (isQueryNode(outputNode)) {
		if (!isTextGenerationNode(inputNode) && !isImageGenerationNode(inputNode)) {
			return {
				canConnect: false,
				message:
					"Query node can only be connected to text generation or image generation",
			};
		}
	}

	// data query can only be connected to text generation, image generation, or content generation
	if (isDataQueryNode(outputNode)) {
		if (
			!isTextGenerationNode(inputNode) &&
			!isImageGenerationNode(inputNode) &&
			!isContentGenerationNode(inputNode)
		) {
			return {
				canConnect: false,
				message:
					"Data query node can only be connected to text generation, image generation, or content generation",
			};
		}
	}

	// data query node can only receive inputs from data store, text, trigger, action, app entry, or generation nodes
	if (isDataQueryNode(inputNode)) {
		const allowedOutputTypes = [
			"dataStore",
			"text",
			"trigger",
			"action",
			"appEntry",
			"textGeneration",
			"contentGeneration",
		];
		if (!allowedOutputTypes.includes(outputNode.content.type)) {
			return {
				canConnect: false,
				message: "This node is not supported as an input for Data Query",
			};
		}
	}

	return {
		canConnect: true,
	};
}
