nsfwapp/backend/sse.go

// backend\sse.go
package main

import (
	"bytes"
	"encoding/json"
	"net/http"
	"path/filepath"
	"sort"
	"strings"
	"sync/atomic"
	"time"

	"github.com/r3labs/sse/v2"
)

// -------------------- SSE server --------------------

type appSSE struct {
	server *sse.Server
}

var sseApp *appSSE

type jobEvent struct {
	Type             string    `json:"type"`
	Model            string    `json:"model"`
	JobID            string    `json:"jobId"`
	Status           JobStatus `json:"status"`
	Phase            string    `json:"phase,omitempty"`
	Progress         int       `json:"progress,omitempty"`
	SourceURL        string    `json:"sourceUrl,omitempty"`
	Output           string    `json:"output,omitempty"`
	StartedAt        string    `json:"startedAt,omitempty"`
	StartedAtMs      int64     `json:"startedAtMs,omitempty"`
	EndedAt          string    `json:"endedAt,omitempty"`
	EndedAtMs        int64     `json:"endedAtMs,omitempty"`
	SizeBytes        int64     `json:"sizeBytes,omitempty"`
	DurationSeconds  float64   `json:"durationSeconds,omitempty"`
	PreviewState     string    `json:"previewState,omitempty"`
	RoomStatus       string    `json:"roomStatus,omitempty"`
	IsOnline         bool      `json:"isOnline,omitempty"`
	ModelImageURL    string    `json:"modelImageUrl,omitempty"`
	ModelChatRoomURL string    `json:"modelChatRoomUrl,omitempty"`

	PostWorkKey string `json:"postWorkKey,omitempty"`
	PostWork    any    `json:"postWork,omitempty"`

	TS int64 `json:"ts"`
}

type ssePublishItem struct {
	EventName string
	CacheKey  string
	Data      []byte
}

type finishedPostworkEvent struct {
	Type        string `json:"type"`  // "finished_postwork"
	Model       string `json:"model"` // model event name, z. B. "maypeach"
	File        string `json:"file"`
	AssetID     string `json:"assetId,omitempty"`
	Queue       string `json:"queue"` // "postwork" | "enrich"
	State       string `json:"state"` // "queued" | "running" | "done" | "error" | "missing"
	Phase       string `json:"phase,omitempty"`
	Label       string `json:"label,omitempty"`
	Error       string `json:"error,omitempty"`
	Position    int    `json:"position,omitempty"`
	Waiting     int    `json:"waiting,omitempty"`
	Running     int    `json:"running,omitempty"`
	MaxParallel int    `json:"maxParallel,omitempty"`
	TS          int64  `json:"ts"`
}

type taskStateEvent struct {
	Type                    string `json:"type"` // "task_state"
	GenerateAssetsRunning   bool   `json:"generateAssetsRunning"`
	RegenerateAssetsRunning bool   `json:"regenerateAssetsRunning"`
	CheckVideosRunning      bool   `json:"checkVideosRunning"`
	CleanupRunning          bool   `json:"cleanupRunning"`
	AnyRunning              bool   `json:"anyRunning"`
	TS                      int64  `json:"ts"`
}

func publishFinishedPostworkStateForJob(
	j *RecordJob,
	queue string,
	state string,
	phase string,
	label string,
	position int,
	waiting int,
	running int,
	maxParallel int,
) {
	if j == nil {
		return
	}

	file := strings.TrimSpace(filepath.Base(j.Output))
	if file == "" {
		return
	}

	assetID := assetIDFromVideoPath(j.Output)

	publishFinishedPostworkStatusForJob(j, finishedPostworkEvent{
		File:        file,
		AssetID:     assetID,
		Queue:       queue,
		State:       state,
		Phase:       phase,
		Label:       label,
		Position:    position,
		Waiting:     waiting,
		Running:     running,
		MaxParallel: maxParallel,
	})
}

