AI-tianyan/internal/stream/ingestor.go

package stream

import (
	"context"
	"fmt"
	"io"
	"log"
	"os/exec"
	"sync"
	"time"

	"bytes"

	"tianyan-edge/internal/config"
)

type Frame struct {
	StreamID int
	DeviceID string
	URL      string
	JPEG     []byte
	TS       float64
}

type streamProcess struct {
	cmd *exec.Cmd
	fps int
}

type StreamManager struct {
	cfg       *config.Config
	puller    *Puller // 可选的动态拉流客户端
	frames    chan<- Frame
	processes map[string]*streamProcess // url -> {cmd, fps}
	mu        sync.Mutex
	updates   <-chan *config.Config
}

func NewStreamManager(cfg *config.Config, puller *Puller, frames chan<- Frame, updates <-chan *config.Config) *StreamManager {
	return &StreamManager{
		cfg:       cfg,
		puller:    puller,
		frames:    frames,
		processes: make(map[string]*streamProcess),
		updates:   updates,
	}
}

func (sm *StreamManager) Run(ctx context.Context) {
	// 1. 动态拉流模式
	if sm.puller != nil {
		log.Println("stream: running in auto-pull mode")
		sm.runAutoPullMode(ctx)
		return
	}

	// 2. 静态配置模式
	sm.applyStreams(ctx, sm.cfg.RTSPURLs, sm.cfg.InferFPS)

	for {
		select {
		case <-ctx.Done():
			sm.stopAll()
			return
		case newCfg := <-sm.updates:
			log.Printf("stream: applying new config, urls=%d fps=%d", len(newCfg.RTSPURLs), newCfg.InferFPS)
			sm.applyStreams(ctx, newCfg.RTSPURLs, newCfg.InferFPS)
		}
	}
}

// runAutoPullMode 动态拉流模式主循环
func (sm *StreamManager) runAutoPullMode(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			sm.stopAll()
			return
		default:
		}

		// 从云端获取拉流地址
		url, err := sm.puller.Pull()
		if err != nil {
			log.Printf("stream: pull failed: %v, retry in 10s", err)
			time.Sleep(10 * time.Second)
			continue
		}

		log.Printf("stream: got dynamic url=%s, starting ingestion", url)
		
		// 启动该 URL 的拉流
		sm.applyStreams(ctx, []string{url}, sm.cfg.InferFPS)

		// 等待上下文结束或 URL 被移除（通常意味着需要重新 Pull）
		// 在 auto-pull 模式下，如果 ffmpeg 断开，streamLoop 会自动重试，
		// 但我们需要在这里监听 ctx，或者等待 puller 心跳失败后重新拉取。
		// 简单起见，这里阻塞直到 ctx 取消，streamLoop 内部会处理 ffmpeg 重连。
		// 如果 streamLoop 发现 URL 失效，会退出并删除 process。
		<-ctx.Done()
		sm.stopAll()
	}
}

func (sm *StreamManager) applyStreams(ctx context.Context, urls []string, fps int) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	// Stream enabled check
	enabled := true
	if sm.cfg != nil && !sm.cfg.StreamEnabled {
		enabled = false
	}

	// If disabled, stop all streams and return
	if !enabled {
		log.Println("stream: disabled by config, stopping all")
		for url, sp := range sm.processes {
			log.Printf("stream: stopping %s", url)
			sm.stopProcess(sp)
			delete(sm.processes, url)
		}
		return
	}

	// Identify URLs to remove (present in sm.processes but not in new urls)
	toRemove := map[string]*streamProcess{}
	for url, sp := range sm.processes {
		found := false
		for _, u := range urls {
			if u == url {
				found = true
				break
			}
		}
		if !found {
			toRemove[url] = sp
		}
	}

	// Stop removed streams
	for url, sp := range toRemove {
		log.Printf("stream: stopping %s", url)
		sm.stopProcess(sp)
		delete(sm.processes, url)
	}

	// Restart streams with changed fps
	for url, sp := range sm.processes {
		if sp.fps != fps {
			log.Printf("stream: restarting %s (fps %d -> %d)", url, sp.fps, fps)
			sm.stopProcess(sp)
			delete(sm.processes, url)
		}
	}

	// Start new / restarted streams
	for i, url := range urls {
		if _, exists := sm.processes[url]; !exists {
			go sm.streamLoop(ctx, i, url, fps)
		}
	}
}

func (sm *StreamManager) streamLoop(ctx context.Context, idx int, url string, fps int) {
	deviceID := fmt.Sprintf("cam-%03d", idx)
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		// Use current config fps (may have been updated)
		currentFps := fps
		sm.mu.Lock()
		if sm.cfg != nil && sm.cfg.InferFPS > 0 {
			currentFps = sm.cfg.InferFPS
		}
		sm.mu.Unlock()

		cmd := exec.CommandContext(ctx, "ffmpeg",
			"-i", url,
			"-vf", fmt.Sprintf("fps=%d", currentFps),
			"-f", "image2pipe", "-vcodec", "mjpeg", "-q:v", "5", "-",
		)
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Printf("stream[%d] stdout pipe failed: %v", idx, err)
			time.Sleep(5 * time.Second)
			continue
		}

		sp := &streamProcess{cmd: cmd, fps: currentFps}
		sm.mu.Lock()
		sm.processes[url] = sp
		sm.mu.Unlock()

		if err := cmd.Start(); err != nil {
			log.Printf("stream[%d] start failed: %v", idx, err)
			sm.mu.Lock()
			delete(sm.processes, url)
			sm.mu.Unlock()
			time.Sleep(5 * time.Second)
			continue
		}

		log.Printf("stream[%d] connected fps=%d url=%s", idx, currentFps, url)
		sm.readFrames(ctx, stdout, idx, deviceID, url)
		cmd.Wait()

		// Clean up from map after exit
		sm.mu.Lock()
		delete(sm.processes, url)
		sm.mu.Unlock()

		select {
		case <-ctx.Done():
			return
		default:
			log.Printf("stream[%d] lost, retry 5s", idx)
			time.Sleep(5 * time.Second)
		}
	}
}

func (sm *StreamManager) readFrames(ctx context.Context, r io.Reader, idx int, deviceID, url string) {
	soi, eoi := []byte{0xFF, 0xD8}, []byte{0xFF, 0xD9}
	buf, tmp := make([]byte, 0, 1<<20), make([]byte, 32768)
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		n, err := r.Read(tmp)
		if n > 0 {
			buf = append(buf, tmp[:n]...)
			for {
				s := bytes.Index(buf, soi)
				if s < 0 {
					buf = buf[:0]
					break
				}
				buf = buf[s:]
				e := bytes.Index(buf[2:], eoi)
				if e < 0 {
					break
				}
				end := e + 4
				frame := make([]byte, end)
				copy(frame, buf[:end])
				buf = buf[end:]
				select {
				case sm.frames <- Frame{StreamID: idx, DeviceID: deviceID, URL: url,
					JPEG: frame, TS: float64(time.Now().UnixMilli()) / 1000.0}:
				default:
				}
			}
		}
		if err != nil {
			return
		}
	}
}

func (sm *StreamManager) stopProcess(sp *streamProcess) {
	if sp == nil || sp.cmd == nil || sp.cmd.Process == nil {
		return
	}
	sp.cmd.Process.Kill()
	sp.cmd.Wait() // Prevent zombie processes
}

func (sm *StreamManager) stopAll() {
	sm.mu.Lock()
	defer sm.mu.Unlock()
	for url, sp := range sm.processes {
		log.Printf("stream: stopping all %s", url)
		sm.stopProcess(sp)
		delete(sm.processes, url)
	}
}