glpi-neural-brain/internal/sourceagent/agent.go

package sourceagent

import (
	"bytes"
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"html"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/local/glpi-neural-brain/internal/research"
	_ "modernc.org/sqlite"
)

type RunnerConfig struct {
	BrainURL      string
	AgentID       string
	Token         string
	DataDir       string
	ConfigFile    string
	ConfigRefresh time.Duration
	HTTPTimeout   time.Duration
	Concurrency   int
	BatchSize     int
	AllowPrivate  bool
	Version       string
}

type Runner struct {
	cfg        RunnerConfig
	http       *http.Client
	sourceHTTP *http.Client
	state      *localState
	mu         sync.RWMutex
	remote     RemoteConfig
	wake       chan struct{}
}

type bootstrapConfig struct {
	BrainURL string `json:"brain_url"`
	AgentID  string `json:"agent_id"`
	Token    string `json:"token"`
}

func NewRunner(cfg RunnerConfig) (*Runner, error) {
	if strings.TrimSpace(cfg.ConfigFile) != "" {
		if data, err := os.ReadFile(cfg.ConfigFile); err == nil {
			var b bootstrapConfig
			if json.Unmarshal(data, &b) == nil {
				if cfg.BrainURL == "" {
					cfg.BrainURL = b.BrainURL
				}
				if cfg.AgentID == "" {
					cfg.AgentID = b.AgentID
				}
				if cfg.Token == "" {
					cfg.Token = b.Token
				}
			}
		}
	}
	cfg.BrainURL = strings.TrimRight(strings.TrimSpace(cfg.BrainURL), "/")
	cfg.AgentID = strings.TrimSpace(cfg.AgentID)
	cfg.Token = strings.TrimSpace(cfg.Token)
	if cfg.BrainURL == "" || cfg.AgentID == "" || cfg.Token == "" {
		return nil, errors.New("agent mode requires BRAIN_AGENT_BRAIN_URL, BRAIN_AGENT_ID and BRAIN_AGENT_TOKEN (or BRAIN_AGENT_CONFIG_FILE)")
	}
	u, err := url.Parse(cfg.BrainURL)
	if err != nil || u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") {
		return nil, errors.New("BRAIN_AGENT_BRAIN_URL must be an absolute http(s) URL")
	}
	if u.User != nil {
		return nil, errors.New("BRAIN_AGENT_BRAIN_URL must not contain userinfo")
	}
	if cfg.ConfigRefresh < time.Minute {
		cfg.ConfigRefresh = 5 * time.Minute
	}
	if cfg.HTTPTimeout < 5*time.Second {
		cfg.HTTPTimeout = 30 * time.Second
	}
	if cfg.Concurrency < 1 {
		cfg.Concurrency = 3
	}
	if cfg.Concurrency > 16 {
		cfg.Concurrency = 16
	}
	if cfg.BatchSize < 1 {
		cfg.BatchSize = 50
	}
	if cfg.BatchSize > 500 {
		cfg.BatchSize = 500
	}
	state, err := openLocalState(cfg.DataDir)
	if err != nil {
		return nil, err
	}
	return &Runner{
		cfg:        cfg,
		http:       &http.Client{Timeout: cfg.HTTPTimeout},
		sourceHTTP: research.NewSafeHTTPClient(cfg.AllowPrivate, cfg.HTTPTimeout),
		state:      state,
		wake:       make(chan struct{}, 1),
	}, nil
}
func (r *Runner) Close() error {
	if r == nil || r.state == nil {
		return nil
	}
	return r.state.Close()
}

func (r *Runner) Start(ctx context.Context) {
	r.loadCachedConfig()
	go r.loop(ctx)
}

func (r *Runner) Status() map[string]any {
	r.mu.RLock()
	remote := r.remote
	r.mu.RUnlock()
	return map[string]any{"ok": true, "mode": "agent", "agent_id": r.cfg.AgentID, "brain_url": r.cfg.BrainURL, "configured_tasks": len(remote.Tasks), "config_issued_at": remote.IssuedAt, "version": r.cfg.Version}
}

func (r *Runner) loop(ctx context.Context) {
	refresh := time.NewTicker(r.cfg.ConfigRefresh)
	defer refresh.Stop()
	run := time.NewTicker(30 * time.Second)
	defer run.Stop()
	_ = r.refreshConfig(ctx)
	r.runDue(ctx)
	for {
		select {
		case <-ctx.Done():
			return
		case <-refresh.C:
			_ = r.refreshConfig(ctx)
		case <-run.C:
			r.runDue(ctx)
		case <-r.wake:
			r.runDue(ctx)
		}
	}
}

