PiCloud/dashboard/src/lib/rhai-mode.ts

// CodeMirror StreamLanguage for Rhai.
//
// Keyword and operator lists are sourced from the upstream TextMate
// grammar maintained by the Rhai authors:
//   https://github.com/rhaiscript/vscode-rhai
//   syntax/rhai.tmLanguage.json (MPL-2.0)
// This file does NOT copy the upstream grammar bytes — only the
// symbol lists. The matching logic is a simple regex tokenizer
// tailored to CodeMirror's StreamLanguage shape; if richer
// highlighting is wanted later, swap this out for a full tmLanguage
// loader (vscode-textmate + oniguruma) without touching callers.
//
// SDK completions (`ctx.*`, `log::*`) come from our own SDK contract
// in crates/executor-core/tests/sdk_contract.rs — that file is the
// authoritative list of what scripts can do.

import { StreamLanguage, LanguageSupport } from '@codemirror/language';
import { autocompletion, type CompletionContext, type CompletionResult } from '@codemirror/autocomplete';
import { EditorSelection, StateEffect, StateField, type Extension } from '@codemirror/state';
import { EditorView, keymap, showPanel, type Panel } from '@codemirror/view';
import { parse, buildSymbolTable, type ParseResult, type Range, type SymbolTable } from './rhai';

// Keywords that drive control flow (`if`, `for`, ...) — these get the
// `controlKeyword` tag so the theme can color them distinctly from
// declaration-style keywords like `let` or `fn`.
const CONTROL_KEYWORDS = new Set([
	'if',
	'else',
	'for',
	'while',
	'loop',
	'do',
	'switch',
	'case',
	'default',
	'return',
	'break',
	'continue',
	'try',
	'catch',
	'throw'
]);

const DECLARATION_KEYWORDS = new Set([
	'let',
	'const',
	'fn',
	'private',
	'in',
	'as',
	'is'
]);

// Reserved-but-not-currently-valid keywords from the upstream grammar.
// We still highlight them so users notice them; the parser will reject
// at execute time.
const RESERVED_KEYWORDS = new Set([
	'var',
	'match',
	'public',
	'protected',
	'new',
	'use',
	'with',
	'module',
	'package',
	'super',
	'spawn',
	'thread',
	'go',
	'sync',
	'async',
	'await',
	'yield',
	'void',
	'null',
	'nil',
	'debug',
	'eval',
	'print',
	'import',
	'export'
]);

const BOOLEAN_LITERALS = new Set(['true', 'false']);

const SPECIAL_VARIABLES = new Set(['ctx']);
const NAMESPACES = new Set(['log']);

interface RhaiState {
	inBlockComment: boolean;
	inString: false | '"' | '`';
}

