sift/internal/api/api.go

package api

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/auth"
	"git.wntrmute.dev/kyle/sift/internal/clean"
	"git.wntrmute.dev/kyle/sift/internal/digest"
	"git.wntrmute.dev/kyle/sift/internal/fetch"
	"git.wntrmute.dev/kyle/sift/internal/fleetctx"
	"git.wntrmute.dev/kyle/sift/internal/guard"
	"git.wntrmute.dev/kyle/sift/internal/recover"
	"git.wntrmute.dev/kyle/sift/internal/selfip"
)

const maxReqBody = 64 << 10

type Server struct {
	Version string
	Log     *slog.Logger
	Guard   *guard.Guard
	Fetch   *fetch.Fetcher
	Auth    *auth.Store
	Mux     *http.ServeMux
	OutCap  int
	SelfIP  *selfip.Discoverer
	Digest  *digest.Engine
	Ctx     *fleetctx.Client

	global chan struct{}
	perMu  sync.Mutex
	per    map[string]chan struct{}

	reqN, blockN, scrapeN atomic.Uint64
}

func New(version string, g *guard.Guard, f *fetch.Fetcher, keys *auth.Store, log *slog.Logger) *Server {
	if log == nil {
		log = slog.Default()
	}
	s := &Server{
		Version: version,
		Log:     log,
		Guard:   g,
		Fetch:   f,
		Auth:    keys,
		Mux:     http.NewServeMux(),
		OutCap:  200_000,
		global:  make(chan struct{}, 8),
		per:     map[string]chan struct{}{},
	}
	s.Mux.HandleFunc("GET /healthz", s.healthz)
	s.Mux.HandleFunc("GET /healthz/", s.healthz)
	s.Mux.HandleFunc("GET /metrics", s.metrics)
	s.Mux.HandleFunc("POST /v2/scrape", s.scrape)
	s.Mux.HandleFunc("POST /v1/fetch", s.v1fetch)
	return s
}

func (s *Server) HTTPServer(addr string) *http.Server {
	return &http.Server{
		Addr:              addr,
		Handler:           s,
		ReadHeaderTimeout: 5 * time.Second,
		ReadTimeout:       15 * time.Second,
		WriteTimeout:      30 * time.Second,
		IdleTimeout:       30 * time.Second,
		MaxHeaderBytes:    16 << 10,
	}
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, maxReqBody)
	s.Mux.ServeHTTP(w, r)
}

func (s *Server) healthz(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	out := map[string]any{"status": "ok", "version": s.Version}
	if s.SelfIP != nil {
		out["self_ip"] = s.SelfIP.Status()
	}
	_ = json.NewEncoder(w).Encode(out)
}

func (s *Server) metrics(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	fails := uint64(0)
	if s.SelfIP != nil {
		fails = s.SelfIP.Failures.Load()
	}
	_ = json.NewEncoder(w).Encode(map[string]any{
		"requests":         s.reqN.Load(),
		"blocks":           s.blockN.Load(),
		"scrapes":          s.scrapeN.Load(),
		"self_ip_failures": fails,
	})
}

type scrapeReq struct {
	URL                 string   `json:"url"`
	Formats             []string `json:"formats"`
	OnlyMainContent     *bool    `json:"onlyMainContent"`
	SkipTlsVerification *bool    `json:"skipTlsVerification"`
	RemoveBase64Images  *bool    `json:"removeBase64Images"`
	FastMode            *bool    `json:"fastMode"`
	BlockAds            *bool    `json:"blockAds"`
	StoreInCache        *bool    `json:"storeInCache"`
	MaxAge              *int     `json:"maxAge"`
	Mobile              *bool    `json:"mobile"`
	Origin              string   `json:"origin"`
}

type v1Req struct {
	URL      string `json:"url"`
	Question string `json:"question"`
	Mode     string `json:"mode"`
	Context  *bool  `json:"context"`
	JQ       string `json:"jq"`
	Recover  bool   `json:"recover"`
}

func (s *Server) caller(w http.ResponseWriter, r *http.Request) *auth.Caller {
	h := r.Header.Get("Authorization")
	tok := strings.TrimSpace(strings.TrimPrefix(h, "Bearer"))
	tok = strings.TrimSpace(tok)
	if tok == "" || s.Auth == nil {
		s.Log.Warn("unauthorized", "caller", "-", "host", "-", "reason", "missing key")
		writeErr(w, 401, "unauthorized", "")
		return nil
	}
	c := s.Auth.LookupBearer(tok)
	if c == nil {
		s.Log.Warn("unauthorized", "caller", "-", "host", "-", "reason", "unknown key")
		writeErr(w, 401, "unauthorized", "")
		return nil
	}
	return c
}

