main.ts

import { App, Editor, EditorPosition, MarkdownView, Modal, Notice, Plugin, PluginSettingTab, Setting, FileSystemAdapter } from 'obsidian';
import { ChildProcess, exec, spawn } from 'child_process';
import { existsSync, writeFileSync, mkdirSync } from 'fs';
import { Socket } from 'net';
import { resolve } from 'path';
const connect = require('nrepl-client').connect;

type nreplClientExtraFns = {
	clone: Function,
	close: Function,
	describe: Function,
	eval: Function,
	interrupt: Function,
	loadFile: Function,
	lsSessions: Function,
	stdin: Function,
	send: Function,
};
type nreplClient = Socket & nreplClientExtraFns;

const DEFAULT_OUTPUT_LIMIT = 1000;
const REPL_CLIENTS: { [k: string]: nreplClient } = {};
const PROCESSES: ChildProcess[] = [];
const logPrefix = '[obsidian-babashka]';


function getVaultRoot(app: App) {
	return (app.vault.adapter as FileSystemAdapter).getBasePath();
}

interface Codeblock {
	lang: string,
	code: string,
	insideInsertPos: EditorPosition,
	outsideInsertPos: EditorPosition,
}

function getCodeblockForCursor(editor: Editor): Codeblock | null {
	const src = editor.getValue();
	const cursor = editor.getCursor();
	const cljRe = /^```(clojure|clojurescript)$\s([\s\S]*?)^```$/gmd;
	const cljCodeblocks = Array.from(src.matchAll(cljRe));
	const cursorOffset = editor.posToOffset(cursor);
	const insideMatch = (offset: any) => (match: any) => match.indices[0][0] <= offset && offset <= match.indices[0][1];
	const match: any = cljCodeblocks.find(insideMatch(cursorOffset));

	if (match) {
		return {
			lang: match[1],
			code: match[2],
			insideInsertPos: editor.offsetToPos(match.indices[2][1]),
			outsideInsertPos: editor.offsetToPos(match.indices[0][1]),
		}
	} else {
		return null;
	}
}

function validateSettingsForCodeblock(codeblock: Codeblock, settings: PluginSettings) {
	const { lang } = codeblock;
	const { bbPath, nbbPath, nodePath } = settings;

	if (lang == 'clojure' && !bbPath) {
		new Notice('Please set Babashka path in settings');
		return false;
	}

	if (lang == 'clojurescript' && (!nbbPath || !nodePath)) {
		new Notice('Please set both Node Babashka and Node paths in settings');
		return false;
	}

	return true;
}

function capCharsAt(str: string, maxChars: number) {
	if (str.length > maxChars) {
		return str.substring(0, maxChars) + `\n... and ${str.length - maxChars} more chars`;
	} else {
		return str;
	}
}

function resolvePluginFolder(vaultPath: string, settings: PluginSettings) {
	return resolve(vaultPath, settings.bbDir);
}

function writePluginFolder(app: App, settings: PluginSettings) {
	const vaultPath = getVaultRoot(app);
	const pluginFolder = resolvePluginFolder(vaultPath, settings);
	const genFolder = `${pluginFolder}/gen`;

	if (!existsSync(genFolder)) {
		mkdirSync(genFolder, { recursive: true });
	}
}

function writeVaultBindingsNs(app: App, view: MarkdownView, settings: PluginSettings) {
	const vaultPath = getVaultRoot(this.app);
	const pluginFolder = resolvePluginFolder(vaultPath, settings);
	const bindingsNsPath = `${pluginFolder}/gen/vault_bindings.cljc`;

	const src = `(ns vault-bindings)
;; This file is auto-generated by the Babashka plugin when you eval a codeblock.
;; Don't edit it directly, it will be overwritten.
(def *vault-name* "${app.vault.getName()}")
(def *vault-path* "${vaultPath}")
(def *vault-babashka-path* "${pluginFolder}")
(def *last-folder-path* "${vaultPath}/${view.file.parent.path}")
(def *last-file-name* "${view.file.name}")
(def *last-file-path* "${vaultPath}/${view.file.path}")
(def *vault-plugin-folder* "${pluginFolder}")
`;

	writeFileSync(bindingsNsPath, src);
}

