From d0c3169aa0176d13fc0268ccd49290a53cb74140 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Thu, 7 Sep 2023 18:26:41 +0100
Subject: [PATCH] Support unicode and tag aliases

---
 py/autocomplete.py            |  2 +-
 web/js/autocompleter.js       | 94 +++++++++++++++++++++++++++++++----
 web/js/common/autocomplete.js |  8 ++-
 3 files changed, 90 insertions(+), 14 deletions(-)

diff --git a/py/autocomplete.py b/py/autocomplete.py
index 26810e0..abf838d 100644
--- a/py/autocomplete.py
+++ b/py/autocomplete.py
@@ -17,6 +17,6 @@ async def get_autocomplete(request):
 
 @PromptServer.instance.routes.post("/pysssss/autocomplete")
 async def update_autocomplete(request):
-    with open(file, "w") as f:
+    with open(file, "w", encoding="utf-8") as f:
         f.write(await request.text())
     return web.Response(status=200)
diff --git a/web/js/autocompleter.js b/web/js/autocompleter.js
index 6ef9772..4811554 100644
--- a/web/js/autocompleter.js
+++ b/web/js/autocompleter.js
@@ -5,6 +5,59 @@ import { $el, ComfyDialog } from "../../../../scripts/ui.js";
 import { TextAreaAutoComplete } from "./common/autocomplete.js";
 import { ModelInfoDialog } from "./common/modelInfoDialog.js";
 
+function parseCSV(csvText) {
+	const rows = [];
+	const delimiter = ",";
+	const quote = '"';
+	let currentField = "";
+	let inQuotedField = false;
+
+	function pushField() {
+		rows[rows.length - 1].push(currentField);
+		currentField = "";
+		inQuotedField = false;
+	}
+
+	rows.push([]); // Initialize the first row
+
+	for (let i = 0; i < csvText.length; i++) {
+		const char = csvText[i];
+		const nextChar = csvText[i + 1];
+
+		if (!inQuotedField) {
+			if (char === quote) {
+				inQuotedField = true;
+			} else if (char === delimiter) {
+				pushField();
+			} else if (char === "\r" || char === "\n" || i === csvText.length - 1) {
+				pushField();
+				if (nextChar === "\n") {
+					i++; // Handle Windows line endings (\r\n)
+				}
+				rows.push([]); // Start a new row
+			} else {
+				currentField += char;
+			}
+		} else {
+			if (char === quote && nextChar === quote) {
+				currentField += quote;
+				i++; // Skip the next quote
+			} else if (char === quote) {
+				inQuotedField = false;
+			} else {
+				currentField += char;
+			}
+		}
+	}
+
+	// Remove the last row if it's empty
+	if (rows[rows.length - 1].length === 0) {
+		rows.pop();
+	}
+
+	return rows;
+}
+
 async function getCustomWords() {
 	const resp = await api.fetchApi("/pysssss/autocomplete", { cache: "no-store" });
 	if (resp.status === 200) {
@@ -20,18 +73,37 @@ async function addCustomWords(text) {
 	if (text) {
 		TextAreaAutoComplete.updateWords(
 			"pysssss.customwords",
-			text.split("\n").reduce((p, n) => {
-				n = n.trim();
-				const pos = n.lastIndexOf(",");
-				let priority = undefined;
-				if (pos > -1) {
-					const v = parseInt(n.substr(pos + 1).trim());
-					if (!isNaN(v)) {
-						priority = v;
-						n = n.substr(0, pos).trim();
-					}
+			parseCSV(text).reduce((p, n) => {
+				let text;
+				let priority;
+				let value;
+				let num;
+				switch (n.length) {
+					case 0:
+						return;
+					case 1:
+						// Single word
+						text = n[0];
+						break;
+					case 2:
+						// Word,[priority|alias]
+						num = +n[1];
+						if (isNaN(num)) {
+							text = n[1];
+							value = n[0];
+						} else {
+							text = n[0];
+							priority = num;
+						}
+						break;
+					default:
+						// Word,alias,priority
+						text = n[1];
+						value = n[0];
+						priority = +n[2];
+						break;
 				}
-				p[n] = { text: n, priority };
+				p[text] = { text, priority, value };
 				return p;
 			}, {})
 		);
diff --git a/web/js/common/autocomplete.js b/web/js/common/autocomplete.js
index ecfea7b..9d0a204 100644
--- a/web/js/common/autocomplete.js
+++ b/web/js/common/autocomplete.js
@@ -462,7 +462,7 @@ export class TextAreaAutoComplete {
 	#update() {
 		let before = this.helper.getBeforeCursor();
 		if (before?.length) {
-			const m = before.match(/\b([\w-_:\\\/]+)$/);
+			const m = before.match(/([^\s|,|;|"]+)$/);
 			if (m) {
 				before = m[0];
 			} else {
@@ -523,7 +523,10 @@ export class TextAreaAutoComplete {
 				{
 					onclick: () => {
 						this.el.focus();
-						this.helper.insertAtCursor(wordInfo.text + TextAreaAutoComplete.separator, -before.length);
+						this.helper.insertAtCursor(
+							(wordInfo.value ?? wordInfo.text) + TextAreaAutoComplete.separator,
+							-before.length
+						);
 						setTimeout(() => {
 							this.#update();
 						}, 150);
@@ -563,6 +566,7 @@ export class TextAreaAutoComplete {
 	}
 
 	static updateWords(id, words) {
+		console.log(id, words);
 		const isUpdate = id in TextAreaAutoComplete.groups;
 		TextAreaAutoComplete.groups[id] = words;
 		if (isUpdate) {