func (r *Runner) refreshConfig(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, r.cfg.BrainURL+"/api/v1/agent/config", nil)
	if err != nil {
		return err
	}
	r.auth(req)
	resp, err := r.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		return fmt.Errorf("brain config returned HTTP %d", resp.StatusCode)
	}
	var cfg RemoteConfig
	if err := json.NewDecoder(io.LimitReader(resp.Body, 4<<20)).Decode(&cfg); err != nil {
		return err
	}
	if cfg.Agent.ID != "" && cfg.Agent.ID != r.cfg.AgentID {
		return fmt.Errorf("brain returned config for unexpected agent %q", cfg.Agent.ID)
	}
	r.mu.Lock()
	r.remote = cfg
	r.mu.Unlock()
	r.saveCachedConfig(cfg)
	_ = r.sendHeartbeat(ctx, Heartbeat{AgentID: r.cfg.AgentID, Version: r.cfg.Version, Status: "online", Metadata: map[string]any{"configured_tasks": len(cfg.Tasks)}})
	return nil
}

func (r *Runner) runDue(ctx context.Context) {
	r.mu.RLock()
	tasks := append([]Task(nil), r.remote.Tasks...)
	r.mu.RUnlock()
	if len(tasks) == 0 {
		return
	}
	sem := make(chan struct{}, r.cfg.Concurrency)
	var wg sync.WaitGroup
	for _, task := range tasks {
		task := task
		if !task.Enabled || !r.state.Due(ctx, task) {
			continue
		}
		wg.Add(1)
		go func() {
			defer wg.Done()
			select {
			case sem <- struct{}{}:
			case <-ctx.Done():
				return
			}
			defer func() { <-sem }()
			r.runTask(ctx, task)
		}()
	}
	wg.Wait()
}

func (r *Runner) runTask(ctx context.Context, task Task) {
	started := time.Now().UTC()
	docs, err := r.pollTask(ctx, task)
	if err == nil && len(docs) > 0 {
		for start := 0; start < len(docs); start += r.cfg.BatchSize {
			end := start + r.cfg.BatchSize
			if end > len(docs) {
				end = len(docs)
			}
			if sendErr := r.sendBatch(ctx, task.ID, docs[start:end]); sendErr != nil {
				err = sendErr
				break
			}
			for _, d := range docs[start:end] {
				_ = r.state.MarkSeen(ctx, task.ID, d.CanonicalURL, d.ContentSHA256)
			}
		}
	}
	_ = r.state.FinishTask(ctx, task.ID, started, err)
	h := Heartbeat{AgentID: r.cfg.AgentID, Version: r.cfg.Version, Status: "ok", LastRunAt: started, TasksChecked: 1, Documents: len(docs)}
	if err != nil {
		h.Status = "error"
		h.LastError = err.Error()
		slog.Warn("source agent task failed", "task", task.ID, "error", err)
	} else {
		slog.Info("source agent task completed", "task", task.ID, "documents", len(docs))
	}
	_ = r.sendHeartbeat(ctx, h)
}

func (r *Runner) pollTask(ctx context.Context, task Task) ([]Document, error) {
	switch task.Type {
	case "rss", "atom":
		return r.pollFeed(ctx, task)
	case "sitemap":
		return r.pollSitemap(ctx, task)
	case "web":
		return r.pollWeb(ctx, task)
	default:
		return nil, fmt.Errorf("unsupported task type %q", task.Type)
	}
}

type feedEnvelope struct {
	Channel struct {
		Items []struct {
			Title       string `xml:"title"`
			Link        string `xml:"link"`
			Description string `xml:"description"`
			PubDate     string `xml:"pubDate"`
			GUID        string `xml:"guid"`
		} `xml:"item"`
	} `xml:"channel"`
	Entries []struct {
		Title     string `xml:"title"`
		ID        string `xml:"id"`
		Updated   string `xml:"updated"`
		Published string `xml:"published"`
		Summary   string `xml:"summary"`
		Content   string `xml:"content"`
		Links     []struct {
			Href string `xml:"href,attr"`
			Rel  string `xml:"rel,attr"`
		} `xml:"link"`
	} `xml:"entry"`
}

type feedItem struct {
	Title, Link, Summary string
	Published            time.Time
}