function writeDefaultBbEdn(app: App, settings: PluginSettings) {
	const vaultPath = getVaultRoot(this.app);
	const pluginFolder = resolvePluginFolder(vaultPath, settings);
	const bbEdnPath = `${pluginFolder}/bb.edn`;
	const nbbEdnPath = `${pluginFolder}/nbb.edn`;
	const src = `{:paths ["gen"]}`;

	if (!existsSync(bbEdnPath)) {
		writeFileSync(bbEdnPath, src);
	}

	if (!existsSync(nbbEdnPath)) {
		writeFileSync(nbbEdnPath, src);
	}
}

function setupPluginFolder(app: App, view: MarkdownView, settings: PluginSettings) {
	writePluginFolder(app, settings);
	writeVaultBindingsNs(app, view, settings);
	writeDefaultBbEdn(app, settings);
}

function debug(str: string) {
	console.debug(`${logPrefix} ${str}`);
}

function printOutput(output: string, type: "stdout" | "stderr" | "value", limitOutput: boolean, notice = true) {
	const cappedOutput = limitOutput ? capCharsAt(output, DEFAULT_OUTPUT_LIMIT) : output;
	debug(`${type}:\n${cappedOutput}`);
	if (cappedOutput && notice) {
		new Notice(`${type} during eval:\n${cappedOutput}`, 5000);
	}
	return cappedOutput;
}

function printError(error: Error, during: string) {
	const errorMsg = `Error during ${during}:\n` + error;
	debug(errorMsg)
	new Notice(errorMsg, 5000);
}

function insertOutput(editor: Editor, output: string, insertAt: "inside" | "outside", codeblock: Codeblock) {
	const { insideInsertPos, outsideInsertPos } = codeblock;

	if (insertAt == "inside") {
		const outputAsComments = output.trim().replaceAll(/^/gm, ';; ');
		editor.replaceRange(`\n${outputAsComments}\n`, insideInsertPos, insideInsertPos);
	} else if (insertAt == "outside") {
		editor.replaceRange(`\n\n${output}`, outsideInsertPos, outsideInsertPos);
	}
}

function removeProcess(p: ChildProcess) {
	const index = PROCESSES.indexOf(p);
	if (index > -1) {
		PROCESSES.splice(index, 1);
	}
}

function evalCodeblockInCLI(
	codeblock: Codeblock, insertAt: "inside" | "outside",
	vaultPath: string, editor: Editor, settings: PluginSettings) {
	const { lang, code } = codeblock;
	const { bbPath, nbbPath, nodePath, limitOutput } = settings;

	const bin = lang == 'clojure' ? bbPath : `${nodePath} ${nbbPath}`;
	const pluginCwd = resolvePluginFolder(vaultPath, settings);

	const codeShellStr = code.replaceAll("\"", "\\\"");
	const cmd = `${bin} -e "${codeShellStr}"`;

	debug(`eval: \`${cmd}\` on ${pluginCwd}`);
	const evalProcess = exec(
		cmd,
		{ cwd: `${pluginCwd}` },
		(err, stdout, stderr) => {
			if (err) {
				printError(err, 'eval');
				return;
			}

			printOutput(stderr, "stderr", limitOutput);
			const out = printOutput(stdout, "stdout", limitOutput, false);
			if (out) {
				insertOutput(editor, out, insertAt, codeblock);
			}
			removeProcess(evalProcess);
		});

	PROCESSES.push(evalProcess);
}

function evalCodeblockInRepl(
	client: nreplClient, codeblock: Codeblock, insertAt: "inside" | "outside",
	editor: Editor, settings: PluginSettings) {
	const { code } = codeblock;
	const { limitOutput } = settings;

	debug(`repl eval: \`${code}\``);
	client.eval(code, function (err: Error, result: any[]) {
		if (err) {
			printError(err, 'repl eval');
			return;
		}

		const stdout = result.filter((r) => r.out).map((r) => r.out).join('');
		const stderr = result.filter((r) => r.err).map((r) => r.err).join('');
		const value = result.filter((r) => r.value).map((r) => r.value).at(-1);

		printOutput(stderr, "stderr", limitOutput);
		printOutput(stdout, "stdout", limitOutput, false);
		printOutput(value, "value", limitOutput, false);

		const outAndValue = stdout + value;
		const cappedOutAndValue = limitOutput ? capCharsAt(outAndValue, DEFAULT_OUTPUT_LIMIT) : outAndValue;
		if (cappedOutAndValue) {
			insertOutput(editor, cappedOutAndValue, insertAt, codeblock);
		}
	});
}

interface PluginSettings {
	bbDir: string,
	bbPath: string,
	nbbPath: string,
	nodePath: string,
	limitOutput: boolean,
	bbnReplPort: number,
	nbbnReplPort: number,
}