func anyGenerateAssetsTaskRunning() bool {
	for _, st := range snapshotAssetsJobs() {
		if st.Queued || st.Running {
			return true
		}
	}
	return false
}

func anyRegenerateAssetsTaskRunning() bool {
	for _, job := range getRegenerateAssetsJobsSnapshot() {
		state := strings.TrimSpace(strings.ToLower(anyToString(job["state"])))
		if state == "queued" || state == "running" {
			return true
		}
	}
	return false
}

func anyCheckVideosTaskRunning() bool {
	for _, st := range snapshotCheckVideosJobs() {
		if st.Queued || st.Running {
			return true
		}
	}
	return false
}

func anyCleanupTaskRunning() bool {
	for _, st := range snapshotCleanupJobs() {
		if st.Queued || st.Running {
			return true
		}
	}
	return false
}

func publishTaskState() {
	assetsRunning := anyGenerateAssetsTaskRunning()
	regenerateRunning := anyRegenerateAssetsTaskRunning()
	checkVideosRunning := anyCheckVideosTaskRunning()
	cleanupRunning := anyCleanupTaskRunning()

	ev := taskStateEvent{
		Type:                    "task_state",
		GenerateAssetsRunning:   assetsRunning,
		RegenerateAssetsRunning: regenerateRunning,
		CheckVideosRunning:      checkVideosRunning,
		CleanupRunning:          cleanupRunning,
		AnyRunning:              assetsRunning || regenerateRunning || checkVideosRunning || cleanupRunning,
		TS:                      time.Now().UnixMilli(),
	}

	b, _ := json.Marshal(ev)
	publishSSE("taskState", b)
}

func publishFinishedPostworkStatusForJob(j *RecordJob, ev finishedPostworkEvent) {
	if j == nil {
		return
	}
	if strings.TrimSpace(ev.File) == "" {
		return
	}

	eventName := sseModelEventNameForJob(j)
	if eventName == "" {
		return
	}

	ev.Type = "finished_postwork"
	ev.Model = eventName
	ev.TS = time.Now().UnixMilli()

	b, _ := json.Marshal(ev)

	// nur global für FinishedDownloads
	publishSSE("finishedPostwork", b)
}

func isTerminalJobStatusForSSE(status JobStatus) bool {
	s := strings.ToLower(strings.TrimSpace(string(status)))
	return s == "stopped" ||
		s == "finished" ||
		s == "failed" ||
		s == "done" ||
		s == "completed" ||
		s == "canceled" ||
		s == "cancelled"
}

func isPostworkJobForSSE(j *RecordJob) bool {
	if j == nil {
		return false
	}

	phase := strings.ToLower(strings.TrimSpace(j.Phase))
	pwKey := strings.TrimSpace(j.PostWorkKey)

	if pwKey != "" {
		return true
	}

	if j.PostWork != nil {
		// falls PostWork als struct/map vorliegt, reicht für SSE der generelle Hinweis:
		return true
	}

	if j.EndedAt != nil && phase != "" {
		return true
	}

	if phase == "postwork" {
		return true
	}

	return false
}

func isVisibleDownloadJobForSSE(j *RecordJob) bool {
	if j == nil {
		return false
	}
	if isPostworkJobForSSE(j) {
		return false
	}
	if isTerminalJobStatusForSSE(j.Status) {
		return false
	}
	if j.EndedAt != nil {
		return false
	}
	return true
}

func isVisiblePostworkJobForSSE(j *RecordJob) bool {
	if j == nil {
		return false
	}
	if !isPostworkJobForSSE(j) {
		return false
	}
	if isTerminalJobStatusForSSE(j.Status) {
		return false
	}
	return true
}

func shouldPublishModelEventForJob(j *RecordJob) bool {
	return isVisibleDownloadJobForSSE(j) || isVisiblePostworkJobForSSE(j)
}

