wechat_ai/demos/go-ffmpeg-wechat-stream/main.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"
)

type windowBounds struct {
	X      float64 `json:"x"`
	Y      float64 `json:"y"`
	Width  float64 `json:"width"`
	Height float64 `json:"height"`
	Owner  string  `json:"owner"`
	Title  string  `json:"title"`
	ID     int     `json:"id"`
}

type streamState struct {
	condition *sync.Cond
	jpeg      []byte
	sequence  int64
	window    windowBounds
	startedAt time.Time
	updatedAt time.Time
}

func newStreamState(window windowBounds) *streamState {
	return &streamState{
		condition: sync.NewCond(&sync.Mutex{}),
		window:    window,
		startedAt: time.Now(),
	}
}

func (state *streamState) update(jpeg []byte) {
	state.condition.L.Lock()
	state.jpeg = append(state.jpeg[:0], jpeg...)
	state.sequence++
	state.updatedAt = time.Now()
	state.condition.Broadcast()
	state.condition.L.Unlock()
}

func (state *streamState) snapshot() (int64, time.Time, windowBounds) {
	state.condition.L.Lock()
	defer state.condition.L.Unlock()
	return state.sequence, state.updatedAt, state.window
}

func (state *streamState) waitNext(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 main() {
	if runtime.GOOS != "darwin" {
		log.Fatal("this demo currently supports macOS only because it uses avfoundation + CoreGraphics")
	}

	var host string
	var port int
	var input string
	var fps float64
	var openBrowser bool
	var waitWindow bool
	flag.StringVar(&host, "host", "127.0.0.1", "HTTP server host")
	flag.IntVar(&port, "port", 8765, "HTTP server port")
	flag.StringVar(&input, "input", "3:none", "FFmpeg avfoundation input, e.g. 3:none")
	flag.Float64Var(&fps, "fps", 15, "capture frame rate")
	flag.BoolVar(&openBrowser, "open", true, "open browser automatically")
	flag.BoolVar(&waitWindow, "wait", true, "wait until a visible WeChat window is found")
	flag.Parse()

	window, err := waitForWeChatWindow(waitWindow)
	if err != nil {
		log.Fatalf("find WeChat window failed: %v", err)
	}
	log.Printf("WeChat window: owner=%q title=%q id=%d rect=[%.0f,%.0f,%.0f,%.0f]", window.Owner, window.Title, window.ID, window.X, window.Y, window.Width, window.Height)

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

	go runCapture(ctx, state, input, fps, window)

	server := &http.Server{
		Addr:              fmt.Sprintf("%s:%d", host, port),
		Handler:           handler(state),
		ReadHeaderTimeout: 5 * time.Second,
	}

	url := fmt.Sprintf("http://%s:%d/", host, port)
	log.Printf("Serving %s", url)
	if openBrowser {
		go func() {
			time.Sleep(500 * time.Millisecond)
			_ = exec.Command("open", url).Start()
		}()
	}

	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatal(err)
	}
}

func runCapture(ctx context.Context, state *streamState, input string, fps float64, window windowBounds) {
	ffmpegCtx, cancelFFmpeg := context.WithCancel(ctx)
	ffmpegErr := make(chan error, 1)
	go func() {
		ffmpegErr <- runFFmpeg(ffmpegCtx, state, input, fps, window)
	}()

	timer := time.NewTimer(3 * time.Second)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		cancelFFmpeg()
		return
	case err := <-ffmpegErr:
		cancelFFmpeg()
		log.Printf("FFmpeg stopped before first frame, fallback to screencapture: %v", err)
	case <-timer.C:
		sequence, _, _ := state.snapshot()
		if sequence > 0 {
			log.Printf("FFmpeg is producing frames")
			if err := <-ffmpegErr; err != nil && ctx.Err() == nil {
				log.Printf("FFmpeg stopped: %v", err)
			}
			return
		}
		cancelFFmpeg()
		log.Printf("FFmpeg produced no frames in 3s, fallback to window-only screencapture")
	}

	if err := runScreencapture(ctx, state, fps, window); err != nil && ctx.Err() == nil {
		log.Printf("screencapture stopped: %v", err)
	}
}

func waitForWeChatWindow(wait bool) (windowBounds, error) {
	for {
		window, err := findWeChatWindow()
		if err == nil {
			return window, nil
		}
		if !wait {
			return windowBounds{}, err
		}
		log.Printf("waiting for visible WeChat window: %v", err)
		time.Sleep(2 * time.Second)
	}
}

func findWeChatWindow() (windowBounds, error) {
	script := `
import CoreGraphics
import Foundation

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

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.contains("微信") || owner.lowercased().contains("wechat") || title == "微信") {
    if let bounds = info[kCGWindowBounds as String] as? [String: Any],
       let x = bounds["X"] as? Double,
       let y = bounds["Y"] as? Double,
       let w = bounds["Width"] as? Double,
       let h = bounds["Height"] as? Double {
      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)!)
      exit(0)
    }
  }
}

exit(2)
`
	command := exec.Command("swift", "-")
	command.Stdin = strings.NewReader(script)
	output, err := command.Output()
	if err != nil {
		return windowBounds{}, err
	}

	var window windowBounds
	if err := json.Unmarshal(output, &window); err != nil {
		return windowBounds{}, err
	}
	return window, nil
}

