deskview/internal/agent/engine.go

// Package agent 是应用的业务核心：组织消息、控制抓屏监听与调用大模型。
package agent

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

	"viewsagent/internal/capture"
	"viewsagent/internal/config"
	"viewsagent/internal/llm"
	"viewsagent/internal/vision"
)

// autoPrompt 是自动触发时附加的指令，避免模型对着同一道题反复作答。
const autoPrompt = "画面已更新。请针对当前画面中最新出现的内容作答，不要重复回答之前已经答过的题目；如果没有新内容，只回复「(无新内容)」。"

// ErrBusy 表示上一条回答尚未结束。
var ErrBusy = errors.New("上一条回答还没结束，请稍候")

// State 是引擎对外暴露的运行时快照。
type State struct {
	Watching    bool
	Busy        bool
	Triggers    int
	TotalTokens int
	LastLatency time.Duration
	Region      capture.Region
	Model       string
	Profile     string
}

// Sink 是引擎向界面推送事件的回调集合。
// 所有回调都可能来自任意 goroutine，实现方必须自行切换到 UI 线程。
type Sink interface {
	OnUserMessage(text string, withImage bool)
	OnAssistantStart(auto bool)
	OnAssistantDelta(delta string)
	OnAssistantEnd(usage llm.Usage, err error)
	OnState(state State)
}

// Engine 把抓屏、意图与模型调用串起来。
type Engine struct {
	mu       sync.Mutex
	cfg      *config.Config
	store    *config.Store
	screen   capture.Screen
	provider llm.Provider
	sink     Sink

	region      capture.Region
	watching    bool
	stopWatch   chan struct{}
	busy        bool
	triggers    int
	lastLatency time.Duration
	usage       llm.Usage

	// history 只保存纯文本，不保留历史图片，避免图片重复计费与内存膨胀。
	history []llm.Message
}

// New 创建引擎。
func New(cfg *config.Config, store *config.Store) *Engine {
	e := &Engine{
		cfg:    cfg,
		store:  store,
		region: cfg.Region,
	}
	if prov, err := llm.New(toOptions(cfg.Active())); err != nil {
		slog.Warn("模型接入点不可用", "err", err)
	} else {
		e.provider = prov
	}
	return e
}

func toOptions(p config.Profile) llm.Options {
	return llm.Options{
		Name:         p.Name,
		BaseURL:      p.BaseURL,
		APIKey:       p.APIKey,
		Model:        p.Model,
		ExtraHeaders: p.ExtraHeaders,
		MaxTokens:    p.MaxTokens,
		Temperature:  p.Temperature,
	}
}

// SetSink 绑定事件接收方。
func (e *Engine) SetSink(s Sink) {
	e.mu.Lock()
	e.sink = s
	e.mu.Unlock()
	e.emitState()
}

// Config 返回当前配置。
func (e *Engine) Config() *config.Config {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.cfg
}

// Store 返回配置存储。
func (e *Engine) Store() *config.Store { return e.store }

// State 返回当前状态快照。
func (e *Engine) State() State {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.stateLocked()
}

func (e *Engine) stateLocked() State {
	active := config.Profile{}
	if e.cfg != nil {
		active = e.cfg.Active()
	}
	return State{
		Watching:    e.watching,
		Busy:        e.busy,
		Triggers:    e.triggers,
		TotalTokens: e.usage.TotalTokens,
		LastLatency: e.lastLatency,
		Region:      e.region,
		Model:       active.Model,
		Profile:     active.Name,
	}
}

func (e *Engine) emitState() {
	e.mu.Lock()
	sink := e.sink
	st := e.stateLocked()
	e.mu.Unlock()
	if sink != nil {
		sink.OnState(st)
	}
}

// Region 返回当前监视区域。
func (e *Engine) Region() capture.Region {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.region
}

// SetRegion 更新监视区域并持久化。
func (e *Engine) SetRegion(r capture.Region) error {
	e.mu.Lock()
	e.region = r
	e.cfg.Region = r
	cfg := e.cfg
	e.mu.Unlock()

	e.emitState()
	if e.store != nil {
		if err := e.store.Save(cfg); err != nil {
			return fmt.Errorf("保存区域失败: %w", err)
		}
	}
	return nil
}