const DEFAULT_SETTINGS: PluginSettings = {
	bbDir: '.babashka',
	bbPath: '',
	nbbPath: '',
	nodePath: '',
	limitOutput: true,
	bbnReplPort: 1667,
	nbbnReplPort: 1668,
}

function validateAndSetup(codeblock: Codeblock, view: MarkdownView, settings: PluginSettings): boolean {
	if (validateSettingsForCodeblock(codeblock, settings)) {
		setupPluginFolder(app, view, settings);
		return true;
	}
	return false;
}

function evalAndInsert(
	app: App, editor: Editor, view: MarkdownView,
	settings: PluginSettings, insertAt: "inside" | "outside") {
	const codeblock = getCodeblockForCursor(editor);
	if (!codeblock) {
		new Notice('No clojure(script) codeblock found at cursor');
	} else if (validateAndSetup(codeblock, view, settings)) {
		const client = REPL_CLIENTS[codeblock.lang];
		if (client) {
			evalCodeblockInRepl(client, codeblock, insertAt, editor, settings);
		} else {
			evalCodeblockInCLI(codeblock, insertAt, getVaultRoot(app), editor, settings);
		}
	}
}

function statusBarText() {
	return Object.keys(REPL_CLIENTS).map(lang => lang == 'clojure' ? 'CLJ' : 'CLJS').join('+');
}

function killClient(lang: string, client = REPL_CLIENTS[lang]) {
	if (client && !client.destroyed) {
		client.destroy();
	}
	if (REPL_CLIENTS[lang] = client) {
		delete REPL_CLIENTS[lang];
	}
}

function startAndConnectRepl(app: App, editor: Editor, view: MarkdownView, settings: PluginSettings, statusBarItemEl: HTMLElement) {
	const codeblock = getCodeblockForCursor(editor);
	if (!codeblock) {
		new Notice('No clojure(script) codeblock found at cursor');
	} else if (validateAndSetup(codeblock, view, settings)) {
		const lang = codeblock.lang;
		if (REPL_CLIENTS[lang]) {
			new Notice(`A ${lang} repl is already connected.`);
			return;
		}

		const vaultPath = getVaultRoot(app);
		const { bbPath, nbbPath, nodePath, bbnReplPort, nbbnReplPort } = settings;
		const pluginCwd = resolvePluginFolder(vaultPath, settings);

		const bin = lang == 'clojure' ? bbPath : nodePath;
		const port = lang == 'clojure' ? bbnReplPort : nbbnReplPort;
		const args = lang == 'clojure' ? ['nrepl-server', '' + port] : [nbbPath, 'nrepl-server', ':port', '' + port];

		debug(`repl: \`${bin + ' ' + args.join(' ')}\` on ${pluginCwd}`);
		const replProcess = spawn(bin, args, { cwd: `${pluginCwd}` });
		PROCESSES.push(replProcess);

		replProcess.stdout.on('data', (data) => {
			debug(`repl stdout: ${data}`);
			// On success,
			// bb says
			//   Started nREPL server at 127.0.0.1:1667
			// nbb says
			//   nREPL server started on port 1667 on host 127.0.0.1 - nrepl://127.0.0.1:1667
			// On port error, neither say nREPL
			// Matching on 'nREPL' seems ok, really hope this doesn't change D:
			if (data.toString().includes('nREPL')) {
				debug(`repl stdout: nREPL detected, connecting...`);
				killClient(lang);
				const client = connect({ port });
				REPL_CLIENTS[lang] = client;
				client.once('error', (err: Error) => {
					printError(err, 'repl connect');
					killClient(lang, client);
				});
				client.once('close', () => {
					killClient(lang, client);
					statusBarItemEl.setText(statusBarText());
				});
				client.once('connect', () => {
					debug(`repl stdout: nREPL connected`);
					new Notice(`Connected to ${lang} nREPL server`);
					statusBarItemEl.setText(statusBarText());
				});
			}
		});

		replProcess.stderr.on('data', (data) => {
			debug(`repl stderr: ${data}`);
		});

		replProcess.on('close', (code) => {
			debug(`repl close: child process exited with code ${code}`);
			removeProcess(replProcess);
			killClient(lang);
		});
	}
}

function killAllProcesses() {
	new Notice(`Sent kill signal to ${PROCESSES.length} processes and ${Object.keys(REPL_CLIENTS).length} nREPL clients`);
	for (const lang in REPL_CLIENTS) {
		killClient(lang);
	}

	for (const p of PROCESSES) {
		p.kill();
	}
}