export const rhaiLanguage = StreamLanguage.define<RhaiState>({
	name: 'rhai',
	startState: () => ({ inBlockComment: false, inString: false }),
	token(stream, state) {
		// --- inside a /* … */ block comment ---
		if (state.inBlockComment) {
			if (stream.match(/.*?\*\//)) {
				state.inBlockComment = false;
			} else {
				stream.skipToEnd();
			}
			return 'comment';
		}

		// --- inside a multi-line string (rare but possible with ` ` strings) ---
		if (state.inString) {
			const quote = state.inString;
			while (!stream.eol()) {
				const ch = stream.next();
				if (ch === '\\') {
					stream.next();
					continue;
				}
				if (ch === quote) {
					state.inString = false;
					return 'string';
				}
			}
			return 'string';
		}

		// Skip whitespace
		if (stream.eatSpace()) return null;

		// --- line comment ---
		if (stream.match('//')) {
			stream.skipToEnd();
			return 'comment';
		}
		// --- block comment ---
		if (stream.match('/*')) {
			state.inBlockComment = true;
			if (stream.match(/.*?\*\//)) {
				state.inBlockComment = false;
			} else {
				stream.skipToEnd();
			}
			return 'comment';
		}

		// --- strings ---
		const quote = stream.peek();
		if (quote === '"' || quote === '`') {
			stream.next();
			while (!stream.eol()) {
				const ch = stream.next();
				if (ch === '\\') {
					stream.next();
					continue;
				}
				if (ch === quote) {
					return 'string';
				}
			}
			// String continues on next line (only really valid for ` `).
			state.inString = quote;
			return 'string';
		}

		// --- numbers (hex, binary, decimal, float) ---
		if (stream.match(/^0x[0-9a-fA-F_]+/)) return 'number';
		if (stream.match(/^0b[01_]+/)) return 'number';
		if (stream.match(/^\d[\d_]*(?:\.\d[\d_]*)?(?:[eE][+-]?\d+)?/)) return 'number';

		// --- module path: name::name (used for log::info etc.) ---
		// Recognized before plain identifiers so we can return 'namespace'.
		if (stream.match(/^[a-zA-Z_]\w*(?=::)/)) {
			const word = stream.current();
			if (NAMESPACES.has(word)) return 'namespace';
			return 'variableName';
		}

		// --- identifiers + keywords ---
		if (stream.match(/^[a-zA-Z_]\w*/)) {
			const word = stream.current();
			if (CONTROL_KEYWORDS.has(word)) return 'controlKeyword';
			if (DECLARATION_KEYWORDS.has(word)) return 'keyword';
			if (RESERVED_KEYWORDS.has(word)) return 'keyword';
			if (BOOLEAN_LITERALS.has(word)) return 'bool';
			if (SPECIAL_VARIABLES.has(word)) return 'variableName.special';
			if (NAMESPACES.has(word)) return 'namespace';
			// Followed by `(` → function call. We highlight as a function name.
			if (stream.peek() === '(') return 'function(variableName)';
			// Property after `.`
			const before = stream.string.slice(0, stream.start);
			if (before.endsWith('.')) return 'propertyName';
			return 'variableName';
		}

		// --- operators / punctuation ---
		if (stream.match(/^(\?\?|\.\.=|\.\.|::|==|!=|<=|>=|&&|\|\||<<|>>|=>|->|[+\-*/%<>!&|^~=])/)) {
			return 'operator';
		}
		if (stream.match(/^[(){}[\];,.]/)) return 'punctuation';

		// Unrecognized — advance one char and bail.
		stream.next();
		return null;
	},
	languageData: {
		commentTokens: { line: '//', block: { open: '/*', close: '*/' } },
		closeBrackets: { brackets: ['(', '[', '{', '"', '`'] },
		indentOnInput: /^\s*[}\])]$/
	}
});

// ---------------------------------------------------------------------------
// Autocomplete
// ---------------------------------------------------------------------------

interface CompletionItem {
	label: string;
	detail?: string;
	type?: 'keyword' | 'variable' | 'function' | 'property' | 'namespace';
}

const KEYWORD_COMPLETIONS: CompletionItem[] = [
	...['let', 'const', 'fn'].map((k) => ({ label: k, type: 'keyword' as const, detail: 'declaration' })),
	...['if', 'else', 'for', 'while', 'loop', 'switch', 'return', 'break', 'continue', 'try', 'catch', 'throw'].map(
		(k) => ({ label: k, type: 'keyword' as const, detail: 'control flow' })
	),
	...['in', 'as', 'is'].map((k) => ({ label: k, type: 'keyword' as const })),
	{ label: 'true', type: 'keyword', detail: 'boolean' },
	{ label: 'false', type: 'keyword', detail: 'boolean' }
];

// ctx.* — keep aligned with `build_ctx_map` in
// crates/executor-core/src/engine.rs.
const CTX_TOP_COMPLETIONS: CompletionItem[] = [
	{ label: 'execution_id', type: 'property', detail: 'string' },
	{ label: 'script_id', type: 'property', detail: 'string' },
	{ label: 'script_name', type: 'property', detail: 'string' },
	{ label: 'request_id', type: 'property', detail: 'string' },
	{ label: 'invocation_type', type: 'property', detail: '"http" | "function" | "scheduled"' },
	{ label: 'sdk_version', type: 'property', detail: 'string ("major.minor")' },
	{ label: 'request', type: 'property', detail: 'object' }
];