func (r *Runner) pollFeed(ctx context.Context, task Task) ([]Document, error) {
	body, finalURL, _, _, err := r.fetchRaw(ctx, task.URL, 5<<20)
	if err != nil {
		return nil, err
	}
	var env feedEnvelope
	if err := xml.Unmarshal(body, &env); err != nil {
		return nil, fmt.Errorf("feed XML: %w", err)
	}
	var items []feedItem
	for _, it := range env.Channel.Items {
		link := strings.TrimSpace(it.Link)
		if link == "" {
			link = strings.TrimSpace(it.GUID)
		}
		items = append(items, feedItem{Title: strings.TrimSpace(it.Title), Link: resolveLink(finalURL, link), Summary: stripMarkup(it.Description), Published: parsePublished(it.PubDate)})
	}
	for _, it := range env.Entries {
		link := ""
		for _, l := range it.Links {
			if l.Rel == "" || l.Rel == "alternate" {
				link = l.Href
				break
			}
		}
		pub := parsePublished(it.Published)
		if pub.IsZero() {
			pub = parsePublished(it.Updated)
		}
		summary := it.Summary
		if summary == "" {
			summary = it.Content
		}
		items = append(items, feedItem{Title: strings.TrimSpace(it.Title), Link: resolveLink(finalURL, link), Summary: stripMarkup(summary), Published: pub})
	}
	sort.SliceStable(items, func(i, j int) bool { return items[i].Published.After(items[j].Published) })
	if len(items) > task.MaxItems {
		items = items[:task.MaxItems]
	}
	return r.materializeItems(ctx, task, items)
}

type sitemapEnvelope struct {
	URLs []struct {
		Loc     string `xml:"loc"`
		LastMod string `xml:"lastmod"`
	} `xml:"url"`
	Sitemaps []struct {
		Loc string `xml:"loc"`
	} `xml:"sitemap"`
}

func (r *Runner) pollSitemap(ctx context.Context, task Task) ([]Document, error) {
	items, err := r.collectSitemapItems(ctx, task.URL, task.MaxItems, 0)
	if err != nil {
		return nil, err
	}
	return r.materializeItems(ctx, task, items)
}

func (r *Runner) collectSitemapItems(ctx context.Context, rawURL string, limit, depth int) ([]feedItem, error) {
	if limit <= 0 || depth > 1 {
		return nil, nil
	}
	body, finalURL, _, _, err := r.fetchRaw(ctx, rawURL, 8<<20)
	if err != nil {
		return nil, err
	}
	var sm sitemapEnvelope
	if err := xml.Unmarshal(body, &sm); err != nil {
		return nil, err
	}
	items := make([]feedItem, 0, limit)
	for _, u := range sm.URLs {
		link := resolveLink(finalURL, u.Loc)
		if link == "" {
			continue
		}
		items = append(items, feedItem{Link: link, Published: parsePublished(u.LastMod)})
		if len(items) >= limit {
			return items, nil
		}
	}
	// Sitemap indexes are common on larger publishers. Follow a bounded number of
	// child maps once; source HTTP safety rules apply to every child URL.
	for i, child := range sm.Sitemaps {
		if len(items) >= limit || i >= 12 {
			break
		}
		childURL := resolveLink(finalURL, child.Loc)
		if childURL == "" {
			continue
		}
		more, childErr := r.collectSitemapItems(ctx, childURL, limit-len(items), depth+1)
		if childErr != nil {
			continue
		}
		items = append(items, more...)
	}
	return items, nil
}

var hrefPattern = regexp.MustCompile(`(?is)<a\b[^>]*href\s*=\s*["']([^"'#]+)["'][^>]*>(.*?)</a>`)

func (r *Runner) pollWeb(ctx context.Context, task Task) ([]Document, error) {
	body, finalURL, _, _, err := r.fetchRaw(ctx, task.URL, 5<<20)
	if err != nil {
		return nil, err
	}
	matches := hrefPattern.FindAllStringSubmatch(string(body), -1)
	seen := map[string]bool{}
	items := make([]feedItem, 0, task.MaxItems)
	base, _ := url.Parse(finalURL)
	for _, m := range matches {
		link := resolveLink(finalURL, m[1])
		if link == "" || seen[link] {
			continue
		}
		u, err := url.Parse(link)
		if err != nil || u.Host != base.Host {
			continue
		}
		if !looksArticleLink(u.Path, stripMarkup(m[2])) {
			continue
		}
		seen[link] = true
		items = append(items, feedItem{Title: stripMarkup(m[2]), Link: link})
		if len(items) >= task.MaxItems {
			break
		}
	}
	return r.materializeItems(ctx, task, items)
}

