sift/internal/search/sources.go

package search

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"sort"
	"strings"

	"git.wntrmute.dev/kyle/sift/internal/search/dial"
)

func warnStatus(name string, status int) string {
	if status == 0 {
		return name + ": error"
	}
	return fmt.Sprintf("%s: %d", name, status)
}

func warnErr(name string, err error) string {
	if err == nil {
		return ""
	}
	if strings.Contains(strings.ToLower(err.Error()), "timeout") || strings.Contains(err.Error(), "deadline") {
		return name + ": timeout"
	}
	return name + ": error"
}

func pinErr(ctx context.Context, c *dial.Client, name string) string {
	if c == nil {
		return name + ": error"
	}
	if err := c.Ensure(ctx); err != nil {
		return warnErr(name, err)
	}
	return ""
}

func readJSON(resp *http.Response, dest any) error {
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 512))
		return fmt.Errorf("status %d", resp.StatusCode)
	}
	dec := json.NewDecoder(io.LimitReader(resp.Body, 1<<20))
	return dec.Decode(dest)
}

// SearxNG is GET {url}/search?q=&format=json&pageno=1
type SearxNG struct {
	Client *dial.Client
}

func (s *SearxNG) Name() string { return NameSearxNG }
func (s *SearxNG) Ready() bool  { return s != nil && s.Client.Ready() }
func (s *BraveLLM) Ready() bool { return s != nil && s.Client.Ready() }

func (s *SearxNG) Search(ctx context.Context, query string, k int) SourceOut {
	if s == nil || s.Client == nil {
		return SourceOut{Warn: NameSearxNG + ": error"}
	}
	if w := pinErr(ctx, s.Client, NameSearxNG); w != "" {
		return SourceOut{Warn: w}
	}
	base := strings.TrimRight(s.Client.URL.String(), "/")
	u, err := url.Parse(base + "/search")
	if err != nil {
		return SourceOut{Warn: warnErr(NameSearxNG, err)}
	}
	q := u.Query()
	q.Set("q", query)
	q.Set("format", "json")
	q.Set("pageno", "1")
	u.RawQuery = q.Encode()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
	if err != nil {
		return SourceOut{Warn: warnErr(NameSearxNG, err)}
	}
	req.Header.Set("Accept", "application/json")
	resp, err := s.Client.HTTP.Do(req)
	if err != nil {
		return SourceOut{Warn: warnErr(NameSearxNG, err)}
	}
	if resp.StatusCode != 200 {
		_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 256))
		resp.Body.Close()
		return SourceOut{Warn: warnStatus(NameSearxNG, resp.StatusCode)}
	}
	var env struct {
		Results []struct {
			URL     string  `json:"url"`
			Title   string  `json:"title"`
			Content string  `json:"content"`
			Score   float64 `json:"score"`
		} `json:"results"`
	}
	if err := readJSON(resp, &env); err != nil {
		return SourceOut{Warn: warnErr(NameSearxNG, err)}
	}
	sort.SliceStable(env.Results, func(i, j int) bool { return env.Results[i].Score > env.Results[j].Score })
	hits := make([]Hit, 0, len(env.Results))
	for _, r := range env.Results {
		if r.URL == "" {
			continue
		}
		hits = append(hits, Hit{URL: r.URL, Title: r.Title, Snippet: r.Content, Score: r.Score, Source: NameSearxNG})
	}
	return SourceOut{Hits: hits}
}

// BraveLLM is GET {url}?q= with X-Subscription-Token.
type BraveLLM struct {
	Client *dial.Client
	Key    string
}

func (s *BraveLLM) Name() string { return NameBraveLLM }

func (s *BraveLLM) Search(ctx context.Context, query string, k int) SourceOut {
	if s == nil || s.Client == nil {
		return SourceOut{Warn: NameBraveLLM + ": error"}
	}
	if w := pinErr(ctx, s.Client, NameBraveLLM); w != "" {
		return SourceOut{Warn: w}
	}
	u, err := url.Parse(s.Client.URL.String())
	if err != nil {
		return SourceOut{Warn: warnErr(NameBraveLLM, err)}
	}
	q := u.Query()
	q.Set("q", query)
	u.RawQuery = q.Encode()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
	if err != nil {
		return SourceOut{Warn: warnErr(NameBraveLLM, err)}
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set("X-Subscription-Token", s.Key)
	resp, err := s.Client.HTTP.Do(req)
	if err != nil {
		return SourceOut{Warn: warnErr(NameBraveLLM, err)}
	}
	if resp.StatusCode != 200 {
		code := resp.StatusCode
		resp.Body.Close()
		return SourceOut{Warn: warnStatus(NameBraveLLM, code)}
	}
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	resp.Body.Close()
	if err != nil {
		return SourceOut{Warn: warnErr(NameBraveLLM, err)}
	}
	hits := mapBrave(raw)
	credits := 0
	if len(hits) >= 0 {
		credits = 1
	}
	return SourceOut{Hits: hits, Credits: credits}
}

func mapBrave(raw []byte) []Hit {
	var top map[string]any
	if err := json.Unmarshal(raw, &top); err != nil {
		return nil
	}
	var buckets []any
	for _, key := range []string{"sources", "results", "web"} {
		if v, ok := top[key]; ok {
			buckets = append(buckets, v)
		}
	}
	var hits []Hit
	seen := map[string]bool{}
	var walk func(any)
	walk = func(v any) {
		switch t := v.(type) {
		case []any:
			for _, x := range t {
				walk(x)
			}
		case map[string]any:
			u := strField(t, "url", "link", "href")
			if u == "" {
				for _, x := range t {
					walk(x)
				}
				return
			}
			if seen[u] {
				return
			}
			seen[u] = true
			hits = append(hits, Hit{
				URL:     u,
				Title:   strField(t, "title", "name"),
				Snippet: strField(t, "snippet", "description", "content", "text"),
				Source:  NameBraveLLM,
			})
		}
	}
	for _, b := range buckets {
		walk(b)
	}
	return hits
}

func strField(m map[string]any, keys ...string) string {
	for _, k := range keys {
		if v, ok := m[k]; ok {
			if s, ok := v.(string); ok && s != "" {
				return s
			}
		}
	}
	return ""
}