const CTX_REQUEST_COMPLETIONS: CompletionItem[] = [
	{ label: 'path', type: 'property', detail: 'string' },
	{ label: 'headers', type: 'property', detail: 'map of string→string' },
	{ label: 'body', type: 'property', detail: 'parsed JSON value' },
	{ label: 'params', type: 'property', detail: 'map (param-route captures, SDK 1.1+)' },
	{ label: 'query', type: 'property', detail: 'map (parsed query string, SDK 1.1+)' },
	{ label: 'rest', type: 'property', detail: 'string (prefix-route tail, SDK 1.1+)' }
];

const LOG_COMPLETIONS: CompletionItem[] = [
	{ label: 'info', type: 'function', detail: 'log::info(msg, data?)' },
	{ label: 'warn', type: 'function', detail: 'log::warn(msg, data?)' },
	{ label: 'error', type: 'function', detail: 'log::error(msg, data?)' },
	{ label: 'trace', type: 'function', detail: 'log::trace(msg, data?) — use instead of "debug" (reserved keyword)' }
];

const TOP_LEVEL_GLOBALS: CompletionItem[] = [
	{ label: 'ctx', type: 'variable', detail: 'invocation context' },
	{ label: 'log', type: 'namespace', detail: 'log::info/warn/error/trace' }
];

function toCMCompletions(items: CompletionItem[]) {
	return items.map((c) => ({
		label: c.label,
		type: c.type,
		detail: c.detail
	}));
}

export function rhaiCompletions(context: CompletionContext): CompletionResult | null {
	// `log::` namespace
	const ns = context.matchBefore(/log::\w*/);
	if (ns) {
		return {
			from: ns.from + 'log::'.length,
			options: toCMCompletions(LOG_COMPLETIONS),
			validFor: /^\w*$/
		};
	}

	// `ctx.request.` properties
	const ctxReq = context.matchBefore(/ctx\.request\.\w*/);
	if (ctxReq) {
		return {
			from: ctxReq.from + 'ctx.request.'.length,
			options: toCMCompletions(CTX_REQUEST_COMPLETIONS),
			validFor: /^\w*$/
		};
	}

	// `ctx.` properties
	const ctx = context.matchBefore(/ctx\.\w*/);
	if (ctx) {
		return {
			from: ctx.from + 'ctx.'.length,
			options: toCMCompletions(CTX_TOP_COMPLETIONS),
			validFor: /^\w*$/
		};
	}

	// Member access on something other than `ctx`/`log` — let the
	// script-aware source decide whether it has fields to offer. We bow
	// out so we don't flood the popup with keywords after a `.`.
	if (context.matchBefore(/\.\w*$/)) return null;

	// Plain word at the cursor → keywords + top-level names.
	const word = context.matchBefore(/\w+/);
	if (!word && !context.explicit) return null;
	return {
		from: word ? word.from : context.pos,
		options: toCMCompletions([...KEYWORD_COMPLETIONS, ...TOP_LEVEL_GLOBALS]),
		validFor: /^\w*$/
	};
}

// ---------------------------------------------------------------------------
// Script-aware analysis
// ---------------------------------------------------------------------------

// One AST + symbol table per editor state. Rebuilt on every doc change.
// Scripts in the dashboard are small (20–200 lines is the target); the
// parse + walk is sub-millisecond at that size, so we don't bother
// debouncing for now.
interface RhaiAnalysis {
	parse: ParseResult;
	table: SymbolTable;
}

function analyze(source: string): RhaiAnalysis {
	const parsed = parse(source);
	return { parse: parsed, table: buildSymbolTable(parsed) };
}

export const rhaiAnalysisField = StateField.define<RhaiAnalysis>({
	create: (state) => analyze(state.doc.toString()),
	update: (value, tr) => (tr.docChanged ? analyze(tr.newDoc.toString()) : value)
});

/**
 * Script-aware completion source.
 *
 * Two things on top of the static `ctx.*` / `log::*` list:
 *   1. After `name.`, if `name` was initialized to an object-map
 *      literal, suggest that literal's field names.
 *   2. At a plain word position, suggest user-defined symbols in scope
 *      (locals, function parameters, top-level `fn` decls). `fn` decls
 *      get their signature in the `detail` field so the popup shows
 *      `process(order, user)` rather than just `process`.
 *
 * Composes with `rhaiCompletions` via `autocompletion({ override:
 * [scopeCompletionSource, rhaiCompletions] })` — CodeMirror merges the
 * results.
 */