func (r *Runner) materializeItems(ctx context.Context, task Task, items []feedItem) ([]Document, error) {
	fetcher := research.New("")
	out := make([]Document, 0, len(items))
	for _, item := range items {
		if strings.TrimSpace(item.Link) == "" {
			continue
		}
		if r.state.SeenURL(ctx, task.ID, item.Link) && !strings.EqualFold(strings.TrimSpace(task.Config["refetch_seen"]), "true") {
			continue
		}
		page, diag, err := fetcher.FetchPage(ctx, item.Link, research.FetchOptions{MaxBytes: 2 << 20, MaxChars: 20000, Timeout: r.cfg.HTTPTimeout, AllowPrivate: r.cfg.AllowPrivate})
		text := strings.TrimSpace(item.Summary)
		title := strings.TrimSpace(item.Title)
		ctype := "text/html"
		final := item.Link
		if err == nil {
			if page.Content != "" {
				text = page.Content
			}
			if page.Title != "" {
				title = page.Title
			}
			if page.URL != "" {
				final = page.URL
			}
			ctype = page.ContentType
		} else if len([]rune(text)) < 80 {
			continue
		}
		if title == "" {
			title = final
		}
		sum := sha256.Sum256([]byte(text))
		sha := hex.EncodeToString(sum[:])
		if r.state.SeenHash(ctx, task.ID, final, sha) {
			continue
		}
		baseURL := ""
		if u, e := url.Parse(task.URL); e == nil {
			baseURL = u.Scheme + "://" + u.Host
		}
		out = append(out, Document{ExternalID: final, URL: final, CanonicalURL: final, Title: title, PublishedAt: item.Published, DiscoveredAt: time.Now().UTC(), ContentType: ctype, Text: text, ContentSHA256: sha, SourceName: task.Name, SourceBaseURL: baseURL, Categories: task.Categories, Metadata: map[string]any{"agent_task_type": task.Type, "fetch_error_kind": diag.ErrorKind}})
		if len(out) >= task.MaxItems {
			break
		}
	}
	return out, nil
}

func (r *Runner) fetchRaw(ctx context.Context, raw string, max int64) ([]byte, string, string, string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, raw, nil)
	if err != nil {
		return nil, "", "", "", err
	}
	req.Header.Set("User-Agent", "glpi-neural-brain-source-agent/1.0")
	req.Header.Set("Accept", "application/rss+xml, application/atom+xml, application/xml, text/xml, text/html;q=0.9, */*;q=0.5")
	resp, err := r.sourceHTTP.Do(req)
	if err != nil {
		return nil, "", "", "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		return nil, "", "", "", fmt.Errorf("source returned HTTP %d", resp.StatusCode)
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, max+1))
	if err != nil {
		return nil, "", "", "", err
	}
	if int64(len(body)) > max {
		return nil, "", "", "", errors.New("source response too large")
	}
	return body, resp.Request.URL.String(), resp.Header.Get("ETag"), resp.Header.Get("Last-Modified"), nil
}

func (r *Runner) sendBatch(ctx context.Context, taskID string, docs []Document) error {
	payload := IngestBatch{SchemaVersion: SchemaVersion, AgentID: r.cfg.AgentID, TaskID: taskID, Documents: docs}
	data, _ := json.Marshal(payload)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.cfg.BrainURL+"/api/v1/agent/ingest", bytes.NewReader(data))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	r.auth(req)
	resp, err := r.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return fmt.Errorf("brain ingest HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(b)))
	}
	return nil
}
func (r *Runner) sendHeartbeat(ctx context.Context, h Heartbeat) error {
	data, _ := json.Marshal(h)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.cfg.BrainURL+"/api/v1/agent/heartbeat", bytes.NewReader(data))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	r.auth(req)
	resp, err := r.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		return fmt.Errorf("heartbeat HTTP %d", resp.StatusCode)
	}
	return nil
}
func (r *Runner) auth(req *http.Request) {
	req.Header.Set("Authorization", "Bearer "+r.cfg.Token)
	req.Header.Set("X-Brain-Agent-ID", r.cfg.AgentID)
}

func resolveLink(base, ref string) string {
	ref = strings.TrimSpace(html.UnescapeString(ref))
	if ref == "" {
		return ""
	}
	u, err := url.Parse(ref)
	if err != nil {
		return ""
	}
	b, err := url.Parse(base)
	if err != nil {
		return ""
	}
	return b.ResolveReference(u).String()
}
func stripMarkup(v string) string {
	v = regexp.MustCompile(`(?is)<[^>]+>`).ReplaceAllString(v, " ")
	return strings.Join(strings.Fields(html.UnescapeString(v)), " ")
}
func looksArticleLink(path, label string) bool {
	p := strings.ToLower(path + " " + label)
	if strings.Contains(p, "/tag/") || strings.Contains(p, "/category/") || strings.Contains(p, "/author/") || strings.Contains(p, "login") || strings.Contains(p, "privacy") || strings.Contains(p, "impress") || strings.Contains(p, "kontakt") {
		return false
	}
	segments := strings.Split(strings.Trim(path, "/"), "/")
	return len(segments) >= 2 || strings.Contains(p, "news") || strings.Contains(p, "blog") || strings.Contains(p, "advis") || strings.Contains(p, "release") || strings.Contains(p, "security")
}
func parsePublished(v string) time.Time {
	v = strings.TrimSpace(v)
	for _, layout := range []string{time.RFC3339, time.RFC1123Z, time.RFC1123, time.RFC822Z, time.RFC822, "2006-01-02"} {
		if t, err := time.Parse(layout, v); err == nil {
			return t.UTC()
		}
	}
	return time.Time{}
}