func runFFmpeg(ctx context.Context, state *streamState, input string, fps float64, window windowBounds) error {
	args := []string{
		"-hide_banner",
		"-loglevel", "warning",
		"-f", "avfoundation",
		"-framerate", fmt.Sprintf("%.3f", fps),
		"-pixel_format", "bgr0",
		"-i", input,
		"-vf", cropFilter(window),
		"-f", "image2pipe",
		"-vcodec", "mjpeg",
		"-q:v", "5",
		"-",
	}
	log.Printf("FFmpeg command: ffmpeg %s", 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 err
	}
	go logLines("ffmpeg", stderr)

	readErr := readJPEGFrames(stdout, state.update)
	waitErr := command.Wait()
	if ctx.Err() != nil {
		return nil
	}
	if readErr != nil {
		return readErr
	}
	return waitErr
}

func runScreencapture(ctx context.Context, state *streamState, fps float64, window windowBounds) error {
	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("", "go-ffmpeg-wechat-stream-*")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	framePath := filepath.Join(tempDir, "frame.jpg")
	rect := captureRect(window)
	log.Printf("window-only screencapture fallback started window_id=%d rect=%s", window.ID, rect)
	for ctx.Err() == nil {
		command := exec.CommandContext(ctx, "screencapture", "-x", "-l", fmt.Sprintf("%d", window.ID), framePath)
		if output, err := command.CombinedOutput(); err != nil {
			if ctx.Err() != nil {
				return nil
			}
			log.Printf("screencapture failed: %v output=%s", err, string(output))
		} else if jpeg, err := os.ReadFile(framePath); err != nil {
			log.Printf("read frame failed: %v", err)
		} else if len(jpeg) > 0 {
			state.update(jpeg)
		}

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

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 cropFilter(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 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
				onFrame(append([]byte(nil), buffer[start:end]...))
				buffer = buffer[end:]
			}
		}
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
	}
}

func handler(state *streamState) 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, `<!doctype html>
<html>
<head>
  <meta charset="utf-8" />
  <title>Go FFmpeg WeChat Stream</title>
  <style>
    html, body { margin: 0; height: 100%; background: #050505; color: #d7ffe4; font-family: ui-monospace, SFMono-Regular, Menlo, monospace; }
    .wrap { height: 100%; display: grid; grid-template-rows: auto 1fr; }
    .bar { display: flex; gap: 16px; align-items: center; padding: 10px 14px; background: #101510; border-bottom: 1px solid #263526; font-size: 13px; }
    img { width: 100%; height: 100%; object-fit: contain; background: #000; }
  </style>
</head>
<body>
  <div class="wrap">
    <div class="bar"><strong>Go FFmpeg WeChat Stream</strong><span id="status">connecting...</span></div>
    <img src="/stream.mjpg" />
  </div>
  <script>
    let prevSeq = null, prevTime = performance.now();
    async function tick() {
      try {
        const res = await fetch('/latest.json?t=' + Date.now(), {cache: 'no-store'});
        const data = await res.json();
        const now = performance.now();
        let fps = '--';
        if (prevSeq !== null && data.sequence >= prevSeq) fps = ((data.sequence - prevSeq) / ((now - prevTime) / 1000)).toFixed(1);
        prevSeq = data.sequence;
        prevTime = now;
        document.getElementById('status').textContent = 'seq=' + data.sequence + ' fps=' + fps + ' window=' + data.window.width + 'x' + data.window.height;
      } catch (err) {
        document.getElementById('status').textContent = 'waiting for stream';
      }
    }
    tick(); setInterval(tick, 1000);
  </script>
</body>
</html>`)
	})
	mux.HandleFunc("/latest.json", func(writer http.ResponseWriter, request *http.Request) {
		sequence, updatedAt, window := state.snapshot()
		writer.Header().Set("Content-Type", "application/json; charset=utf-8")
		writer.Header().Set("Cache-Control", "no-store")
		_ = json.NewEncoder(writer).Encode(map[string]any{
			"source":     "go-ffmpeg-demo",
			"sequence":   sequence,
			"updated_at": updatedAt.Format(time.RFC3339Nano),
			"window":     window,
		})
	})
	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.waitNext(lastSequence)
			lastSequence = sequence
			if len(jpeg) == 0 {
				continue
			}
			if _, err := fmt.Fprintf(writer, "--frame\r\nContent-Type: image/jpeg\r\nContent-Length: %d\r\n\r\n", len(jpeg)); 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 logLines(prefix string, reader io.Reader) {
	buffer := make([]byte, 4096)
	var line []byte
	for {
		count, err := reader.Read(buffer)
		if count > 0 {
			line = append(line, buffer[:count]...)
			for {
				index := bytes.IndexByte(line, '\n')
				if index < 0 {
					break
				}
				log.Printf("%s: %s", prefix, strings.TrimSpace(string(line[:index])))
				line = line[index+1:]
			}
		}
		if err != nil {
			if len(line) > 0 {
				log.Printf("%s: %s", prefix, strings.TrimSpace(string(line)))
			}
			return
		}
	}
}