export function scopeCompletionSource(context: CompletionContext): CompletionResult | null {
	const analysis = context.state.field(rhaiAnalysisField, false);
	if (!analysis) return null;

	// Member access — `obj.fie|`. Only fire if `obj` resolves to a known
	// object-literal in scope; otherwise leave the popup empty rather
	// than guess wrong fields.
	const member = context.matchBefore(/(\w+)\.(\w*)/);
	if (member) {
		const dotIdx = member.text.indexOf('.');
		const objectName = member.text.slice(0, dotIdx);
		// `ctx.*` is handled by the static source — don't double up.
		if (objectName !== 'ctx') {
			const fields = analysis.table.objectFieldsOf(objectName, member.from);
			if (fields.length > 0) {
				return {
					from: member.from + dotIdx + 1,
					options: fields.map((label) => ({ label, type: 'property' })),
					validFor: /^\w*$/
				};
			}
		}
		return null;
	}

	// Skip when right after `::` (handled by the static source for log::).
	if (context.matchBefore(/::\w*$/)) return null;

	// Plain word — surface in-scope decls.
	const word = context.matchBefore(/\w+/);
	if (!word && !context.explicit) return null;
	const decls = analysis.table.scopeCompletions(context.pos);
	if (decls.length === 0) return null;
	return {
		from: word ? word.from : context.pos,
		options: decls.map((d) => ({
			label: d.name,
			type: d.kind === 'fn' ? 'function' : 'variable',
			detail: d.signature ?? d.kind
		})),
		validFor: /^\w*$/
	};
}

// ---------------------------------------------------------------------------
// Go-to-definition, Ctrl/Cmd+Click, find-usages
// ---------------------------------------------------------------------------

interface UsagesPanelState {
	name: string;
	ranges: Range[];
}

const setUsagesPanel = StateEffect.define<UsagesPanelState | null>();

const usagesPanelField = StateField.define<UsagesPanelState | null>({
	create: () => null,
	update(value, tr) {
		for (const e of tr.effects) if (e.is(setUsagesPanel)) return e.value;
		return value;
	},
	provide: (f) => showPanel.from(f, (value) => (value ? buildUsagesPanel : null))
});

function buildUsagesPanel(view: EditorView): Panel {
	const state = view.state.field(usagesPanelField);
	const dom = document.createElement('div');
	dom.className = 'cm-rhai-usages';

	const head = document.createElement('div');
	head.className = 'cm-rhai-usages-head';
	const count = state ? state.ranges.length : 0;
	const label = document.createElement('span');
	label.textContent = `${count} occurrence${count === 1 ? '' : 's'} of "${state?.name ?? ''}"`;
	head.appendChild(label);
	const close = document.createElement('button');
	close.type = 'button';
	close.textContent = '×';
	close.title = 'Close (Esc)';
	close.onclick = () => {
		view.dispatch({ effects: setUsagesPanel.of(null) });
		view.focus();
	};
	head.appendChild(close);
	dom.appendChild(head);

	if (state) {
		for (const r of state.ranges) {
			const line = view.state.doc.lineAt(r.start);
			const row = document.createElement('button');
			row.type = 'button';
			row.className = 'cm-rhai-usages-row';
			const num = document.createElement('span');
			num.className = 'cm-rhai-usages-line';
			num.textContent = String(line.number);
			const snip = document.createElement('span');
			snip.className = 'cm-rhai-usages-snip';
			snip.textContent = line.text.trim();
			row.appendChild(num);
			row.appendChild(snip);
			row.onclick = () => {
				view.dispatch({
					selection: EditorSelection.cursor(r.start),
					scrollIntoView: true
				});
				view.focus();
			};
			dom.appendChild(row);
		}
	}

	return { dom, top: false };
}

