palette/guess.go

package main

import (
	"encoding/json"
	"net/http"
	"regexp"
	"strings"

	"github.com/go-enry/go-enry/v2"
)

// hintRule is a lightweight regex hint that nudges detection. Hints don't
// decide on their own: matching hints are passed to enry's classifier as
// candidate languages, and enry (trained on Linguist samples) makes the final
// call. Adding a language later is usually a one-line addition here plus the
// dropdown in web/templates/new.html.
type hintRule struct {
	lang string // enry language name
	re   *regexp.Regexp
}

// hintRules are evaluated in order; keep more specific languages earlier so
// ties break in their favor.
var hintRules = []hintRule{
	{"Dockerfile", regexp.MustCompile(`(?mi)^(FROM\s+\S+:\S*|RUN\s+\S+|COPY\s+\S+\s+\S+|ENTRYPOINT\s+|WORKDIR\s+/)`)},
	{"Diff", regexp.MustCompile(`(?m)^(diff --git|--- a/|\+\+\+ b/|@@ -\d+)`)},
	{"PHP", regexp.MustCompile(`(?m)<\?php|\$\w+\s*=\s*[^=]|\becho\s+["'$]`)},
	{"HTML", regexp.MustCompile(`(?i)<(!DOCTYPE|html|head|body|div|span|script|p|a)\b`)},
	{"XML", regexp.MustCompile(`(?m)^<\?xml\b|<\/?[a-zA-Z][\w.-]*:[\w.-]*[>\s]`)},
	{"CSS", regexp.MustCompile(`(?m)(^|\})\s*[^{}@]+\{[^}]*:[^}]*\}|@(media|import|font-face)\b`)},
	{"TypeScript", regexp.MustCompile(`(?m)(:\s*(string|number|boolean|any)\b|\binterface \w+ \{|\btype \w+ =|\bimplements \w+)`)},
	{"TOML", regexp.MustCompile(`(?m)^\[[\w."-]+\]\s*$|^\w[\w-]*\s*=\s*("[^"]*"|\d+|true|false|\[[^\]]*\])\s*$`)},
	{"INI", regexp.MustCompile(`(?m)^\[[\w.-]+\]\s*$|^\w[\w.-]*\s*=\s*\S+\s*$`)},
	{"Ruby", regexp.MustCompile(`(?m)(\bdef \w+($|\s)|\brequire ['"]|\bputs \w+|@\w+\s*=\s*[^=]|\bend\b\s*$)`)},
	{"Perl", regexp.MustCompile(`(?m)(\buse strict\b|\bmy \$\w+|->\{|sub \w+ \{)`)},
	{"Lua", regexp.MustCompile(`(?m)(\bfunction\s+\w+\s*\(|\blocal \w+\s*=|\bthen\b|\belseif\b|\.\.\.)`)},
	{"Go", regexp.MustCompile(`(?m)^\s*(package \w+|import \(|func (\w+|\() )`)},
	{"Python", regexp.MustCompile(`(?m)^\s*(def \w+|class \w+|import \w+|from \w+ import |@\w+)`)},
	{"JavaScript", regexp.MustCompile(`(?m)(\bconst \w+ = |require\(|import \w+ from |=> \{|\bconsole\.log\()` )},
	{"Rust", regexp.MustCompile(`(?m)(\bfn \w+|let mut \b|\bimpl \b|use std::)`)},
	{"Java", regexp.MustCompile(`(?m)(\bpublic (static |final |class )|\bSystem\.out\.print|import java\.)`)},
	{"C", regexp.MustCompile(`(?m)(#include\s*<\w+\.h>|printf\(|\bint main\()` )},
	{"C++", regexp.MustCompile(`(?m)(#include\s*<(iostream|vector|string)>|std::|\bcout\s*<<)`)},
	{"SQL", regexp.MustCompile(`(?i)\b(SELECT .+ FROM|INSERT INTO|CREATE TABLE|UPDATE \w+ SET)\b`)},
	{"YAML", regexp.MustCompile(`(?m)^(\w[\w-]*:\s*(\||\S)|  \w[\w-]*: |---\s*$)`)},
	{"Markdown", regexp.MustCompile("(?m)^(#{1,6} \\S|\\|.*\\||-\\s\\[\\s?\\]|```)")},
	{"Shell", regexp.MustCompile(`(?m)^(#!.*bash|#!.*sh|\w+\(\)\s*\{)` )},
}

// canonical maps enry display names to the lowercase ids we store and render.
var canonical = map[string]string{
	"Dockerfile": "dockerfile",
	"C#":         "csharp",
	"Shell":      "bash",
}

// guessLang detects a language from pasted content. Order: fast decisive
// paths (empty, JSON, unambiguous markers enry can't see without a filename),
// then enry strategies (shebangs, XML decl, modelines, content heuristics),
// then enry's classifier seeded by our regex hints.
func guessLang(s string) string {
	src := strings.TrimSpace(s)
	if src == "" {
		return ""
	}

	// JSON: must start with { or [ and parse — cheaper and more decisive
	// than the classifier for pasted JSON, and handles compact single-line
	// JSON that content heuristics miss.
	if src[0] == '{' || src[0] == '[' {
		var v any
		if json.Unmarshal([]byte(src), &v) == nil {
			return "json"
		}
	}

	// enry's built-in strategies: shebangs, XML declaration, modelines,
	// content heuristics.
	if lang := enry.GetLanguage("", []byte(src)); lang != "" && lang != enry.OtherLanguage {
		return normalizeLang(lang)
	}

	// collect hint-matched languages as classifier candidates
	cands := []string{}
	for _, h := range hintRules {
		if h.re.MatchString(src) {
			cands = append(cands, h.lang)
		}
	}
	if len(cands) > 0 {
		// enry's Bayesian classifier (trained on Linguist samples) picks the
		// best of the hint candidates; fall back to the first hint if it
		// can't decide.
		if lang, _ := enry.GetLanguageByClassifier([]byte(src), cands); lang != "" {
			return normalizeLang(lang)
		}
		return normalizeLang(cands[0])
	}

	return "text"
}

// normalizeLang maps enry display names to our lowercase stored ids.
func normalizeLang(lang string) string {
	if c, ok := canonical[lang]; ok {
		return c
	}
	return strings.ToLower(lang)
}

func (a *apiServer) handleGuessLang(w http.ResponseWriter, r *http.Request) {
	setRateLimitHeaders(w, 1, 5)
	if !rateLimitGuess(r) {
		writeRateLimited(w, 1)
		return
	}
	var req struct {
		Content string `json:"content"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeErr(w, http.StatusBadRequest, "invalid json body")
		return
	}
	lang := guessLang(req.Content)
	writeJSON(w, http.StatusOK, map[string]any{"language": lang})
}