wechat_ai/agent/vision_stream.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"image"
	_ "image/jpeg"
	"io"
	"math"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

type visionFrameState struct {
	condition   *sync.Cond
	jpeg        []byte
	frameWidth  int
	frameHeight int
	sequence    int64
	frameIndex  int64
	updatedAt   time.Time
	target      previewTarget
	annotation  annotationMeta
	event       map[string]any
}

type annotationMeta struct {
	Width       int     `json:"width"`
	Height      int     `json:"height"`
	ScaleFactor float64 `json:"scaleFactor"`
}

type captureSource struct {
	ID          string  `json:"id"`
	Kind        string  `json:"kind"`
	Label       string  `json:"label"`
	AppName     string  `json:"appName"`
	Title       string  `json:"title"`
	PID         int     `json:"pid"`
	X           float64 `json:"x"`
	Y           float64 `json:"y"`
	Width       float64 `json:"width"`
	Height      float64 `json:"height"`
	ScaleFactor float64 `json:"scaleFactor"`
}

type previewTarget struct {
	Kind     string
	Label    string
	WindowID int
	Rect     WindowBounds
}

func newVisionFrameState(target previewTarget) *visionFrameState {
	return &visionFrameState{
		condition: sync.NewCond(&sync.Mutex{}),
		target:    target,
		event:     map[string]any{"type": "waiting"},
	}
}

func (state *visionFrameState) update(jpeg []byte) {
	frameWidth, frameHeight := jpegSize(jpeg)
	state.condition.L.Lock()
	state.jpeg = append(state.jpeg[:0], jpeg...)
	if frameWidth > 0 && frameHeight > 0 {
		state.frameWidth = frameWidth
		state.frameHeight = frameHeight
	}
	state.sequence++
	state.frameIndex++
	state.updatedAt = time.Now().UTC()
	state.event = map[string]any{
		"type":         "preview_frame",
		"timestamp":    state.updatedAt.Format(time.RFC3339Nano),
		"frame_index":  state.frameIndex,
		"target_kind":  state.target.Kind,
		"target_label": state.target.Label,
	}
	state.condition.Broadcast()
	state.condition.L.Unlock()
}

func (state *visionFrameState) updateAnnotation(annotation annotationMeta) {
	state.condition.L.Lock()
	state.annotation = annotation
	state.condition.L.Unlock()
}

func (state *visionFrameState) updateStatus(event map[string]any, target previewTarget) {
	state.condition.L.Lock()
	state.updatedAt = time.Now().UTC()
	state.event = event
	state.target = target
	state.condition.Broadcast()
	state.condition.L.Unlock()
}

func (state *visionFrameState) snapshot() ([]byte, map[string]any, int64, time.Time, int, int, annotationMeta) {
	state.condition.L.Lock()
	defer state.condition.L.Unlock()
	jpeg := append([]byte(nil), state.jpeg...)
	event := make(map[string]any, len(state.event))
	for key, value := range state.event {
		event[key] = value
	}
	return jpeg, event, state.sequence, state.updatedAt, state.frameWidth, state.frameHeight, state.annotation
}

func (state *visionFrameState) waitForNext(lastSequence int64) ([]byte, int64) {
	state.condition.L.Lock()
	defer state.condition.L.Unlock()
	for state.sequence == lastSequence {
		state.condition.Wait()
	}
	return append([]byte(nil), state.jpeg...), state.sequence
}

func runVisionStream(args []string) error {
	flags := flag.NewFlagSet("vision-stream", flag.ContinueOnError)
	host := flags.String("host", "127.0.0.1", "HTTP server host")
	port := flags.Int("port", 8765, "HTTP server port")
	input := flags.String("input", "3:none", "FFmpeg avfoundation input for experiments")
	fps := flags.Float64("fps", 30, "capture FPS")
	cropWechat := flags.Bool("crop-wechat", true, "capture only the WeChat window")
	regionsPath := flags.String("regions", defaultRegionsPath(), "annotation regions JSON path")
	if err := flags.Parse(args); err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	state := newVisionFrameState(previewTarget{Kind: "waiting", Label: "等待标注来源"})
	server := &http.Server{
		Addr:              fmt.Sprintf("%s:%d", *host, *port),
		Handler:           visionHTTPHandler(state),
		ReadHeaderTimeout: 5 * time.Second,
	}

	go func() {
		<-ctx.Done()
		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer shutdownCancel()
		_ = server.Shutdown(shutdownCtx)
	}()

	errCh := make(chan error, 2)
	go func() {
		errCh <- server.ListenAndServe()
	}()
	go func() {
		errCh <- runVisionCapture(ctx, state, *input, *fps, *cropWechat, *regionsPath)
	}()

	logInfo(fmt.Sprintf("Go annotation source stream started http://%s:%d/stream.mjpg regions=%s fps=%.1f", *host, *port, *regionsPath, *fps))

	for err := range errCh {
		if err == nil || err == http.ErrServerClosed || ctx.Err() != nil {
			return nil
		}
		cancel()
		return err
	}

	return nil
}