func (s *Server) scrape(w http.ResponseWriter, r *http.Request) {
	s.reqN.Add(1)
	c := s.caller(w, r)
	if c == nil {
		return
	}
	var req scrapeReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.URL == "" {
		writeErr(w, 400, "invalid url", "")
		return
	}
	_ = req.SkipTlsVerification // MUST ignore; never disable TLS verify
	onlyMain := true
	if req.OnlyMainContent != nil {
		onlyMain = *req.OnlyMainContent
	}
	wantMD, wantHTML := true, false
	if len(req.Formats) > 0 {
		wantMD, wantHTML = false, false
		for _, f := range req.Formats {
			switch strings.ToLower(f) {
			case "markdown":
				wantMD = true
			case "html":
				wantHTML = true
			}
		}
		if !wantMD && !wantHTML {
			wantMD = true
		}
	}
	s.scrapeN.Add(1)
	res, cl, ferr := s.doFetch(r, c, req.URL, "", onlyMain)
	if ferr != nil {
		s.writeFetchErr(w, c, req.URL, ferr)
		return
	}
	md, htmlOut := cl.Markdown, ""
	if wantHTML {
		htmlOut = cl.HTML
	}
	if !wantMD {
		md = ""
	}
	rawTok := len(res.Body) / 4
	retTok := len(md) / 4
	s.Log.Info("scrape", "caller", c.Name, "host", hostOf(res.FinalURL), "status", res.Status, "bytes", len(res.Body))
	writeJSON(w, 200, map[string]any{
		"success": true,
		"data": map[string]any{
			"markdown": md,
			"html":     htmlOut,
			"metadata": map[string]any{
				"title":       cl.Title,
				"sourceURL":   res.FinalURL,
				"statusCode":  res.Status,
				"contentType": res.ContentType,
				"sift": map[string]any{
					"hash":      res.SHA256(),
					"fetchedAt": res.FetchedAt.Format(time.RFC3339),
					"cached":    false,
					"truncated": cl.Trunc,
					"tokens":    map[string]int{"raw": rawTok, "returned": retTok},
				},
			},
		},
	})
}