func visibleJobEventsJSON() []ssePublishItem {
	out := make([]ssePublishItem, 0, 64)

	jobsMu.RLock()
	defer jobsMu.RUnlock()

	for _, j := range jobs {
		if j == nil || j.Hidden {
			continue
		}
		if !shouldPublishModelEventForJob(j) {
			continue
		}

		eventName := sseModelEventNameForJob(j)
		if eventName == "" {
			continue
		}

		payload := jobEvent{
			Type:            "job_upsert",
			Model:           eventName,
			JobID:           j.ID,
			Status:          j.Status,
			Phase:           j.Phase,
			Progress:        j.Progress,
			SourceURL:       j.SourceURL,
			Output:          j.Output,
			StartedAt:       j.StartedAt.Format(time.RFC3339Nano),
			StartedAtMs:     j.StartedAtMs,
			SizeBytes:       j.SizeBytes,
			DurationSeconds: j.DurationSeconds,
			PreviewState:    j.PreviewState,
			PostWorkKey:     strings.TrimSpace(j.PostWorkKey),
			PostWork:        j.PostWork,
			TS:              0, // wichtig: stabil halten, damit Dedupe funktioniert
		}

		if sm := sseStoredModelForJob(j); sm != nil {
			payload.RoomStatus = strings.ToLower(strings.TrimSpace(sm.RoomStatus))
			payload.IsOnline = sm.IsOnline
			payload.ModelImageURL = strings.TrimSpace(sm.ImageURL)
			payload.ModelChatRoomURL = strings.TrimSpace(sm.ChatRoomURL)
		}

		if j.EndedAt != nil {
			payload.EndedAt = j.EndedAt.Format(time.RFC3339Nano)
			payload.EndedAtMs = j.EndedAtMs
		}

		b, err := json.Marshal(payload)
		if err != nil {
			continue
		}

		out = append(out, ssePublishItem{
			EventName: eventName,
			CacheKey:  eventName + "|" + j.ID,
			Data:      b,
		})
	}

	return out
}

func publishJobRemove(j *RecordJob) {
	if j == nil || j.Hidden {
		return
	}

	eventName := sseModelEventNameForJob(j)
	if eventName == "" {
		return
	}

	b, _ := json.Marshal(jobEvent{
		Type:  "job_remove",
		Model: eventName,
		JobID: j.ID,
		TS:    time.Now().UnixMilli(),
	})

	publishSSE(eventName, b)
}

func sseModelEventNameForJob(j *RecordJob) string {
	if j == nil {
		return ""
	}

	src := strings.TrimSpace(j.SourceURL)

	switch detectProvider(src) {
	case "chaturbate":
		if u := strings.TrimSpace(extractUsername(src)); u != "" {
			return strings.ToLower(u)
		}
	case "mfc":
		if u := strings.TrimSpace(extractMFCUsername(src)); u != "" {
			return strings.ToLower(u)
		}
	}

	return ""
}

func sseStoredModelForJob(j *RecordJob) *StoredModel {
	if j == nil || cbModelStore == nil {
		return nil
	}

	src := strings.TrimSpace(j.SourceURL)
	if src == "" {
		return nil
	}

	host := ""
	modelKey := ""

	switch detectProvider(src) {
	case "chaturbate":
		host = "chaturbate.com"
		modelKey = strings.ToLower(strings.TrimSpace(extractUsername(src)))
	default:
		return nil
	}

	if host == "" || modelKey == "" {
		return nil
	}

	m, ok := cbModelStore.GetByHostAndModelKey(host, modelKey)
	if !ok {
		return nil
	}

	return m
}