// ApplyConfig 在设置页保存后刷新引擎内部状态。
func (e *Engine) ApplyConfig(cfg *config.Config) error {
	prov, err := llm.New(toOptions(cfg.Active()))
	if err != nil {
		return err
	}
	e.mu.Lock()
	wasWatching := e.watching
	// 区域只由框选修改，不被设置页的旧副本覆盖。
	cfg.Region = e.region
	e.cfg = cfg
	e.provider = prov
	e.mu.Unlock()

	if wasWatching {
		e.StopWatch()
		if err := e.StartWatch(); err != nil {
			slog.Warn("配置变更后恢复监听失败", "err", err)
		}
	}
	e.emitState()
	return nil
}

// PingProfile 用给定接入点做一次连通性测试。
func (e *Engine) PingProfile(ctx context.Context, p config.Profile) error {
	prov, err := llm.New(toOptions(p))
	if err != nil {
		return err
	}
	return prov.Ping(ctx)
}

// Snapshot 抓取当前区域的 JPEG 预览（用于界面缩略图）。
func (e *Engine) Snapshot() ([]byte, error) {
	e.mu.Lock()
	region := e.region
	capCfg := e.cfg.Capture
	e.mu.Unlock()

	if !region.Valid() {
		return nil, fmt.Errorf("还没有框选监视区域")
	}
	img, err := e.screen.Grab(region)
	if err != nil {
		return nil, err
	}
	return vision.EncodeJPEG(img, capCfg.MaxEdge, capCfg.JPEGQuality)
}

// Ask 发起一次手动提问。attach 为 true 时附带当前区域截图。
func (e *Engine) Ask(text string, attach bool) error {
	text = strings.TrimSpace(text)
	if text == "" {
		return errors.New("请输入内容")
	}
	return e.run(text, attach, false)
}

// StartWatch 开始按帧差自动监听。
func (e *Engine) StartWatch() error {
	e.mu.Lock()
	if e.watching {
		e.mu.Unlock()
		return nil
	}
	if !e.region.Valid() {
		e.mu.Unlock()
		return errors.New("请先框选需要监视的屏幕区域")
	}
	stop := make(chan struct{})
	e.watching = true
	e.stopWatch = stop
	e.mu.Unlock()

	go e.watchLoop(stop)
	e.emitState()
	return nil
}

// StopWatch 停止自动监听。
func (e *Engine) StopWatch() {
	e.mu.Lock()
	if !e.watching {
		e.mu.Unlock()
		return
	}
	e.watching = false
	if e.stopWatch != nil {
		close(e.stopWatch)
		e.stopWatch = nil
	}
	e.mu.Unlock()
	e.emitState()
}

// ToggleWatch 在开始/停止之间切换。
func (e *Engine) ToggleWatch() error {
	e.mu.Lock()
	on := e.watching
	e.mu.Unlock()
	if on {
		e.StopWatch()
		return nil
	}
	return e.StartWatch()
}

// watchLoop 以低频抓屏 + 帧差判断是否需要提问。
func (e *Engine) watchLoop(stop chan struct{}) {
	e.mu.Lock()
	capCfg := e.cfg.Capture
	e.mu.Unlock()

	fps := capCfg.FPS
	if fps < 0.5 {
		fps = 0.5
	}
	interval := time.Duration(float64(time.Second) / fps)
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	var last vision.Fingerprint
	var haveLast bool
	var candidate bool
	var stable int
	var lastFire time.Time

	slog.Info("开始监听区域", "interval", interval, "threshold", capCfg.DiffThreshold)

	for {
		select {
		case <-stop:
			slog.Info("停止监听")
			return
		case <-ticker.C:
		}

		e.mu.Lock()
		region := e.region
		e.mu.Unlock()
		if !region.Valid() {
			continue
		}

		img, err := e.screen.Grab(region)
		if err != nil {
			slog.Debug("抓屏失败", "err", err)
			continue
		}
		fp := vision.Of(img)
		if !haveLast {
			last = fp
			haveLast = true
			continue
		}

		if fp.Changed(last, capCfg.DiffThreshold) {
			candidate = true
			stable = 0
			last = fp
			continue
		}
		last = fp

		if !candidate {
			continue
		}
		// 画面还在变化时先等它稳定，避免截到翻页/动画的中间帧。
		stable++
		if stable < capCfg.StableFrames {
			continue
		}
		candidate = false
		stable = 0

		if time.Since(lastFire) < time.Duration(capCfg.CooldownMs)*time.Millisecond {
			continue
		}
		if e.isBusy() {
			continue
		}

		lastFire = time.Now()
		e.mu.Lock()
		e.triggers++
		e.mu.Unlock()

		if err := e.run(autoPrompt, true, true); err != nil && !errors.Is(err, ErrBusy) {
			slog.Warn("自动提问失败", "err", err)
		}
		e.emitState()
	}
}