func (s *Server) v1fetch(w http.ResponseWriter, r *http.Request) {
	s.reqN.Add(1)
	c := s.caller(w, r)
	if c == nil {
		return
	}
	var req v1Req
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.URL == "" {
		writeErr(w, 400, "invalid url", "")
		return
	}
	q := req.Question
	if len([]rune(q)) > 500 {
		rns := []rune(q)
		q = string(rns[:500])
	}
	mode := req.Mode
	if mode == "" {
		if q != "" {
			mode = "digest"
		} else {
			mode = "clean"
		}
	}
	if mode != "clean" && mode != "raw" && mode != "digest" {
		writeErr(w, 400, "invalid url", "")
		return
	}
	wantCtx := q != ""
	if req.Context != nil {
		wantCtx = *req.Context
	}

	t0 := time.Now()
	onlyMain := true
	res, cl, ferr := s.doFetch(r, c, req.URL, req.JQ, onlyMain)
	var snap *recover.Result
	if ferr != nil && req.Recover && recoverWorthy(ferr) {
		sr, rerr := recover.Try(r.Context(), s.Fetch, req.URL)
		if rerr == nil && sr != nil && sr.Res != nil {
			res, cl, ferr = s.cleanFetched(sr.Res, req.JQ, onlyMain)
			snap = sr
		}
	}
	if ferr != nil {
		s.writeFetchErr(w, c, req.URL, ferr)
		return
	}
	fetchMS := time.Since(t0).Milliseconds()

	content := cl.Markdown
	if cl.Kind == "json" {
		content = cl.HTML
	}
	if mode == "raw" {
		content = string(res.Body)
		if len(content) > s.OutCap {
			content = clean.TruncateUTF8(content, s.OutCap)
			cl.Trunc = true
			cl.Warning = "truncated to output cap"
		}
	}
	warn := []string{}
	if cl.Warning != "" {
		warn = append(warn, cl.Warning)
	}

	var digMeta map[string]any
	var digestMS int64
	if mode == "digest" {
		if s.Digest == nil {
			warn = append(warn, "digest fell back to clean: unset")
			mode = "clean"
			digMeta = map[string]any{"fallback": "unset"}
		} else {
			dr := s.Digest.Run(r.Context(), content, q, res.FinalURL)
			digestMS = dr.DigestMS
			if dr.Fallback != "" {
				warn = append(warn, "digest fell back to clean: "+dr.Fallback)
				mode = "clean"
				content = clean.TruncateUTF8(content, s.OutCap)
				digMeta = map[string]any{"fallback": dr.Fallback, "model": dr.Model}
			} else {
				content = dr.Text
				digMeta = map[string]any{
					"model":        dr.Model,
					"chunks_used":  dr.ChunksUsed,
					"chunks_total": dr.ChunksTotal,
					"fallback":     false,
					"cached":       dr.Cached,
				}
			}
		}
	}

	var ctxItems []fleetctx.Item
	var ctxMS int64
	if wantCtx && s.Ctx != nil {
		ct0 := time.Now()
		query := fleetctx.Query(q, cl.Title, headings(cl.Markdown))
		items, cw := s.Ctx.Search(r.Context(), c.Workspace, query)
		ctxMS = time.Since(ct0).Milliseconds()
		if cw != "" {
			warn = append(warn, cw)
		}
		ctxItems = items
	}

	qLen := len(q)
	if qLen > 80 {
		qLen = 80
	}
	fb, _ := "", false
	if digMeta != nil {
		if s, ok := digMeta["fallback"].(string); ok {
			fb = s
		}
	}
	s.Log.Info("fetch", "caller", c.Name, "host", hostOf(res.FinalURL), "mode", mode,
		"q_len", qLen, "raw_tokens", len(res.Body)/4, "returned_tokens", len(content)/4,
		"digest_ms", digestMS, "fallback", fb)

	src := map[string]any{
		"url":          res.URL,
		"final_url":    res.FinalURL,
		"status":       res.Status,
		"content_type": res.ContentType,
		"fetched_at":   res.FetchedAt.Format(time.RFC3339),
		"hash":         res.SHA256(),
		"bytes":        len(res.Body),
		"cached":       false,
		"kind":         "live",
	}
	if snap != nil {
		src["kind"] = "snapshot"
		src["via"] = string(snap.Via)
		src["original_url"] = snap.OriginalURL
		if snap.SnapshotAt != "" {
			src["snapshot_at"] = snap.SnapshotAt
		} else {
			src["snapshot_at"] = nil
			warn = append(warn, "snapshot date unknown")
		}
	}
	out := map[string]any{
		"content":   content,
		"mode":      mode,
		"untrusted": true,
		"source":    src,
		"warnings":  warn,
		"sizes": map[string]int{
			"raw_tokens":      len(res.Body) / 4,
			"clean_tokens":    digest.EstTokens(cl.Markdown),
			"returned_tokens": digest.EstTokens(content),
		},
		"timings": map[string]int64{"fetch_ms": fetchMS, "digest_ms": digestMS, "context_ms": ctxMS},
	}
	if digMeta != nil {
		out["digest"] = digMeta
	}
	if wantCtx {
		out["context"] = ctxItems
	}
	writeJSON(w, 200, out)
}

func recoverWorthy(err error) bool {
	var fe *fetch.Error
	if !errors.As(err, &fe) {
		return false
	}
	if strings.Contains(fe.Msg, "blocked:") {
		return false
	}
	return fe.Hint == "browser" || fe.Status == 422 || fe.Status == 429 || fe.Status == 504
}

func headings(md string) []string {
	var h []string
	for _, line := range strings.Split(md, "\n") {
		if strings.HasPrefix(line, "#") {
			h = append(h, strings.TrimLeft(line, "# "))
		}
	}
	return h
}

func (s *Server) cleanFetched(res *fetch.Resource, jq string, onlyMain bool) (*fetch.Resource, *clean.Result, error) {
	kind := clean.MediaKind(res.ContentType)
	if kind == "sniff" {
		kind = clean.Sniff(res.Body)
	}
	if kind == "other" {
		return nil, nil, &fetch.Error{Status: 415, Hint: "browser", Msg: "unsupported content type"}
	}
	var cl *clean.Result
	var cerr error
	switch kind {
	case "html":
		cl, cerr = clean.HTML(res, onlyMain, s.OutCap)
	case "json":
		cl, cerr = clean.JSON(res, jq, s.OutCap)
	default:
		cl, cerr = clean.Text(res, s.OutCap)
	}
	if cerr != nil {
		return nil, nil, &fetch.Error{Status: 422, Msg: cerr.Error()}
	}
	return res, cl, nil
}

