petal/internal/suggestions/handlers.go

// Package suggestions implements the grammar-checkpoint endpoint and the
// accept/dismiss surface for LLM-proposed edits. Checkpoints run a single LLM
// pass over a document and persist the resulting pending suggestions; the
// frontend re-anchors each one by its `original` string at render time, so the
// stored positions are advisory only (spec Note #6).
package suggestions

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"

	"github.com/go-chi/chi/v5"

	"gitea.parodia.dev/drwily/petal/internal/db"
	"gitea.parodia.dev/drwily/petal/internal/httputil"
	"gitea.parodia.dev/drwily/petal/internal/llm"
)

// Handler holds the dependencies for the checkpoint + suggestion routes. The
// grammar checkpoint and the voice pass each get their own per-document rate
// limiter — they are independent passes with different cadences.
type Handler struct {
	DB               *db.DB
	Client           llm.LLMClient
	Limit            *llm.RateLimiter // grammar checkpoint floor
	VoiceLimit       *llm.RateLimiter // voice-consistency floor
	CollocationLimit *llm.RateLimiter // collocation-coach floor
}

// New constructs a Handler with per-document checkpoint, voice, and collocation
// rate limiters.
func New(database *db.DB, client llm.LLMClient) *Handler {
	return &Handler{
		DB:               database,
		Client:           client,
		Limit:            llm.NewRateLimiter(llm.CheckpointInterval),
		VoiceLimit:       llm.NewRateLimiter(llm.VoiceInterval),
		CollocationLimit: llm.NewRateLimiter(llm.CollocationInterval),
	}
}

// RegisterDocRoutes adds the document-scoped routes (check + voice + list) onto
// the existing /api/docs router so they share its base path.
func (h *Handler) RegisterDocRoutes(r chi.Router) {
	r.Post("/{id}/check", h.check)
	r.Post("/{id}/mechanics", h.mechanics)
	r.Post("/{id}/voice", h.voice)
	r.Post("/{id}/collocation", h.collocation)
	r.Post("/{id}/rewrite", h.rewrite)
	r.Get("/{id}/suggestions", h.listForDoc)
}

// Routes returns the router mounted at /api/suggestions for per-suggestion
// actions.
func (h *Handler) Routes() chi.Router {
	r := chi.NewRouter()
	r.Post("/{id}/accept", h.accept)
	r.Post("/{id}/dismiss", h.dismiss)
	r.Post("/{id}/chat", h.chat)
	r.Post("/{id}/translate", h.translate)
	return r
}

// check runs a grammar checkpoint over the document. Fast, typing-cadence pass.
// The deterministic mechanics family is owned by a separate pass (see mechanics),
// detected client-side; a grammar checkpoint leaves those flags untouched.
func (h *Handler) check(w http.ResponseWriter, r *http.Request) {
	h.runPass(w, r, h.Limit, llm.RunCheckpoint, grammarScope)
}

// mechanicsFinding is one deterministic, rule-based fix detected client-side (see
// web Companion/prose.ts). Detection lives in the frontend — the same rules that
// power the companion's prose notes — so the server only persists these; it does
// not compute them. Offsets are exact plaintext spans from the detector.
type mechanicsFinding struct {
	From        int    `json:"from"`
	To          int    `json:"to"`
	Original    string `json:"original"`
	Replacement string `json:"replacement"`
	Explanation string `json:"explanation"`
}

// maxMechanicsFindings caps a single submission so a runaway client can't flood
// the table; far above any realistic count for one document.
const maxMechanicsFindings = 500

// mechanics persists the client-detected deterministic fixes as the 'mechanics'
// family and returns the document's unified pending set. Free and not rate-
// limited — it runs alongside the grammar checkpoint. It mirrors a single LLM
// family: it replaces only the pending mechanics rows, honours actioned-
// suppression, and (unlike the LLM passes) keeps the detector's exact offsets
// rather than re-locating by string, which matters when the same word repeats.
func (h *Handler) mechanics(w http.ResponseWriter, r *http.Request) {
	docID := chi.URLParam(r, "id")

	// Confirm the document exists (and is the local user's) for clean 404s.
	var exists bool
	err := h.DB.QueryRow(
		`SELECT EXISTS(SELECT 1 FROM documents WHERE id = ? AND user_id = ?)`,
		docID, db.LocalUserID,
	).Scan(&exists)
	if err != nil {
		httputil.ServerError(w, err)
		return
	}
	if !exists {
		httputil.ErrorJSON(w, http.StatusNotFound, "document not found")
		return
	}

	var body struct {
		Findings []mechanicsFinding `json:"findings"`
	}
	if err := json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&body); err != nil {
		httputil.BadRequest(w, "invalid request body")
		return
	}
	if len(body.Findings) > maxMechanicsFindings {
		body.Findings = body.Findings[:maxMechanicsFindings]
	}

	if err := h.replaceMechanics(docID, body.Findings); err != nil {
		httputil.ServerError(w, err)
		return
	}

	out, err := h.fetchPending(docID)
	if err != nil {
		httputil.ServerError(w, err)
		return
	}
	httputil.WriteJSON(w, http.StatusOK, out)
}