func visionHTTPHandler(state *visionFrameState) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "text/html; charset=utf-8")
		_, _ = io.WriteString(writer, `<html><body style="margin:0;background:#050505;color:#d7ffe4;font-family:monospace"><img src="/stream.mjpg" style="width:100vw;height:100vh;object-fit:contain" /></body></html>`)
	})
	mux.HandleFunc("/latest.json", func(writer http.ResponseWriter, request *http.Request) {
		_, event, sequence, updatedAt, frameWidth, frameHeight, annotation := state.snapshot()
		payload := map[string]any{
			"source":     "go-window-capture",
			"sequence":   sequence,
			"updated_at": updatedAt.Format(time.RFC3339Nano),
			"event":      event,
			"frame": map[string]any{
				"width":  frameWidth,
				"height": frameHeight,
			},
			"annotation": annotation,
			"target": map[string]any{
				"kind":      state.target.Kind,
				"label":     state.target.Label,
				"window_id": state.target.WindowID,
				"rect":      state.target.Rect,
			},
		}
		if updatedAt.IsZero() {
			payload["updated_at"] = time.Now().UTC().Format(time.RFC3339Nano)
		}
		writer.Header().Set("Content-Type", "application/json; charset=utf-8")
		writer.Header().Set("Cache-Control", "no-store")
		_ = json.NewEncoder(writer).Encode(payload)
	})
	mux.HandleFunc("/frame.jpg", func(writer http.ResponseWriter, request *http.Request) {
		jpeg, _, sequence, _, _, _, _ := state.snapshot()
		if len(jpeg) == 0 {
			http.Error(writer, "frame not ready", http.StatusServiceUnavailable)
			return
		}
		writer.Header().Set("Content-Type", "image/jpeg")
		writer.Header().Set("Cache-Control", "no-store")
		writer.Header().Set("X-Frame-Sequence", fmt.Sprintf("%d", sequence))
		_, _ = writer.Write(jpeg)
	})
	mux.HandleFunc("/stream.mjpg", func(writer http.ResponseWriter, request *http.Request) {
		flusher, ok := writer.(http.Flusher)
		if !ok {
			http.Error(writer, "streaming unsupported", http.StatusInternalServerError)
			return
		}

		writer.Header().Set("Content-Type", "multipart/x-mixed-replace; boundary=frame")
		writer.Header().Set("Cache-Control", "no-store")
		lastSequence := int64(-1)
		for request.Context().Err() == nil {
			jpeg, sequence := state.waitForNext(lastSequence)
			lastSequence = sequence
			if len(jpeg) == 0 {
				continue
			}
			_, err := fmt.Fprintf(writer, "--frame\r\nContent-Type: image/jpeg\r\nContent-Length: %d\r\n\r\n", len(jpeg))
			if err != nil {
				return
			}
			if _, err := writer.Write(jpeg); err != nil {
				return
			}
			if _, err := writer.Write([]byte("\r\n")); err != nil {
				return
			}
			flusher.Flush()
		}
	})
	return mux
}

func runVisionCapture(ctx context.Context, state *visionFrameState, input string, fps float64, cropWechat bool, regionsPath string) error {
	for ctx.Err() == nil {
		target, annotation, err := previewTargetFromAnnotation(regionsPath)
		if err != nil {
			message := "等待标注来源窗口出现：" + err.Error()
			logWarning(message)
			state.updateStatus(map[string]any{
				"type":      "source_waiting",
				"timestamp": time.Now().UTC().Format(time.RFC3339Nano),
				"message":   message,
			}, previewTarget{Kind: "waiting", Label: "等待标注来源"})
			select {
			case <-ctx.Done():
				return nil
			case <-time.After(2 * time.Second):
				continue
			}
		}

		state.updateAnnotation(annotation)
		return runResolvedVisionCapture(ctx, state, input, fps, cropWechat, target)
	}
	return nil
}

