init.lua

-- fettle — prototype
-- The loop: hotkey -> Cmd+A, Cmd+C -> correct(text) -> paste back -> restore clipboard.
--
-- Correction engine: your own API key + a provider/model you pick from the
-- menu-bar dropdown. No bundled model, so the whole thing stays tiny.
--
-- Two providers supported:
--   * Anthropic — highest quality (Haiku/Sonnet/Opus).
--   * Groq      — much lower latency for short text (small models on fast HW).
-- Correction is an easy task, so Groq's small models are usually plenty good
-- and noticeably snappier. Pick whichever suits the moment from the ✍️ menu.

local HOTKEY_MODS = { "alt" } -- Option
local HOTKEY_KEY = "space" -- Option+Space

local MAX_TOKENS = 2048
local ANTHROPIC_VERSION = "2023-06-01"

-- Spelling behaviour, selectable from the ✍️ menu (persisted across reloads).
-- "auto" keeps whatever variety the writer uses; the others force a dialect.
local DIALECTS = {
	{
		key = "auto",
		label = "Match my writing (auto)",
		clause = "Keep the writer's spelling variety consistent: infer from the text whether they write British or American English and normalise stray spellings to that variety; if there is no clear signal, leave regional spellings as written. Only fix genuine misspellings.",
	},
	{
		key = "british",
		label = "British English",
		clause = "Use British English spelling throughout (colour, organise, favour); correct American spellings to British.",
	},
	{
		key = "american",
		label = "American English",
		clause = "Use American English spelling throughout (color, organize, favor); correct British spellings to American.",
	},
}
local SETTING_DIALECT = "fettle.dialect" -- persisted across reloads/restarts

-- Poll settings for waiting on the async clipboard copy to land.
local POLL_INTERVAL = 0.02 -- seconds between checks
local POLL_MAX_TRIES = 25 -- ~0.5s total timeout
local RESTORE_DELAY = 0.35 -- seconds to wait before restoring the old clipboard

local function trim(s)
	return (s:gsub("^%s+", ""):gsub("%s+$", ""))
end