// replaceMechanics swaps the document's pending mechanics rows for the supplied
// findings in one transaction, leaving the LLM families and actioned rows
// untouched. Findings the user already accepted or dismissed are suppressed (the
// detector has no memory between runs), and malformed spans are skipped.
func (h *Handler) replaceMechanics(docID string, findings []mechanicsFinding) error {
	tx, err := h.DB.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	if _, err := tx.Exec(
		`DELETE FROM suggestions WHERE doc_id = ? AND status = ? AND type = ?`,
		docID, db.SuggestionStatusPending, db.SuggestionTypeMechanics,
	); err != nil {
		return err
	}

	actioned, err := actionedKeys(tx, docID)
	if err != nil {
		return err
	}

	for _, f := range findings {
		if f.From < 0 || f.To <= f.From || strings.TrimSpace(f.Original) == "" {
			continue // malformed span — the client re-anchors by string anyway
		}
		if _, seen := actioned[suggestionKey(f.Original, f.Replacement)]; seen {
			continue
		}
		if _, err := tx.Exec(
			`INSERT INTO suggestions (doc_id, from_pos, to_pos, original, replacement, explanation, type)
			 VALUES (?, ?, ?, ?, ?, ?, ?)`,
			docID, f.From, f.To, f.Original, f.Replacement, f.Explanation, db.SuggestionTypeMechanics,
		); err != nil {
			return err
		}
	}

	return tx.Commit()
}

// voice runs a Tier-1 voice-consistency pass over the whole document. Slow,
// explicit-action pass; replaces only the pending voice flags.
func (h *Handler) voice(w http.ResponseWriter, r *http.Request) {
	h.runPass(w, r, h.VoiceLimit, llm.RunVoice, voiceScope)
}

// collocation runs the collocation coach over the whole document, flagging
// non-native word pairings. Explicit-action pass; replaces only the pending
// collocation flags.
func (h *Handler) collocation(w http.ResponseWriter, r *http.Request) {
	h.runPass(w, r, h.CollocationLimit, llm.RunCollocation, collocationScope)
}

// pass is the signature shared by the grammar checkpoint and the voice pass:
// given the document text and the document's tone it returns the model's raw
// suggestions. The voice pass ignores tone (see llm.RunVoice).
type pass func(ctx context.Context, client llm.LLMClient, contentText, tone string) ([]llm.RawSuggestion, error)

// runPass is the shared body for both LLM passes. It loads the document text,
// enforces the pass's per-document rate limit, runs the model, swaps in the
// fresh batch scoped to this family, and returns the document's FULL pending set
// (both families) so the client always renders a unified picture.
func (h *Handler) runPass(w http.ResponseWriter, r *http.Request, limiter *llm.RateLimiter, run pass, scope pendingScope) {
	docID := chi.URLParam(r, "id")

	var contentText, tone string
	err := h.DB.QueryRow(
		`SELECT content_text, tone FROM documents WHERE id = ? AND user_id = ?`,
		docID, db.LocalUserID,
	).Scan(&contentText, &tone)
	if errors.Is(err, sql.ErrNoRows) {
		httputil.ErrorJSON(w, http.StatusNotFound, "document not found")
		return
	}
	if err != nil {
		httputil.ServerError(w, err)
		return
	}

	// Nothing to analyze on an empty document — skip the LLM round-trip.
	if strings.TrimSpace(contentText) == "" {
		httputil.WriteJSON(w, http.StatusOK, []db.Suggestion{})
		return
	}

	ok, _, slotAt := limiter.Allow(docID)
	if !ok {
		// Throttled: return the existing pending set unchanged rather than an
		// error, so the frontend keeps showing current suggestions.
		existing, err := h.fetchPending(docID)
		if err != nil {
			httputil.ServerError(w, err)
			return
		}
		httputil.WriteJSON(w, http.StatusOK, existing)
		return
	}

	raw, err := run(r.Context(), h.Client, contentText, tone)
	if err != nil {
		// Allow ran before the model call, so a failed pass would otherwise hold
		// the per-document slot for the full interval — stranding the frontend's
		// auto-retry on the throttle path. Release it so a retry can re-run.
		limiter.Release(docID, slotAt)
		httputil.ErrorJSON(w, http.StatusBadGateway, "llm pass failed: "+err.Error())
		return
	}

	if err := h.replacePending(docID, contentText, raw, scope); err != nil {
		httputil.ServerError(w, err)
		return
	}

	// Return the unified pending set (grammar + voice), not just this batch, so
	// a grammar check never drops the voice highlights from the client and the
	// throttle path above stays consistent with the success path.
	out, err := h.fetchPending(docID)
	if err != nil {
		httputil.ServerError(w, err)
		return
	}
	httputil.WriteJSON(w, http.StatusOK, out)
}