func (s *Server) doFetch(r *http.Request, c *auth.Caller, rawURL, jq string, onlyMain bool) (*fetch.Resource, *clean.Result, error) {
	if err := s.acquire(c.Name, rawURL); err != nil {
		return nil, nil, err
	}
	defer s.release(c.Name, rawURL)

	res, err := s.Fetch.Fetch(r.Context(), rawURL)
	if err != nil {
		return nil, nil, err
	}
	kind := clean.MediaKind(res.ContentType)
	if kind == "sniff" {
		kind = clean.Sniff(res.Body)
	}
	if kind == "other" {
		ct := res.ContentType
		if ct == "" {
			ct = "unknown"
		}
		return nil, nil, &fetch.Error{Status: 415, Hint: "browser", Msg: "unsupported content type " + ct + "; hint: browser"}
	}
	var cl *clean.Result
	var cerr error
	switch kind {
	case "html":
		cl, cerr = clean.HTML(res, onlyMain, s.OutCap)
	case "json":
		cl, cerr = clean.JSON(res, jq, s.OutCap)
	default:
		cl, cerr = clean.Text(res, s.OutCap)
	}
	if cerr != nil {
		var tl *clean.TooLargeError
		if errors.As(cerr, &tl) {
			return nil, nil, &fetch.Error{Status: 400, Msg: tl.Error()}
		}
		return nil, nil, &fetch.Error{Status: 422, Msg: cerr.Error()}
	}
	return res, cl, nil
}

func (s *Server) writeFetchErr(w http.ResponseWriter, c *auth.Caller, rawURL string, err error) {
	name := "-"
	if c != nil {
		name = c.Name
	}
	host := hostOf(rawURL)
	var fe *fetch.Error
	if errors.As(err, &fe) {
		if fe.Status == 403 {
			s.blockN.Add(1)
		}
		if fe.Status == 502 {
			fe.Status = 500
			fe.Msg = "internal error"
		}
		if fe.Status == 401 || fe.Status == 403 {
			s.Log.Warn("refused", "caller", name, "host", host, "reason", fe.Msg)
		} else {
			s.Log.Warn("fetch failed", "caller", name, "host", host, "reason", fe.Msg, "status", fe.Status)
		}
		writeErr(w, fe.Status, fe.Msg, fe.Hint)
		return
	}
	s.Log.Error("internal", "caller", name, "host", host, "err", err)
	writeErr(w, 500, "internal error", "")
}

func (s *Server) acquire(caller, rawURL string) error {
	host := hostOf(rawURL)
	key := caller + "|" + host
	select {
	case s.global <- struct{}{}:
	case <-time.After(5 * time.Second):
		return &fetch.Error{Status: 429, Msg: "busy"}
	}
	s.perMu.Lock()
	ch, ok := s.per[key]
	if !ok {
		ch = make(chan struct{}, 2)
		s.per[key] = ch
	}
	s.perMu.Unlock()
	select {
	case ch <- struct{}{}:
		return nil
	case <-time.After(5 * time.Second):
		<-s.global
		return &fetch.Error{Status: 429, Msg: "busy"}
	}
}

func (s *Server) release(caller, rawURL string) {
	key := caller + "|" + hostOf(rawURL)
	s.perMu.Lock()
	ch := s.per[key]
	if ch != nil {
		select {
		case <-ch:
		default:
		}
		if len(ch) == 0 {
			delete(s.per, key)
		}
	}
	s.perMu.Unlock()
	select {
	case <-s.global:
	default:
	}
}

func hostOf(raw string) string {
	u, err := url.Parse(raw)
	if err != nil {
		return ""
	}
	return u.Hostname()
}

func writeErr(w http.ResponseWriter, status int, msg, hint string) {
	if status == 502 {
		status = 500
		msg = "internal error"
	}
	body := map[string]any{"success": false, "error": msg}
	if hint != "" {
		body["hint"] = hint
	}
	writeJSON(w, status, body)
}

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 HashBody(b []byte) string {
	sum := sha256.Sum256(b)
	return hex.EncodeToString(sum[:])
}

func Drain(r io.Reader) {
	_, _ = io.Copy(io.Discard, r)
}