const usagesPanelTheme = EditorView.baseTheme({
	'.cm-rhai-usages': {
		background: '#0f172a',
		color: '#e2e8f0',
		borderTop: '1px solid #334155',
		padding: '0.4rem 0.5rem',
		maxHeight: '180px',
		overflowY: 'auto',
		fontSize: '0.85rem'
	},
	'.cm-rhai-usages-head': {
		display: 'flex',
		justifyContent: 'space-between',
		alignItems: 'center',
		color: '#94a3b8',
		marginBottom: '0.25rem'
	},
	'.cm-rhai-usages-head button': {
		background: 'transparent',
		color: '#94a3b8',
		border: 'none',
		fontSize: '1rem',
		cursor: 'pointer',
		padding: '0 0.4rem'
	},
	'.cm-rhai-usages-row': {
		display: 'flex',
		gap: '0.5rem',
		alignItems: 'baseline',
		width: '100%',
		background: 'transparent',
		color: '#e2e8f0',
		border: 'none',
		textAlign: 'left',
		padding: '0.15rem 0.25rem',
		cursor: 'pointer',
		fontFamily: 'inherit',
		fontSize: 'inherit',
		borderRadius: '0.25rem'
	},
	'.cm-rhai-usages-row:hover': {
		background: '#1e293b'
	},
	'.cm-rhai-usages-line': {
		color: '#64748b',
		minWidth: '2.5rem',
		flexShrink: 0
	},
	'.cm-rhai-usages-snip': {
		fontFamily:
			'ui-monospace, SFMono-Regular, Menlo, Consolas, "Liberation Mono", monospace',
		whiteSpace: 'nowrap',
		overflow: 'hidden',
		textOverflow: 'ellipsis'
	}
});

function declAtCursor(view: EditorView, pos: number) {
	const analysis = view.state.field(rhaiAnalysisField, false);
	if (!analysis) return null;
	return analysis.table.declOfUsageAt(pos);
}

export function gotoDefinition(view: EditorView): boolean {
	const pos = view.state.selection.main.head;
	const decl = declAtCursor(view, pos);
	if (!decl) return false;
	view.dispatch({
		selection: EditorSelection.cursor(decl.nameRange.start),
		scrollIntoView: true
	});
	return true;
}

export function findUsages(view: EditorView): boolean {
	const pos = view.state.selection.main.head;
	const analysis = view.state.field(rhaiAnalysisField, false);
	if (!analysis) return false;
	const decl = analysis.table.declOfUsageAt(pos);
	if (!decl) return false;
	view.dispatch({
		effects: setUsagesPanel.of({
			name: decl.name,
			ranges: analysis.table.usagesOf(decl)
		})
	});
	return true;
}

function closeUsagesPanel(view: EditorView): boolean {
	if (!view.state.field(usagesPanelField, false)) return false;
	view.dispatch({ effects: setUsagesPanel.of(null) });
	return true;
}

const rhaiKeymap = keymap.of([
	{ key: 'F12', run: gotoDefinition },
	{ key: 'Shift-F12', run: findUsages },
	{ key: 'Escape', run: closeUsagesPanel }
]);

// Ctrl+Click (Cmd+Click on macOS) on an identifier → jump to its decl.
// Returning `true` suppresses CodeMirror's default selection behavior so
// the click doesn't simultaneously place the caret at the click site.
const ctrlClickHandler = EditorView.domEventHandlers({
	mousedown(event, view) {
		if (!(event.metaKey || event.ctrlKey) || event.button !== 0) return false;
		const pos = view.posAtCoords({ x: event.clientX, y: event.clientY });
		if (pos == null) return false;
		const decl = declAtCursor(view, pos);
		if (!decl) return false;
		view.dispatch({
			selection: EditorSelection.cursor(decl.nameRange.start),
			scrollIntoView: true
		});
		event.preventDefault();
		return true;
	}
});

export function rhai(): LanguageSupport {
	const extensions: Extension[] = [
		rhaiAnalysisField,
		usagesPanelField,
		usagesPanelTheme,
		rhaiKeymap,
		ctrlClickHandler,
		autocompletion({ override: [scopeCompletionSource, rhaiCompletions] })
	];
	return new LanguageSupport(rhaiLanguage, extensions);
}