// pendingScope describes how one LLM pass touches the shared suggestions table:
// which family of pending rows it replaces, and the type to stamp on the rows it
// inserts. The grammar checkpoint and voice pass each own a disjoint family, so
// running one never disturbs the other's pending flags.
type pendingScope struct {
	deleteWhere string // extra WHERE clause scoping the DELETE to this family
	forceType   string // if set, every inserted row gets this type; else normalizeType
}

var (
	// grammarScope owns the grammar/phrasing/idiom/clarity flags — everything but
	// the other self-owned families (voice, collocation, mechanics), which run on
	// their own cadence/pass and must survive a grammar checkpoint. Notably the
	// deterministic mechanics pass writes its rows in the same /check request just
	// before this DELETE runs, so excluding it here is what keeps them alive.
	grammarScope = pendingScope{deleteWhere: "type NOT IN ('voice','collocation','mechanics')", forceType: ""}
	// voiceScope owns the voice flags only.
	voiceScope = pendingScope{deleteWhere: "type = 'voice'", forceType: db.SuggestionTypeVoice}
	// collocationScope owns the collocation flags only.
	collocationScope = pendingScope{deleteWhere: "type = 'collocation'", forceType: db.SuggestionTypeCollocation}
)

// replacePending swaps a document's pending suggestions within one family for a
// fresh batch in a single transaction. Accepted/rejected suggestions and the
// other family's pending rows are left untouched.
//
// Suggestions the user already accepted or dismissed are suppressed from the
// fresh batch: the model has no memory between passes, so without this it would
// re-propose the identical edit on the very next checkpoint — re-nagging a
// sentence the user already resolved. This matters most when an accept silently
// no-ops (the `original` text couldn't be anchored in the editor, so the doc
// never changed): the sentence is unaltered, yet the user shouldn't see the same
// card again.
func (h *Handler) replacePending(docID, contentText string, raw []llm.RawSuggestion, scope pendingScope) error {
	tx, err := h.DB.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	if _, err := tx.Exec(
		`DELETE FROM suggestions WHERE doc_id = ? AND status = ? AND `+scope.deleteWhere,
		docID, db.SuggestionStatusPending,
	); err != nil {
		return err
	}

	actioned, err := actionedKeys(tx, docID)
	if err != nil {
		return err
	}

	for _, s := range raw {
		if _, seen := actioned[suggestionKey(s.Original, s.Replacement)]; seen {
			continue
		}
		typ := scope.forceType
		if typ == "" {
			typ = normalizeType(s.Type)
		}
		from, to := locate(contentText, s.Original)
		if _, err := tx.Exec(
			`INSERT INTO suggestions (doc_id, from_pos, to_pos, original, replacement, explanation, type)
			 VALUES (?, ?, ?, ?, ?, ?, ?)`,
			docID, from, to, s.Original, s.Replacement, s.Explanation, typ,
		); err != nil {
			return err
		}
	}

	return tx.Commit()
}