func runResolvedVisionCapture(ctx context.Context, state *visionFrameState, input string, fps float64, cropWechat bool, target previewTarget) error {
	if target.Kind == "window" {
		logInfo(fmt.Sprintf("using macOS window-only capture window_id=%d", target.WindowID))
		return runScreencaptureMJPEG(ctx, state, fps, &target.Rect, target.WindowID)
	}

	if target.Kind == "display" {
		logInfo(fmt.Sprintf("using macOS display rectangle capture rect=[%.0f,%.0f,%.0f,%.0f]", target.Rect.X, target.Rect.Y, target.Rect.Width, target.Rect.Height))
		return runScreencaptureMJPEG(ctx, state, fps, &target.Rect, 0)
	}

	window, err := maybeFindWeChatWindow(cropWechat)
	if err != nil {
		return err
	}
	logInfo("using macOS window-only capture for fallback WeChat preview")
	return runScreencaptureMJPEG(ctx, state, fps, window, window.ID)

	// Keep the FFmpeg path below as an experiment. It captures the screen and crops
	// by rectangle, so covered windows can leak into the preview. Tauri preview uses
	// screencapture -l instead, which captures the WeChat window itself.
	ffmpegErr := make(chan error, 1)
	ffmpegCtx, cancelFFmpeg := context.WithCancel(ctx)
	go func() {
		ffmpegErr <- runFFmpegMJPEG(ffmpegCtx, state, input, fps, window)
	}()

	timer := time.NewTimer(3 * time.Second)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		cancelFFmpeg()
		return nil
	case err := <-ffmpegErr:
		cancelFFmpeg()
		if err != nil {
			logWarning("FFmpeg exited before first frame, fallback to screencapture: " + err.Error())
		}
		return runScreencaptureMJPEG(ctx, state, fps, window, window.ID)
	case <-timer.C:
		_, _, sequence, _, _, _, _ := state.snapshot()
		if sequence > 0 {
			logInfo("FFmpeg produced frames, keep using FFmpeg capture")
			return <-ffmpegErr
		}
		cancelFFmpeg()
		logWarning("FFmpeg produced no frames in 3s, fallback to screencapture")
		return runScreencaptureMJPEG(ctx, state, fps, window, window.ID)
	}
}

func maybeFindWeChatWindow(cropWechat bool) (*WindowBounds, error) {
	if !cropWechat {
		return nil, nil
	}
	bounds, err := findWeChatWindow()
	if err != nil {
		return nil, fmt.Errorf("find WeChat window failed, fallback to full screen capture: %w", err)
	}
	logInfo(fmt.Sprintf(
		"WeChat window found owner=%s title=%s id=%d rect=[%.0f,%.0f,%.0f,%.0f]",
		bounds.Owner,
		bounds.Title,
		bounds.ID,
		bounds.X,
		bounds.Y,
		bounds.Width,
		bounds.Height,
	))
	return &bounds, nil
}

func defaultRegionsPath() string {
	home, err := os.UserHomeDir()
	if err != nil {
		return "data/regions/wechat.json"
	}
	return filepath.Join(home, "Library", "Application Support", "com.tauri.dev", "data", "regions", "wechat.json")
}

func previewTargetFromAnnotation(path string) (previewTarget, annotationMeta, error) {
	annotation, err := loadAnnotation(path)
	if err != nil {
		return previewTarget{}, annotationMeta{}, fmt.Errorf("load annotation source failed: %w", err)
	}
	meta := annotationMeta{
		Width:       annotation.ScreenshotWidth,
		Height:      annotation.ScreenshotHeight,
		ScaleFactor: annotation.ScaleFactor,
	}

	source := captureSourceFromAnnotation(annotation.Source)
	if source.Kind == "" {
		target, err := fallbackWeChatPreviewTarget()
		return target, meta, err
	}

	if source.Kind == "window" {
		window, err := findWindowForSource(source)
		if err != nil {
			return previewTarget{}, meta, err
		}
		return previewTarget{
			Kind:     "window",
			Label:    source.Label,
			WindowID: window.ID,
			Rect:     window,
		}, meta, nil
	}

	if source.Kind == "display" {
		return previewTarget{
			Kind:  "display",
			Label: source.Label,
			Rect: WindowBounds{
				X:      source.X,
				Y:      source.Y,
				Width:  source.Width,
				Height: source.Height,
				Owner:  "display",
				Title:  source.Label,
			},
		}, meta, nil
	}

	return previewTarget{}, meta, fmt.Errorf("unsupported annotation source kind: %s", source.Kind)
}