-- === providers ==============================================================
-- Each provider knows its endpoint, how to build a request, how to read the
-- reply, and where its key lives (always OUTSIDE the repo). If a provider ever
-- deprecates a model id, just update it in its `models` list below.
local PROVIDERS = {
	anthropic = {
		label = "Anthropic",
		url = "https://api.anthropic.com/v1/messages",
		keyFile = hs.configdir .. "/fettle-key.txt",
		envVar = "ANTHROPIC_API_KEY",
		models = {
			{ label = "Haiku 4.5  —  fast & cheap", id = "claude-haiku-4-5-20251001" },
			{ label = "Sonnet 5  —  balanced", id = "claude-sonnet-5" },
			{ label = "Opus 5  —  best, slower", id = "claude-opus-5" },
		},
		buildRequest = function(key, model, system, text)
			local headers = {
				["x-api-key"] = key,
				["anthropic-version"] = ANTHROPIC_VERSION,
				["content-type"] = "application/json",
			}
			local body = hs.json.encode({
				model = model,
				max_tokens = MAX_TOKENS,
				temperature = 0,
				system = system,
				messages = { { role = "user", content = text } },
			})
			return headers, body
		end,
		parseResponse = function(p)
			if type(p) == "table" and type(p.content) == "table" and p.content[1] then
				return p.content[1].text
			end
			return nil
		end,
	},
	groq = {
		label = "Groq (fast)",
		url = "https://api.groq.com/openai/v1/chat/completions",
		keyFile = hs.configdir .. "/fettle-groq-key.txt",
		envVar = "GROQ_API_KEY",
		models = {
			{ label = "GPT-OSS 20B  —  fastest", id = "openai/gpt-oss-20b" },
			{ label = "GPT-OSS 120B  —  higher quality", id = "openai/gpt-oss-120b" },
		},
		buildRequest = function(key, model, system, text)
			-- Groq is OpenAI-compatible.
			local headers = {
				["Authorization"] = "Bearer " .. key,
				["content-type"] = "application/json",
			}
			local body = hs.json.encode({
				model = model,
				max_tokens = MAX_TOKENS,
				temperature = 0,
				reasoning_effort = "low", -- gpt-oss are reasoning models; keep them snappy
				messages = {
					{ role = "system", content = system },
					{ role = "user", content = text },
				},
			})
			return headers, body
		end,
		parseResponse = function(p)
			if type(p) == "table" and type(p.choices) == "table" and p.choices[1] and p.choices[1].message then
				return p.choices[1].message.content
			end
			return nil
		end,
	},
}
-- Stable display order (PROVIDERS is a hash, so iteration order isn't defined).
local PROVIDER_ORDER = { "anthropic", "groq" }
-- =============================================================================

local function activeProvider()
	local id = hs.settings.get("fettle.provider") or "anthropic"
	if PROVIDERS[id] then
		return PROVIDERS[id], id
	end
	return PROVIDERS.anthropic, "anthropic"
end

local function activeModelId()
	local p = activeProvider()
	local stored = hs.settings.get("fettle.model")
	-- Only honour the stored model if it belongs to the active provider.
	for _, m in ipairs(p.models) do
		if m.id == stored then
			return stored
		end
	end
	return p.models[1].id
end

local function activeLabel()
	local p = activeProvider()
	local mid = activeModelId()
	for _, m in ipairs(p.models) do
		if m.id == mid then
			return p.label .. " · " .. m.label
		end
	end
	return p.label
end

local function getApiKey(provider)
	if provider.envVar then
		local envKey = os.getenv(provider.envVar)
		if envKey and envKey ~= "" then
			return trim(envKey)
		end
	end
	local f = io.open(provider.keyFile, "r")
	if not f then
		return nil
	end
	local contents = f:read("*a")
	f:close()
	local key = trim(contents or "")
	return key ~= "" and key or nil
end

local function currentDialect()
	local stored = hs.settings.get(SETTING_DIALECT)
	for _, d in ipairs(DIALECTS) do
		if d.key == stored then
			return d
		end
	end
	return DIALECTS[1] -- default: auto
end

-- Built fresh per correction so a menu change takes effect immediately; the
-- dialect clause is slotted in from the current menu-bar setting.
local function systemPrompt()
	return table.concat({
		"You fix text. Correct spelling, typos, grammar, and verb tense.",
		"Preserve the original meaning, tone, punctuation style, and line breaks.",
		currentDialect().clause,
		"Do not rephrase or add anything; make the smallest changes needed to make it correct.",
		"Return ONLY the corrected text — no quotes, no commentary, no explanation.",
		"If the text is already correct, return it unchanged.",
	}, " ")
end

-- === the correction engine ==================================================
-- Async so the network round-trip never blocks the UI. Calls done(fixedText).
-- On any error we fall back to the original text (no destructive change).
local function correctAsync(text, done)
	local provider = activeProvider()
	local key = getApiKey(provider)
	if not key then
		hs.alert.show("fettle: no " .. provider.label .. " key (add " .. provider.keyFile .. ")")
		done(text)
		return
	end

	local headers, body = provider.buildRequest(key, activeModelId(), systemPrompt(), text)

	hs.http.asyncPost(provider.url, body, headers, function(status, respBody, _)
		if status ~= 200 or not respBody then
			local detail = respBody and trim(respBody):sub(1, 120) or "no response"
			hs.alert.show("fettle: " .. provider.label .. " error " .. tostring(status) .. " — " .. detail)
			done(text)
			return
		end

		local ok, parsed = pcall(hs.json.decode, respBody)
		local fixed = ok and provider.parseResponse(parsed) or nil
		if not fixed then
			hs.alert.show("fettle: unexpected response")
			done(text)
			return
		end

		fixed = trim(fixed)
		done(fixed ~= "" and fixed or text)
	end)
end
-- =============================================================================

-- Wait for the pasteboard changeCount to move (i.e. our Cmd+C landed), with a
-- timeout so an empty/unsupported field can't hang us forever.
local function waitForCopy(startCount, onReady, triesLeft)
	triesLeft = triesLeft or POLL_MAX_TRIES
	if hs.pasteboard.changeCount() ~= startCount then
		onReady(true)
	elseif triesLeft <= 0 then
		onReady(false)
	else
		hs.timer.doAfter(POLL_INTERVAL, function()
			waitForCopy(startCount, onReady, triesLeft - 1)
		end)
	end
end

local function fixSelection()
	local original = hs.pasteboard.getContents()
	local startCount = hs.pasteboard.changeCount()

	-- Grab whatever's in the focused field.
	hs.eventtap.keyStroke({ "cmd" }, "a", 0)
	hs.eventtap.keyStroke({ "cmd" }, "c", 0)

	waitForCopy(startCount, function(copied)
		if not copied then
			hs.alert.show("fettle: nothing to grab")
			return
		end

		local text = hs.pasteboard.getContents()
		if not text or text == "" then
			hs.alert.show("fettle: empty selection")
			return
		end

		-- Immediate feedback so the network wait feels intentional, not laggy.
		local fixingId = hs.alert.show("✍️  fettling…", hs.alert.defaultStyle, hs.screen.mainScreen(), 10)

		correctAsync(text, function(fixed)
			hs.alert.closeSpecific(fixingId)
			hs.pasteboard.setContents(fixed)
			hs.eventtap.keyStroke({ "cmd" }, "v", 0)

			-- Give the paste a moment to consume the clipboard, then restore
			-- whatever was there before so we don't clobber it.
			hs.timer.doAfter(RESTORE_DELAY, function()
				if original then
					hs.pasteboard.setContents(original)
				end
			end)
		end)
	end)
end

-- === menu bar ===============================================================
local menubar = hs.menubar.new()

local function ensureKeyFileExists(keyFile)
	local f = io.open(keyFile, "r")
	if f then
		f:close()
	else
		local w = io.open(keyFile, "a")
		if w then
			w:close()
		end
	end
end

local function buildMenu()
	local _, activePid = activeProvider()
	local activeMid = activeModelId()

	local items = {
		{ title = "fettle", disabled = true },
		{ title = "-" },
	}

	-- One group of models per provider; clicking sets provider + model together.
	for _, pid in ipairs(PROVIDER_ORDER) do
		local p = PROVIDERS[pid]
		table.insert(items, { title = p.label, disabled = true })
		for _, m in ipairs(p.models) do
			table.insert(items, {
				title = "   " .. m.label,
				checked = (pid == activePid and m.id == activeMid),
				fn = function()
					hs.settings.set("fettle.provider", pid)
					hs.settings.set("fettle.model", m.id)
					hs.alert.show("fettle: " .. p.label .. " · " .. m.label)
				end,
			})
		end
	end

	table.insert(items, { title = "-" })
	table.insert(items, { title = "Spelling", disabled = true })
	local activeDialect = currentDialect().key
	for _, d in ipairs(DIALECTS) do
		table.insert(items, {
			title = "   " .. d.label,
			checked = (d.key == activeDialect),
			fn = function()
				hs.settings.set(SETTING_DIALECT, d.key)
				hs.alert.show("fettle: spelling → " .. d.label)
			end,
		})
	end

	table.insert(items, { title = "-" })
	for _, pid in ipairs(PROVIDER_ORDER) do
		local p = PROVIDERS[pid]
		table.insert(items, {
			title = "Edit " .. p.label .. " key…",
			fn = function()
				ensureKeyFileExists(p.keyFile)
				hs.execute("open -e '" .. p.keyFile .. "'")
			end,
		})
	end
	table.insert(items, { title = "Reload config", fn = function() hs.reload() end })
	return items
end

if menubar then
	menubar:setTitle("✍️")
	menubar:setMenu(buildMenu)
end
-- =============================================================================

hs.hotkey.bind(HOTKEY_MODS, HOTKEY_KEY, fixSelection)
hs.alert.show("fettle loaded — Option+Space (" .. activeLabel() .. ")")