// actionedKeys returns the set of original→replacement keys the user has already
// accepted or rejected for this document, so a fresh checkpoint can skip
// re-proposing them. Keyed on the trimmed original+replacement pair, matching the
// normalization ParseCheckpoint applies, so a genuinely different edit on the same
// sentence is not suppressed.
func actionedKeys(tx *sql.Tx, docID string) (map[string]struct{}, error) {
	rows, err := tx.Query(
		`SELECT original, replacement FROM suggestions
		  WHERE doc_id = ? AND status IN (?, ?)`,
		docID, db.SuggestionStatusAccepted, db.SuggestionStatusRejected,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	keys := make(map[string]struct{})
	for rows.Next() {
		var original, replacement string
		if err := rows.Scan(&original, &replacement); err != nil {
			return nil, err
		}
		keys[suggestionKey(original, replacement)] = struct{}{}
	}
	return keys, rows.Err()
}

// suggestionKey is the dedup identity for a suggestion: its trimmed original and
// replacement text. Two suggestions with the same key are "the same edit."
func suggestionKey(original, replacement string) string {
	return strings.TrimSpace(original) + "\x00" + strings.TrimSpace(replacement)
}

// listForDoc returns the document's current pending suggestions (used when the
// editor loads a document, before any new checkpoint fires).
func (h *Handler) listForDoc(w http.ResponseWriter, r *http.Request) {
	out, err := h.fetchPending(chi.URLParam(r, "id"))
	if err != nil {
		httputil.ServerError(w, err)
		return
	}
	httputil.WriteJSON(w, http.StatusOK, out)
}

func (h *Handler) fetchPending(docID string) ([]db.Suggestion, error) {
	rows, err := h.DB.Query(
		`SELECT id, doc_id, from_pos, to_pos, original, replacement, explanation, type, status, created_at
		   FROM suggestions
		  WHERE doc_id = ? AND status = ?
		  ORDER BY from_pos ASC, created_at ASC`,
		docID, db.SuggestionStatusPending,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := []db.Suggestion{}
	for rows.Next() {
		var s db.Suggestion
		if err := rows.Scan(
			&s.ID, &s.DocID, &s.FromPos, &s.ToPos, &s.Original, &s.Replacement,
			&s.Explanation, &s.Type, &s.Status, &s.CreatedAt,
		); err != nil {
			return nil, err
		}
		out = append(out, s)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return dedupeSpans(out), nil
}

// dedupeSpans resolves collisions between the deterministic mechanics family and
// the LLM families: when a mechanics finding and an LLM suggestion fight over the
// same characters, mechanics wins and the LLM card is dropped. Its span is exact
// (the detector matched it), whereas the LLM positions are only advisory
// (re-anchored by string at render), so the precise fix should own the span.
//
// This deliberately does NOT dedupe LLM-vs-LLM overlaps: voice (awareness-only,
// no replacement) and collocation legitimately co-occupy the same span, and that
// is intended. Suggestions that never anchored (from_pos < 0) occupy no real span
// and are always kept.
func dedupeSpans(in []db.Suggestion) []db.Suggestion {
	type span struct{ from, to int }
	var claimed []span
	for _, s := range in {
		if s.Type == db.SuggestionTypeMechanics && s.FromPos >= 0 {
			claimed = append(claimed, span{s.FromPos, s.ToPos})
		}
	}
	if len(claimed) == 0 {
		return in
	}

	out := make([]db.Suggestion, 0, len(in))
	for _, s := range in {
		if s.Type != db.SuggestionTypeMechanics && s.FromPos >= 0 {
			overlaps := false
			for _, sp := range claimed {
				if s.FromPos < sp.to && sp.from < s.ToPos {
					overlaps = true
					break
				}
			}
			if overlaps {
				continue // an exact mechanics fix owns these characters
			}
		}
		out = append(out, s)
	}
	return out
}

// accept marks a suggestion accepted (the client applies the replacement text).
func (h *Handler) accept(w http.ResponseWriter, r *http.Request) {
	h.setStatus(w, r, db.SuggestionStatusAccepted)
}

// dismiss marks a suggestion rejected.
func (h *Handler) dismiss(w http.ResponseWriter, r *http.Request) {
	h.setStatus(w, r, db.SuggestionStatusRejected)
}

func (h *Handler) setStatus(w http.ResponseWriter, r *http.Request, status string) {
	res, err := h.DB.Exec(
		`UPDATE suggestions SET status = ? WHERE id = ? AND status = ?`,
		status, chi.URLParam(r, "id"), db.SuggestionStatusPending,
	)
	if err != nil {
		httputil.ServerError(w, err)
		return
	}
	if n, _ := res.RowsAffected(); n == 0 {
		httputil.ErrorJSON(w, http.StatusNotFound, "pending suggestion not found")
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

// locate finds the plaintext offsets of original within contentText. Returns
// (-1, -1) when not found; the frontend anchors by string regardless, so a miss
// here is non-fatal.
func locate(contentText, original string) (int, int) {
	idx := strings.Index(contentText, original)
	if idx < 0 {
		return -1, -1
	}
	return idx, idx + len(original)
}

// normalizeType maps the model's type string onto a valid suggestion type,
// defaulting unknown values to grammar so a stray label never trips the CHECK.
func normalizeType(t string) string {
	switch strings.ToLower(strings.TrimSpace(t)) {
	case db.SuggestionTypeGrammar, db.SuggestionTypePhrasing, db.SuggestionTypeIdiom, db.SuggestionTypeClarity, db.SuggestionTypeCollocation:
		return strings.ToLower(strings.TrimSpace(t))
	default:
		return db.SuggestionTypeGrammar
	}
}