nibbleai/main.ts

import { Notice, Plugin, TFile } from "obsidian";
import { Timer } from "src/recorder/Timer";
import { Controls } from "src/recorder/Controls";
import { AudioHandler } from "src/transcription/AudioHandler";
import { NibbleAISettingsTab } from "src/NibbleAISettingsTab";
import { SettingsManager, PluginSettings } from "src/SettingsManager";
import { NativeAudioRecorder } from "src/recorder/AudioRecorder";
import { RecordingStatus, StatusBar } from "src/recorder/StatusBar";
import { getExtensionFromMimeType } from "src/transcription/utils";
import { Proofreader } from "src/proofreading/Proofreader";
import { TitleGenerator } from "src/title/TitleGenerator";
import { PropertyFiller } from "src/property/PropertyFiller";

export default class NibbleAI extends Plugin {
	settings: PluginSettings;
	settingsManager: SettingsManager;
	timer: Timer;
	recorder: NativeAudioRecorder;
	audioHandler: AudioHandler;
	controls: Controls | null = null;
	statusBar: StatusBar;
	proofreader: Proofreader;
	titleGenerator: TitleGenerator;
	propertyFiller: PropertyFiller;

	async onload() {
		this.settingsManager = new SettingsManager(this);
		this.settings = await this.settingsManager.loadSettings();

		this.addRibbonIcon("bot", "Open NibbleAI recording controls", () => {
			this.openControls();
		});

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

		this.timer = new Timer();
		this.audioHandler = new AudioHandler(this);
		this.recorder = new NativeAudioRecorder();
		const deviceId =
			this.settings.audioDeviceId === "default"
				? null
				: this.settings.audioDeviceId;
		this.recorder.setDeviceId(deviceId);

		this.statusBar = new StatusBar(this);

		this.proofreader = new Proofreader(this);
		this.titleGenerator = new TitleGenerator(this);
		this.propertyFiller = new PropertyFiller(this);

		// Timer reacts to recording state changes
		this.statusBar.onChange((status) => {
			switch (status) {
				case RecordingStatus.Recording:
					this.timer.start();
					break;
				case RecordingStatus.Paused:
					this.timer.pause();
					break;
				case RecordingStatus.Processing:
				case RecordingStatus.Idle:
					this.timer.reset();
					break;
			}
		});

		this.addCommands();
		this.registerEvents();
		this.registerUriHandler();
	}

	onunload() {
		if (this.controls) {
			this.controls.close();
		}
		this.statusBar.remove();
	}

	// ── Recording state transitions ──

	async startRecording() {
		if (
			this.statusBar.status === RecordingStatus.Recording ||
			this.statusBar.status === RecordingStatus.Paused
		) {
			new Notice("Already recording");
			return;
		}
		try {
			await this.recorder.startRecording();
			this.statusBar.updateStatus(RecordingStatus.Recording);
			new Notice("Recording...");
		} catch (err) {
			this.statusBar.updateStatus(RecordingStatus.Idle);
			new Notice("✘ Could not start recording");
		}
	}

	async stopRecording() {
		if (
			this.statusBar.status !== RecordingStatus.Recording &&
			this.statusBar.status !== RecordingStatus.Paused
		) {
			return;
		}
		this.statusBar.updateStatus(RecordingStatus.Processing);
		const audioBlob = await this.recorder.stopRecording();
		const extension = getExtensionFromMimeType(this.recorder.getMimeType());
		const fileName = `${new Date()
			.toISOString()
			.replace(/[:.]/g, "-")}.${extension}`;
		await this.audioHandler.sendAudioData(audioBlob, fileName);
		this.statusBar.updateStatus(RecordingStatus.Idle);
	}

	async pauseRecording() {
		if (this.statusBar.status === RecordingStatus.Recording) {
			await this.recorder.pauseRecording();
			this.statusBar.updateStatus(RecordingStatus.Paused);
			new Notice("Recording paused");
		} else if (this.statusBar.status === RecordingStatus.Paused) {
			await this.recorder.pauseRecording();
			this.statusBar.updateStatus(RecordingStatus.Recording);
			new Notice("Recording resumed");
		}
	}

	async cancelRecording() {
		if (
			this.statusBar.status !== RecordingStatus.Recording &&
			this.statusBar.status !== RecordingStatus.Paused
		) {
			return;
		}
		await this.recorder.stopRecording();
		this.statusBar.updateStatus(RecordingStatus.Idle);
		new Notice("Recording cancelled");
	}

	openControls() {
		if (!this.controls) {
			this.controls = new Controls(this);
		}
		this.controls.open();
	}

	// ── Commands ──