func fallbackWeChatPreviewTarget() (previewTarget, error) {
	window, err := findWeChatWindow()
	if err != nil {
		return previewTarget{}, err
	}
	return previewTarget{
		Kind:     "window",
		Label:    "微信",
		WindowID: window.ID,
		Rect:     window,
	}, nil
}

func captureSourceFromAnnotation(raw jsonRawOptional) captureSource {
	data, err := json.Marshal(raw)
	if err != nil {
		return captureSource{}
	}
	var source captureSource
	if err := json.Unmarshal(data, &source); err != nil {
		return captureSource{}
	}
	return source
}

func findWindowForSource(source captureSource) (WindowBounds, error) {
	appName := strings.ReplaceAll(source.AppName, "\\", "\\\\")
	appName = strings.ReplaceAll(appName, `"`, `\"`)
	title := strings.ReplaceAll(source.Title, "\\", "\\\\")
	title = strings.ReplaceAll(title, `"`, `\"`)

	script := fmt.Sprintf(`
import CoreGraphics
import Foundation

let targetOwner = "%s"
let targetTitle = "%s"
let options = CGWindowListOption(arrayLiteral: [.optionOnScreenOnly, .excludeDesktopElements])
guard let infos = CGWindowListCopyWindowInfo(options, kCGNullWindowID) as? [[String: Any]] else { exit(1) }

func printWindow(_ info: [String: Any]) {
  let owner = info[kCGWindowOwnerName as String] as? String ?? ""
  let title = info[kCGWindowName as String] as? String ?? ""
  let bounds = info[kCGWindowBounds as String] as? [String: Any] ?? [:]
  let x = bounds["X"] as? Double ?? 0
  let y = bounds["Y"] as? Double ?? 0
  let w = bounds["Width"] as? Double ?? 0
  let h = bounds["Height"] as? Double ?? 0
  let id = info[kCGWindowNumber as String] as? Int ?? 0
  let result: [String: Any] = ["x": x, "y": y, "width": w, "height": h, "owner": owner, "title": title, "id": id]
  let data = try! JSONSerialization.data(withJSONObject: result)
  print(String(data: data, encoding: .utf8)!)
}

for info in infos {
  let owner = info[kCGWindowOwnerName as String] as? String ?? ""
  let title = info[kCGWindowName as String] as? String ?? ""
  let layer = info[kCGWindowLayer as String] as? Int ?? -1
  if layer == 0 && owner == targetOwner && title == targetTitle { printWindow(info); exit(0) }
}

for info in infos {
  let owner = info[kCGWindowOwnerName as String] as? String ?? ""
  let title = info[kCGWindowName as String] as? String ?? ""
  let layer = info[kCGWindowLayer as String] as? Int ?? -1
  if layer == 0 && targetOwner != "" && owner.contains(targetOwner) { printWindow(info); exit(0) }
}

exit(2)
`, appName, title)

	output, err := runSwift(script)
	if err != nil {
		return WindowBounds{}, fmt.Errorf("find annotation window source failed app=%q title=%q: %w", source.AppName, source.Title, err)
	}
	var window WindowBounds
	if err := json.Unmarshal([]byte(output), &window); err != nil {
		return WindowBounds{}, err
	}
	return window, nil
}

func runFFmpegMJPEG(ctx context.Context, state *visionFrameState, input string, fps float64, window *WindowBounds) error {
	args := []string{
		"-hide_banner",
		"-loglevel", "error",
		"-f", "avfoundation",
		"-framerate", fmt.Sprintf("%.3f", fps),
		"-pixel_format", "bgr0",
		"-i", input,
	}
	if window != nil {
		args = append(args, "-vf", ffmpegCropFilter(*window))
	}
	args = append(args,
		"-f", "image2pipe",
		"-vcodec", "mjpeg",
		"-",
	)
	logInfo("FFmpeg command: ffmpeg " + strings.Join(args, " "))
	command := exec.CommandContext(ctx, "ffmpeg", args...)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		return err
	}

	if err := command.Start(); err != nil {
		return fmt.Errorf("启动 FFmpeg 失败：%w", err)
	}

	go drainFFmpegStderr(stderr)

	readErr := readJPEGFrames(stdout, state.update)
	waitErr := command.Wait()
	if ctx.Err() != nil {
		return nil
	}
	if readErr != nil {
		return readErr
	}
	if waitErr != nil {
		return fmt.Errorf("FFmpeg 退出：%w", waitErr)
	}
	return nil
}