export default class BabashkaPlugin extends Plugin {
	settings: PluginSettings;

	async onload() {
		await this.loadSettings();

		const statusBarItemEl = this.addStatusBarItem();

		this.addSettingTab(new SettingTab(this.app, this));

		this.addCommand({
			id: 'babashka-eval-codeblock',
			name: 'Eval codeblock',
			editorCallback: async (editor: Editor, view: MarkdownView) => {
				evalAndInsert(this.app, editor, view, this.settings, "inside");
			}
		});

		this.addCommand({
			id: 'babashka-eval-codeblock-print-outside',
			name: 'Eval codeblock and print value outside',
			editorCallback: async (editor: Editor, view: MarkdownView) => {
				evalAndInsert(this.app, editor, view, this.settings, "outside");
			}
		});

		this.addCommand({
			id: 'babashka-start-and-connect-nrepl',
			name: 'Start and connect to a nREPL server',
			editorCallback: async (editor: Editor, view: MarkdownView) => {
				startAndConnectRepl(this.app, editor, view, this.settings, statusBarItemEl);
			}
		});

		this.addCommand({
			id: 'babashka-kill-all-processes',
			name: 'Kill all babashka processes',
			editorCallback: async (editor: Editor, view: MarkdownView) => {
				killAllProcesses();
			}
		});
	}

	onunload() {
		killAllProcesses();
	}

	async loadSettings() {
		this.settings = Object.assign({}, DEFAULT_SETTINGS, await this.loadData());
	}

	async saveSettings() {
		await this.saveData(this.settings);
	}
}

class SettingTab extends PluginSettingTab {
	plugin: BabashkaPlugin;

	constructor(app: App, plugin: BabashkaPlugin) {
		super(app, plugin);
		this.plugin = plugin;
	}

	addTextSetting(name: string, desc: string, placeholder: string, k: keyof PluginSettings) {
		const { containerEl } = this;
		new Setting(containerEl)
			.setName(name)
			.setDesc(desc)
			.addText(text => text
				.setPlaceholder(placeholder)
				.setValue(this.plugin.settings[k] as string)
				.onChange(async (value) => {
					(this.plugin.settings as any)[k] = value;
					await this.plugin.saveSettings();
				}));
	}

	addToggleSetting(name: string, desc: string, k: keyof PluginSettings) {
		const { containerEl } = this;
		new Setting(containerEl)
			.setName(name)
			.setDesc(desc)
			.addToggle(toggle => toggle
				.setValue(this.plugin.settings[k] as boolean)
				.onChange(async (value) => {
					(this.plugin.settings as any)[k] = value;
					await this.plugin.saveSettings();
				}));
	}

	addNumberSetting(name: string, desc: string, placeholder: string, k: keyof PluginSettings) {
		const { containerEl } = this;
		new Setting(containerEl)
			.setName(name)
			.setDesc(desc)
			.addText(toggle => toggle
				.setPlaceholder(placeholder)
				.setValue(`${this.plugin.settings[k]}`)
				.onChange(async (value) => {
					const num = parseInt(value);
					if (!isNaN(num)) {
						(this.plugin.settings as any)[k] = num;
						await this.plugin.saveSettings();
					}
				}));
	}

	display(): void {
		const { containerEl } = this;

		containerEl.empty();

		containerEl.createEl('h2', { text: 'General' });
		this.addToggleSetting('Limit output', `Output will be truncated after ${DEFAULT_OUTPUT_LIMIT} characters.`, 'limitOutput');
		this.addNumberSetting('Babashka nREPL port', '', '1667', 'bbnReplPort');
		this.addNumberSetting('Node Babashka nREPL port', '', '1668', 'nbbnReplPort');

		containerEl.createEl('h2', { text: 'Paths' });
		this.addTextSetting('Vault Babashka dir', 'Path to babashka dir from vault root. Can be absolute. Babashka will be run from this dir. You can put bb.edn, nbb.edn, and package.json files there to use dependencies.', '', 'bbDir');
		this.addTextSetting('Babashka path', 'Absolute path to babashka.', 'run `which bb` to see it', 'bbPath');
		this.addTextSetting('Node Babashka path', 'Absolute path to nbb.', 'run `which nbb` to see it', 'nbbPath');
		this.addTextSetting('Node path', 'Absolute path to node.', 'run `which nbb` to see it', 'nodePath');
	}
}