	addCommands() {
		// Transcription / Recording
		this.addCommand({
			id: "start-stop-recording",
			name: "Start/stop recording",
			icon: "mic",
			callback: async () => {
				if (
					this.statusBar.status !== RecordingStatus.Recording &&
					this.statusBar.status !== RecordingStatus.Paused
				) {
					await this.startRecording();
				} else {
					await this.stopRecording();
				}
			},
			hotkeys: [{ modifiers: ["Alt"], key: "Q" }],
		});

		this.addCommand({
			id: "upload-audio-file",
			name: "Upload audio file",
			icon: "upload",
			callback: () => {
				const fileInput = document.createElement("input");
				fileInput.type = "file";
				fileInput.accept =
					"audio/*,video/*,.mp4,.m4a,.wav,.webm,.ogg,.mp3";
				fileInput.onchange = async (event) => {
					const files = (event.target as HTMLInputElement).files;
					if (files && files.length > 0) {
						const file = files[0];
						const audioBlob = file.slice(0, file.size, file.type);
						await this.audioHandler.sendAudioData(
							audioBlob,
							file.name
						);
					}
				};
				fileInput.click();
			},
		});

		this.addCommand({
			id: "pause-resume-recording",
			name: "Pause/resume recording",
			icon: "pause",
			callback: () => this.pauseRecording(),
		});

		this.addCommand({
			id: "open-recording-controls",
			name: "Open recording controls",
			icon: "popup-open",
			callback: () => this.openControls(),
		});

		// Proofreading
		this.addCommand({
			id: "fix-punctuation-grammar",
			name: "Proofread with AI",
			icon: "wand",
			editorCallback: (editor) => {
				const selection = editor.getSelection();
				if (!selection || selection.trim().length === 0) {
					new Notice("No text selected.");
					return;
				}
				this.proofreader.fixSelection(editor);
			},
		});

		// Title Generation
		this.addCommand({
			id: "generate-title",
			name: "Generate title for current note",
			icon: "heading",
			callback: async () => {
				await this.titleGenerator.generateTitle();
			},
		});

		// Property Filler
		this.addCommand({
			id: "fill-property",
			name: "Fill property with AI",
			icon: "pencil",
			callback: async () => {
				await this.propertyFiller.fillProperty();
			},
		});

		// Dynamic preset property commands from settings
		const propertyNames = this.settings.propertyPresets
			? this.settings.propertyPresets.split(",").map(s => s.trim()).filter(s => s.length > 0)
			: [];

		const propertyIcons: Record<string, string> = {
			summary: "list",
			tags: "tags",
		};

		for (const prop of propertyNames) {
			const safeId = prop.toLowerCase().replace(/[^a-z0-9]/g, "-");
			this.addCommand({
				id: `fill-property-${safeId}`,
				name: `Fill "${prop}" property`,
				icon: propertyIcons[prop] || "pencil",
				callback: async () => {
					await this.propertyFiller.fillSpecificProperty(prop);
				},
			});
		}
	}

	// ── Events ──

	registerEvents() {
		// Right-click audio files to transcribe
		this.registerEvent(
			this.app.workspace.on("file-menu", (menu, file) => {
				if (!(file instanceof TFile)) return;

				const audioExtensions = [
					".mp3",
					".mp4",
					".mpeg",
					".mpga",
					".m4a",
					".wav",
					".webm",
					".ogg",
				];
				if (!audioExtensions.some((ext) => file.path.endsWith(ext)))
					return;

				menu.addItem((item) => {
					item.setTitle("Transcribe audio file 🔊")
						.setIcon("document")
						.onClick(async () => {
							const audioBlob = new Blob([
								await file.vault.readBinary(file),
							]);
							await this.audioHandler.sendAudioData(
								audioBlob,
								file.name
							);
						});
				});
			})
		);

		// Right-click selected text to fix punctuation
		this.registerEvent(
			this.app.workspace.on("editor-menu", (menu, editor) => {
				const selection = editor.getSelection();
				if (!selection || selection.trim().length === 0) return;
				if (!this.settings.proofreadingContextMenu) return;

				menu.addItem((item) => {
					item
						.setTitle("Proofread with AI")
						.setIcon("wand")
						.onClick(async () => {
							await this.proofreader.fixSelection(editor);
						});
				});
			})
		);
	}

	// ── URI Handler ──

	registerUriHandler() {
		this.registerObsidianProtocolHandler("nibbleai", async (params) => {
			const command = params.command;
			if (!command) {
				this.openControls();
				return;
			}

			switch (command) {
				case "start":
					await this.startRecording();
					break;
				case "stop":
					await this.stopRecording();
					break;
				case "pause":
					await this.pauseRecording();
					break;
				case "cancel":
					await this.cancelRecording();
					break;
				default:
					new Notice(`✘ Unknown nibbleai command: ${command}`);
			}
		});
	}
}