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 (
	"database/sql"
	"encoding/json"
	"errors"
	"net/http"
	"strings"

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

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

// Handler holds the dependencies for the checkpoint + suggestion routes.
type Handler struct {
	DB     *db.DB
	Client llm.LLMClient
	Limit  *llm.RateLimiter
}

// New constructs a Handler with a per-document checkpoint rate limiter.
func New(database *db.DB, client llm.LLMClient) *Handler {
	return &Handler{
		DB:     database,
		Client: client,
		Limit:  llm.NewRateLimiter(llm.CheckpointInterval),
	}
}

// RegisterDocRoutes adds the document-scoped routes (check + 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.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)
	return r
}

// check runs a grammar checkpoint over the document and returns the fresh
// pending suggestions. It is rate-limited per document (429 when too soon).
func (h *Handler) check(w http.ResponseWriter, r *http.Request) {
	docID := chi.URLParam(r, "id")

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

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

	if ok, _ := h.Limit.Allow(docID); !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 {
			serverError(w, err)
			return
		}
		writeJSON(w, http.StatusOK, existing)
		return
	}

	raw, err := llm.RunCheckpoint(r.Context(), h.Client, contentText)
	if err != nil {
		errorJSON(w, http.StatusBadGateway, "checkpoint failed: "+err.Error())
		return
	}

	saved, err := h.replacePending(docID, contentText, raw)
	if err != nil {
		serverError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, saved)
}

// replacePending swaps a document's pending suggestions for a fresh batch in one
// transaction. Accepted/rejected suggestions are left untouched (history).
func (h *Handler) replacePending(docID, contentText string, raw []llm.RawSuggestion) ([]db.Suggestion, error) {
	tx, err := h.DB.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

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

	out := []db.Suggestion{}
	for _, s := range raw {
		typ := normalizeType(s.Type)
		from, to := locate(contentText, s.Original)
		var saved db.Suggestion
		err := tx.QueryRow(
			`INSERT INTO suggestions (doc_id, from_pos, to_pos, original, replacement, explanation, type)
			 VALUES (?, ?, ?, ?, ?, ?, ?)
			 RETURNING id, doc_id, from_pos, to_pos, original, replacement, explanation, type, status, created_at`,
			docID, from, to, s.Original, s.Replacement, s.Explanation, typ,
		).Scan(
			&saved.ID, &saved.DocID, &saved.FromPos, &saved.ToPos, &saved.Original,
			&saved.Replacement, &saved.Explanation, &saved.Type, &saved.Status, &saved.CreatedAt,
		)
		if err != nil {
			return nil, err
		}
		out = append(out, saved)
	}

	if err := tx.Commit(); err != nil {
		return nil, err
	}
	return out, nil
}

// 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 {
		serverError(w, err)
		return
	}
	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)
	}
	return out, rows.Err()
}

// 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 {
		serverError(w, err)
		return
	}
	if n, _ := res.RowsAffected(); n == 0 {
		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:
		return strings.ToLower(strings.TrimSpace(t))
	default:
		return db.SuggestionTypeGrammar
	}
}

// --- response helpers -------------------------------------------------------

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func errorJSON(w http.ResponseWriter, status int, msg string) {
	writeJSON(w, status, map[string]string{"error": msg})
}

func serverError(w http.ResponseWriter, err error) {
	errorJSON(w, http.StatusInternalServerError, err.Error())
}