AI-tianyan/internal/stream/puller.go

package stream

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"tianyan-edge/internal/config"
)

// Puller 负责从云端网关动态获取拉流地址并维护心跳
type Puller struct {
	cfg       *config.Config
	client    *http.Client
	sessionID string
	stopCh    chan struct{}
}

// PullResponse 云端网关返回结构
type PullResponse struct {
	Code int    `json:"code"`
	Msg  string `json:"msg"`
	Data struct {
		URL       string `json:"url"`
		Protocol  string `json:"protocol"`
		TTL       int    `json:"ttl"`
		DeviceID  string `json:"device_id"`
		ChannelID string `json:"channel_id"`
	} `json:"data"`
}

// HeartbeatPayload 心跳请求结构
type HeartbeatPayload struct {
	EdgeToken string `json:"edge_token"`
}

// PullPayload 拉流请求结构
type PullPayload struct {
	EdgeToken string `json:"edge_token"`
	Protocol  string `json:"protocol"`
}

func NewPuller(cfg *config.Config) *Puller {
	return &Puller{
		cfg: cfg,
		client: &http.Client{
			Timeout: 60 * time.Second, // WVP SIP 信令较慢，设置较长超时
		},
		stopCh: make(chan struct{}),
	}
}

// Pull 请求云端获取拉流 URL
func (p *Puller) Pull() (string, error) {
	if p.cfg.StreamPullURL == "" {
		return "", fmt.Errorf("stream_pull_url not configured")
	}

	payload := PullPayload{
		EdgeToken: p.cfg.EdgeToken,
		Protocol:  p.cfg.StreamProtocol,
	}
	body, err := json.Marshal(payload)
	if err != nil {
		return "", err
	}

	req, err := http.NewRequest("POST", p.cfg.StreamPullURL+"/api/v1/edge/pull", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Edge-Token", p.cfg.EdgeToken)

	resp, err := p.client.Do(req)
	if err != nil {
		return "", fmt.Errorf("pull request failed: %w", err)
	}
	defer resp.Body.Close()

	var pullResp PullResponse
	if err := json.NewDecoder(resp.Body).Decode(&pullResp); err != nil {
		return "", fmt.Errorf("decode pull response failed: %w", err)
	}

	if pullResp.Code != 0 {
		return "", fmt.Errorf("pull failed: %s", pullResp.Msg)
	}

	p.sessionID = p.cfg.EdgeToken
	log.Printf("puller: got stream url protocol=%s ttl=%ds", pullResp.Data.Protocol, pullResp.Data.TTL)
	return pullResp.Data.URL, nil
}

// StartHeartbeatLoop 启动后台心跳保活协程
func (p *Puller) StartHeartbeatLoop() {
	go func() {
		ticker := time.NewTicker(20 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-p.stopCh:
				log.Println("puller: heartbeat loop stopped")
				return
			case <-ticker.C:
				if p.sessionID == "" {
					continue
				}
				if err := p.sendHeartbeat(); err != nil {
					log.Printf("puller: heartbeat failed: %v (will retry next tick)", err)
				}
			}
		}
	}()
}

// sendHeartbeat 发送单次心跳
func (p *Puller) sendHeartbeat() error {
	payload := HeartbeatPayload{EdgeToken: p.cfg.EdgeToken}
	body, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	req, err := http.NewRequest("POST", p.cfg.StreamPullURL+"/api/v1/edge/heartbeat", bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Edge-Token", p.cfg.EdgeToken)

	resp, err := p.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	var respData struct {
		Code int    `json:"code"`
		Msg  string `json:"msg"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&respData); err != nil {
		return err
	}

	if respData.Code != 0 {
		return fmt.Errorf("heartbeat failed: %s", respData.Msg)
	}
	return nil
}

// Release 主动释放云端拉流会话
func (p *Puller) Release() {
	if p.sessionID == "" || p.cfg.StreamPullURL == "" {
		return
	}

	payload := HeartbeatPayload{EdgeToken: p.cfg.EdgeToken}
	body, err := json.Marshal(payload)
	if err != nil {
		log.Printf("puller: release marshal error: %v", err)
		return
	}

	req, err := http.NewRequest("POST", p.cfg.StreamPullURL+"/api/v1/edge/release", bytes.NewReader(body))
	if err != nil {
		log.Printf("puller: release request error: %v", err)
		return
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Edge-Token", p.cfg.EdgeToken)

	resp, err := p.client.Do(req)
	if err != nil {
		log.Printf("puller: release do error: %v", err)
		return
	}
	defer resp.Body.Close()

	var respData struct {
		Code int    `json:"code"`
		Msg  string `json:"msg"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&respData); err != nil {
		log.Printf("puller: release decode error: %v", err)
		return
	}

	if respData.Code == 0 {
		log.Println("puller: session released successfully")
	} else {
		log.Printf("puller: release failed: %s", respData.Msg)
	}
	p.sessionID = ""
}

// Stop 停止拉流管理器
func (p *Puller) Stop() {
	close(p.stopCh)
	p.Release()
}