func initSSE() {
	srv := sse.New()
	srv.SplitData = true

	stream := srv.CreateStream("events")
	stream.AutoReplay = false

	sseApp = &appSSE{
		server: srv,
	}

	// Debounced broadcaster (done changed)
	go func() {
		for range doneNotify {
			time.Sleep(40 * time.Millisecond)
			for {
				select {
				case <-doneNotify:
				default:
					goto SEND
				}
			}
		SEND:
			seq := atomic.AddUint64(&doneSeq, 1)
			b, _ := json.Marshal(map[string]any{
				"type": "doneChanged",
				"seq":  seq,
				"ts":   time.Now().UnixMilli(),
			})
			publishSSE("doneChanged", b)
		}
	}()

	// Debounced broadcaster (assets)
	go func() {
		for range assetsNotify {
			time.Sleep(80 * time.Millisecond)
			for {
				select {
				case <-assetsNotify:
				default:
					goto SEND
				}
			}
		SEND:
			b := assetsSnapshotJSON()
			if len(b) > 0 {
				publishSSE("state", b)
			}
		}
	}()

	// Per running job: 1 SSE event per second
	go func() {
		t := time.NewTicker(2 * time.Second)
		defer t.Stop()

		lastByKey := map[string][]byte{}
		lastHeartbeat := time.Now()

		for range t.C {
			events := visibleJobEventsJSON()
			liveKeys := make(map[string]struct{}, len(events))
			forceHeartbeat := time.Since(lastHeartbeat) >= 15*time.Second

			for _, ev := range events {
				if len(ev.Data) == 0 || ev.EventName == "" || ev.CacheKey == "" {
					continue
				}

				liveKeys[ev.CacheKey] = struct{}{}

				prev, ok := lastByKey[ev.CacheKey]
				if !forceHeartbeat && ok && bytes.Equal(prev, ev.Data) {
					continue
				}

				lastByKey[ev.CacheKey] = append([]byte(nil), ev.Data...)
				publishSSE(ev.EventName, ev.Data)
			}

			for k := range lastByKey {
				if _, ok := liveKeys[k]; !ok {
					delete(lastByKey, k)
				}
			}

			if forceHeartbeat {
				lastHeartbeat = time.Now()
			}
		}
	}()
}

func publishSSE(eventName string, data []byte) {
	if sseApp == nil || sseApp.server == nil {
		return
	}
	if len(data) == 0 {
		return
	}

	sseApp.server.Publish("events", &sse.Event{
		Event: []byte(eventName),
		Data:  data,
	})
}

// -------------------- notify channels --------------------

var (
	doneNotify = make(chan struct{}, 1)
	doneSeq    uint64

	assetsNotify = make(chan struct{}, 1)
)

func notifyDoneChanged() {
	select {
	case doneNotify <- struct{}{}:
	default:
	}
}

func notifyAssetsChanged() {
	select {
	case assetsNotify <- struct{}{}:
	default:
	}
}

// -------------------- snapshots --------------------

func jobsSnapshotJSON() []byte {
	jobsMu.RLock()
	list := make([]*RecordJob, 0, len(jobs))
	for _, j := range jobs {
		if j == nil || j.Hidden {
			continue
		}
		c := *j
		c.cancel = nil
		list = append(list, &c)
	}
	jobsMu.RUnlock()

	sort.Slice(list, func(i, j int) bool {
		return list[i].StartedAt.After(list[j].StartedAt)
	})

	b, _ := json.Marshal(list)
	return b
}

func assetsSnapshotJSON() []byte {
	b, _ := json.Marshal(snapshotAssetsJobs())
	return b
}

// -------------------- unified stream handler --------------------

func eventsStream(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Nur GET erlaubt", http.StatusMethodNotAllowed)
		return
	}

	if sseApp == nil || sseApp.server == nil {
		http.Error(w, "SSE nicht initialisiert", http.StatusInternalServerError)
		return
	}

	q := r.URL.Query()
	q.Set("stream", "events")

	r2 := r.Clone(r.Context())
	r2.URL.RawQuery = q.Encode()

	sseApp.server.ServeHTTP(w, r2)
}

// -------------------- optional helper --------------------

func publishRawSSE(eventName string, buf *bytes.Buffer) {
	if buf == nil {
		return
	}
	publishSSE(eventName, buf.Bytes())
}