streamdeck-obs-replay/replay.go

package main

import (
	"github.com/valyala/fastjson"
	"log"
	"meow.tf/streamdeck/obs-replay/obsws"
	"meow.tf/streamdeck/sdk"
	"strconv"
	"sync"
)

const (
	actionReplayToggle = "tf.meow.obsreplay.replay_toggle"
	actionReplaySave   = "tf.meow.obsreplay.replay_save"
)

var (
	contextMutex sync.RWMutex
	clientMutex  sync.RWMutex
	stateMutex   sync.RWMutex

	clients      = make(map[string]*obsws.Client)
	contexts     = make(map[string]string)
	cachedStates = make(map[string]int)
)

func replayToggle(action, context string, payload *fastjson.Value, deviceId string) {
	c := clientForContext(context)

	if c == nil {
		sdk.ShowAlert(context)
		return
	}

	req := obsws.NewStartStopReplayBufferRequest()

	err := req.Send(c)

	if err != nil {
		sdk.ShowAlert(context)
		return
	}

	sdk.ShowOk(context)
}

func replaySave(action, context string, payload *fastjson.Value, deviceId string) {
	c := clientForContext(context)

	if c == nil {
		sdk.ShowAlert(context)
		return
	}

	stateMutex.RLock()
	state, exists := cachedStates[context]
	stateMutex.RUnlock()

	if exists && state == 0 {
		sdk.ShowAlert(context)
		return
	}

	req := obsws.NewSaveReplayBufferRequest()

	if err := req.Send(c); err != nil {
		sdk.ShowAlert(context)
		return
	}

	sdk.ShowOk(context)
}

func clientForContext(context string) *obsws.Client {
	contextMutex.RLock()
	key, exists := contexts[context]
	contextMutex.RUnlock()

	if !exists {
		return nil
	}

	clientMutex.RLock()
	c, exists := clients[key]
	clientMutex.RUnlock()

	if !exists {
		return nil
	}

	if !c.Connected() {
		err := c.Connect()

		if err != nil {
			return nil
		}
	}

	return c
}

func onWillAppear(e *sdk.WillAppearEvent) {
	settings := e.Payload.Get("settings")

	if settings != nil {
		host := sdk.JsonStringValue(settings, "host")
		port := settings.GetInt("port")
		password := sdk.JsonStringValue(settings, "password")

		key := checkClient(host, port, password)

		if key == "" {
			return
		}

		contextMutex.Lock()
		contexts[e.Context] = key
		contextMutex.Unlock()

		stateMutex.RLock()
		if state, ok := cachedStates[key]; ok {
			sdk.SetState(e.Context, state)
		}
		stateMutex.RUnlock()
	}
}

func checkClient(host string, port int, password string) string {
	if port == 0 {
		port = 4444
	}

	key := host + ":" + strconv.Itoa(port)

	clientMutex.RLock()
	client, ok := clients[key]
	clientMutex.RUnlock()

	if !ok {
		client = &obsws.Client{Host: host, Port: port, Password: password}

		client.AddEventHandler("StreamStatus", streamStatusUpdate(key))
		client.AddEventHandler("ReplayStarted", loopContextState(key, 1))
		client.AddEventHandler("ReplayStopped", loopContextState(key, 0))

		clientMutex.Lock()
		clients[key] = client
		clientMutex.Unlock()
	}

	return key
}

func streamStatusUpdate(key string) func(obsws.Event) {
	return func(e obsws.Event) {
		evt := e.(obsws.StreamStatusEvent)

		state := 0

		if evt.Replay {
			state = 1
		}

		loopContextState(key, state)
	}
}

func loopContextState(key string, state int) func(obsws.Event) {
	stateMutex.Lock()
	cachedStates[key] = state
	stateMutex.Unlock()

	return func(event obsws.Event) {
		contextMutex.RLock()
		defer contextMutex.RUnlock()

		for ctx, ctxKey := range contexts {
			if ctxKey == key {
				sdk.SetState(ctx, state)
			}
		}
	}
}

func onWillDisappear(e *sdk.WillDisappearEvent) {
	contextMutex.Lock()
	defer contextMutex.Unlock()

	// replayToggleContexts
	key, ok := contexts[e.Context]

	delete(contexts, e.Context)

	if !ok {
		return
	}

	for _, k := range contexts {
		if k == key {
			return
		}
	}

	clientMutex.Lock()
	clients[key].Disconnect()

	delete(clients, key)
	clientMutex.Unlock()
}

func onSettingsReceived(e *sdk.ReceiveSettingsEvent) {
	var host, password string

	host = sdk.JsonStringValue(e.Settings, "host")
	port := e.Settings.GetInt("port")
	password = sdk.JsonStringValue(e.Settings, "password")

	if port == 0 {
		port = 4444
	}

	key := checkClient(host, port, password)

	if key == "" {
		return
	}

	contextMutex.Lock()
	contexts[e.Context] = key
	contextMutex.Unlock()

	stateMutex.RLock()
	if state, ok := cachedStates[key]; ok {
		sdk.SetState(e.Context, state)
	}
	stateMutex.RUnlock()
}

func main() {
	sdk.RegisterAction(actionReplayToggle, replayToggle)
	sdk.RegisterAction(actionReplaySave, replaySave)

	sdk.AddHandler(onWillAppear)
	sdk.AddHandler(onWillDisappear)
	sdk.AddHandler(onSettingsReceived)

	// Open and connect the SDK
	err := sdk.Open()

	if err != nil {
		log.Fatalln(err)
	}

	defer cleanupSockets()

	// Wait until the socket is closed, or SIGTERM/SIGINT is received
	sdk.Wait()
}

func cleanupSockets() {
	clientMutex.RLock()
	defer clientMutex.RUnlock()

	for _, client := range clients {
		client.Disconnect()
	}
}