func (r *Runner) configCachePath() string {
	return filepath.Join(r.cfg.DataDir, "source-agent-config-cache.json")
}
func (r *Runner) saveCachedConfig(cfg RemoteConfig) {
	data, err := json.MarshalIndent(cfg, "", "  ")
	if err != nil {
		return
	}
	_ = os.WriteFile(r.configCachePath(), append(data, '\n'), 0o600)
}
func (r *Runner) loadCachedConfig() {
	data, err := os.ReadFile(r.configCachePath())
	if err != nil {
		return
	}
	var cfg RemoteConfig
	if json.Unmarshal(data, &cfg) != nil || cfg.Agent.ID != r.cfg.AgentID {
		return
	}
	r.mu.Lock()
	r.remote = cfg
	r.mu.Unlock()
}

type localState struct {
	db *sql.DB
	mu sync.Mutex
}

func openLocalState(dataDir string) (*localState, error) {
	if strings.TrimSpace(dataDir) == "" {
		dataDir = "./data"
	}
	if err := os.MkdirAll(dataDir, 0o750); err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", "file:"+filepath.ToSlash(filepath.Join(dataDir, "source-agent-local.db"))+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)")
	if err != nil {
		return nil, err
	}
	st := &localState{db: db}
	for _, q := range []string{`CREATE TABLE IF NOT EXISTS task_state(task_id TEXT PRIMARY KEY,last_run_ns INTEGER NOT NULL DEFAULT 0,last_error TEXT NOT NULL DEFAULT '') WITHOUT ROWID`, `CREATE TABLE IF NOT EXISTS seen(task_id TEXT NOT NULL,url TEXT NOT NULL,content_sha256 TEXT NOT NULL,seen_at_ns INTEGER NOT NULL,PRIMARY KEY(task_id,url,content_sha256)) WITHOUT ROWID`, `CREATE INDEX IF NOT EXISTS idx_seen_task_url ON seen(task_id,url,seen_at_ns DESC)`} {
		if _, err := db.Exec(q); err != nil {
			db.Close()
			return nil, err
		}
	}
	return st, nil
}
func (s *localState) Close() error { return s.db.Close() }
func (s *localState) Due(ctx context.Context, t Task) bool {
	d, err := time.ParseDuration(t.PollInterval)
	if err != nil {
		d = 4 * time.Hour
	}
	var last int64
	err = s.db.QueryRowContext(ctx, `SELECT last_run_ns FROM task_state WHERE task_id=?`, t.ID).Scan(&last)
	return err == sql.ErrNoRows || err != nil || last == 0 || time.Since(time.Unix(0, last)) >= d
}
func (s *localState) FinishTask(ctx context.Context, id string, started time.Time, err error) error {
	msg := ""
	if err != nil {
		msg = err.Error()
	}
	_, e := s.db.ExecContext(ctx, `INSERT INTO task_state(task_id,last_run_ns,last_error) VALUES(?,?,?) ON CONFLICT(task_id) DO UPDATE SET last_run_ns=excluded.last_run_ns,last_error=excluded.last_error`, id, started.UnixNano(), msg)
	return e
}
func (s *localState) SeenURL(ctx context.Context, task, urlv string) bool {
	var n int
	_ = s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM seen WHERE task_id=? AND url=?`, task, urlv).Scan(&n)
	return n > 0
}
func (s *localState) SeenHash(ctx context.Context, task, urlv, sha string) bool {
	var n int
	_ = s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM seen WHERE task_id=? AND url=? AND content_sha256=?`, task, urlv, sha).Scan(&n)
	return n > 0
}
func (s *localState) MarkSeen(ctx context.Context, task, urlv, sha string) error {
	_, err := s.db.ExecContext(ctx, `INSERT OR IGNORE INTO seen(task_id,url,content_sha256,seen_at_ns) VALUES(?,?,?,?)`, task, urlv, sha, time.Now().UTC().UnixNano())
	return err
}