func (e *Engine) isBusy() bool {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.busy
}

// run 执行一次完整问答。
func (e *Engine) run(text string, attach, auto bool) error {
	e.mu.Lock()
	if e.busy {
		e.mu.Unlock()
		return ErrBusy
	}
	if e.provider == nil {
		e.mu.Unlock()
		return errors.New("还没有可用的模型接入点，请到设置里填写")
	}
	e.busy = true
	prov := e.provider
	region := e.region
	capCfg := e.cfg.Capture
	sink := e.sink
	system := e.cfg.Template(e.cfg.PromptTemplate).System
	history := append([]llm.Message(nil), e.history...)
	e.mu.Unlock()

	// 立刻广播一次，让界面在请求发出前就显示"正在回答"并禁用发送。
	e.emitState()

	defer func() {
		e.mu.Lock()
		e.busy = false
		e.mu.Unlock()
		e.emitState()
	}()

	var img []byte
	if attach && region.Valid() {
		shot, err := e.screen.Grab(region)
		if err != nil {
			slog.Warn("抓取监视区域失败", "err", err)
		} else if raw, err := vision.EncodeJPEG(shot, capCfg.MaxEdge, capCfg.JPEGQuality); err != nil {
			slog.Warn("图像编码失败", "err", err)
		} else {
			img = raw
		}
	}

	if sink != nil {
		if !auto {
			sink.OnUserMessage(text, len(img) > 0)
		}
		sink.OnAssistantStart(auto)
	}

	msgs := make([]llm.Message, 0, len(history)+1)
	msgs = append(msgs, history...)
	msgs = append(msgs, llm.Message{Role: llm.RoleUser, Text: text, Image: img})

	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()

	var reply strings.Builder
	start := time.Now()
	usage, err := prov.Chat(ctx, system, msgs, func(delta string) error {
		reply.WriteString(delta)
		if sink != nil {
			sink.OnAssistantDelta(delta)
		}
		return nil
	})
	elapsed := time.Since(start)

	e.mu.Lock()
	e.usage.Add(usage)
	e.lastLatency = elapsed
	if err == nil {
		e.appendHistoryLocked(text, len(img) > 0, reply.String())
	}
	e.mu.Unlock()

	if sink != nil {
		sink.OnAssistantEnd(usage, err)
	}
	if err != nil {
		return fmt.Errorf("调用模型失败: %w", err)
	}
	slog.Info("回答完成", "elapsed", elapsed, "tokens", usage.TotalTokens)
	return nil
}

// appendHistoryLocked 把一轮问答写入历史，图片只留占位文本。
func (e *Engine) appendHistoryLocked(question string, withImage bool, answer string) {
	if withImage {
		question += "\n[附带截图]"
	}
	e.history = append(e.history,
		llm.Message{Role: llm.RoleUser, Text: question},
		llm.Message{Role: llm.RoleAssistant, Text: answer},
	)
	limit := e.cfg.Capture.HistoryTurns * 2
	if limit < 2 {
		limit = 2
	}
	if len(e.history) > limit {
		e.history = append([]llm.Message(nil), e.history[len(e.history)-limit:]...)
	}
}

// ResetHistory 清空对话上下文。
func (e *Engine) ResetHistory() {
	e.mu.Lock()
	e.history = nil
	e.mu.Unlock()
}

// Close 释放资源。
func (e *Engine) Close() {
	e.StopWatch()
}