func runScreencaptureMJPEG(ctx context.Context, state *visionFrameState, fps float64, window *WindowBounds, windowID int) error {
	if window == nil {
		return errors.New("screencapture requires a capture target")
	}
	interval := time.Duration(float64(time.Second) / math.Max(fps, 1))
	if interval < 80*time.Millisecond {
		interval = 80 * time.Millisecond
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	tempDir, err := os.MkdirTemp("", "wechat-go-vision-*")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)
	imagePath := filepath.Join(tempDir, "frame.jpg")
	rect := captureRect(*window)
	if windowID > 0 {
		logInfo(fmt.Sprintf("window-only screencapture started window_id=%d", windowID))
	} else {
		logInfo("rectangle screencapture started rect=" + rect)
	}
	for ctx.Err() == nil {
		args := []string{"-x"}
		if windowID > 0 {
			args = append(args, "-l", fmt.Sprintf("%d", windowID))
		} else {
			args = append(args, "-R", rect)
		}
		args = append(args, imagePath)
		command := exec.CommandContext(ctx, "screencapture", args...)
		if output, err := command.CombinedOutput(); err != nil {
			if ctx.Err() != nil {
				return nil
			}
			logWarning(fmt.Sprintf("screencapture failed: %v output=%s", err, string(output)))
		} else if jpeg, err := os.ReadFile(imagePath); err != nil {
			logWarning("read screencapture frame failed: " + err.Error())
		} else if len(jpeg) > 0 {
			state.update(jpeg)
		}

		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
		}
	}
	return nil
}

func ffmpegCropFilter(window WindowBounds) string {
	width := int(math.Max(1, math.Round(window.Width)))
	height := int(math.Max(1, math.Round(window.Height)))
	x := int(math.Max(0, math.Round(window.X)))
	y := int(math.Max(0, math.Round(window.Y)))
	return fmt.Sprintf("crop=%d:%d:%d:%d", width, height, x, y)
}

func captureRect(window WindowBounds) string {
	x := int(math.Round(window.X))
	y := int(math.Round(window.Y))
	width := int(math.Round(window.Width))
	height := int(math.Round(window.Height))
	if x < 0 {
		width += x
		x = 0
	}
	if y < 0 {
		height += y
		y = 0
	}
	if width < 1 {
		width = 1
	}
	if height < 1 {
		height = 1
	}
	return fmt.Sprintf("%d,%d,%d,%d", x, y, width, height)
}

func readJPEGFrames(reader io.Reader, onFrame func([]byte)) error {
	buffer := make([]byte, 0, 1024*1024)
	chunk := make([]byte, 64*1024)
	for {
		readCount, err := reader.Read(chunk)
		if readCount > 0 {
			buffer = append(buffer, chunk[:readCount]...)
			for {
				start := bytes.Index(buffer, []byte{0xff, 0xd8})
				if start < 0 {
					if len(buffer) > 1024*1024 {
						buffer = buffer[len(buffer)-2:]
					}
					break
				}
				endOffset := bytes.Index(buffer[start+2:], []byte{0xff, 0xd9})
				if endOffset < 0 {
					if start > 0 {
						buffer = buffer[start:]
					}
					break
				}
				end := start + 2 + endOffset + 2
				frame := append([]byte(nil), buffer[start:end]...)
				buffer = buffer[end:]
				onFrame(frame)
			}
		}
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
	}
}

func jpegSize(jpeg []byte) (int, int) {
	config, _, err := image.DecodeConfig(bytes.NewReader(jpeg))
	if err != nil {
		return 0, 0
	}
	return config.Width, config.Height
}

func drainFFmpegStderr(reader io.Reader) {
	scanner := bufio.NewScanner(reader)
	for scanner.Scan() {
		logWarning("ffmpeg: " + scanner.Text())
	}
}