ai-operator/internal/ai/tts/elevenlabs_ws_streaming.go

package tts

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"ai-operator/internal/config"
	"github.com/gorilla/websocket"
)

type ElevenLabsWS struct {
	cfg    config.Config
	conn   *websocket.Conn
	audio  chan AudioChunk
	mu     sync.Mutex
	closed bool
}

func NewElevenLabsWS(cfg config.Config) *ElevenLabsWS {
	return &ElevenLabsWS{cfg: cfg, audio: make(chan AudioChunk, 64)}
}

func (p *ElevenLabsWS) Start(ctx context.Context, req TTSStreamRequest) error {
	if p.cfg.Eleven.APIKey == "" {
		return errors.New("ELEVENLABS_API_KEY is required for ElevenLabs streaming TTS")
	}
	if req.VoiceID == "" {
		return errors.New("ElevenLabs voice id is required for streaming TTS")
	}
	base := strings.TrimRight(p.cfg.Eleven.TTSURL, "/") + "/" + url.PathEscape(req.VoiceID) + "/stream-input"
	u, err := url.Parse(base)
	if err != nil {
		return err
	}
	q := u.Query()
	q.Set("model_id", nonEmpty(req.ModelID, p.cfg.Eleven.TTSModelID))
	q.Set("output_format", nonEmpty(req.OutputFormat, p.cfg.Eleven.TTSOutputFormat))
	q.Set("optimize_streaming_latency", strconv.Itoa(p.cfg.Eleven.TTSOptimizeStreamingLatency))
	u.RawQuery = q.Encode()
	h := http.Header{}
	h.Set("xi-api-key", p.cfg.Eleven.APIKey)
	d := websocket.Dialer{HandshakeTimeout: p.cfg.Eleven.TTSTimeout}
	conn, _, err := d.DialContext(ctx, u.String(), h)
	if err != nil {
		return err
	}
	p.mu.Lock()
	p.conn = conn
	p.audio = make(chan AudioChunk, 64)
	p.closed = false
	p.mu.Unlock()
	bos := map[string]any{
		"text": " ",
		"voice_settings": map[string]any{
			"stability":         p.cfg.Eleven.TTSStability,
			"similarity_boost":  p.cfg.Eleven.TTSSimilarityBoost,
			"style":             p.cfg.Eleven.TTSStyle,
			"use_speaker_boost": p.cfg.Eleven.TTSUseSpeakerBoost,
		},
	}
	if err := conn.WriteJSON(bos); err != nil {
		return err
	}
	go p.readLoop()
	return nil
}

func (p *ElevenLabsWS) SendText(ctx context.Context, text string, flush bool) error {
	p.mu.Lock()
	conn := p.conn
	p.mu.Unlock()
	if conn == nil {
		return errors.New("elevenlabs tts websocket not connected")
	}
	msg := map[string]any{"text": text, "try_trigger_generation": flush}
	if flush {
		msg["flush"] = true
	}
	return conn.WriteJSON(msg)
}

func (p *ElevenLabsWS) Audio() <-chan AudioChunk { return p.audio }

func (p *ElevenLabsWS) Close(ctx context.Context) error {
	p.mu.Lock()
	conn := p.conn
	p.conn = nil
	p.mu.Unlock()
	if conn != nil {
		_ = conn.WriteJSON(map[string]any{"text": ""})
		_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
		_ = conn.Close()
	}
	p.closeAudio()
	return nil
}

func (p *ElevenLabsWS) readLoop() {
	for {
		p.mu.Lock()
		conn := p.conn
		p.mu.Unlock()
		if conn == nil {
			p.closeAudio()
			return
		}
		_, data, err := conn.ReadMessage()
		if err != nil {
			p.closeAudio()
			return
		}
		audio := parseAudio(data)
		if len(audio) == 0 {
			continue
		}
		p.emit(AudioChunk{Data: audio, Timestamp: time.Now().UTC()})
	}
}

func parseAudio(data []byte) []byte {
	var m map[string]any
	if json.Unmarshal(data, &m) == nil {
		for _, k := range []string{"audio", "audio_base64"} {
			if s, ok := m[k].(string); ok && s != "" {
				b, _ := base64.StdEncoding.DecodeString(s)
				return b
			}
		}
		return nil
	}
	return data
}

func (p *ElevenLabsWS) emit(ch AudioChunk) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.closed {
		return
	}
	select {
	case p.audio <- ch:
	default:
	}
}

func (p *ElevenLabsWS) closeAudio() {
	p.mu.Lock()
	defer p.mu.Unlock()
	if !p.closed {
		close(p.audio)
		p.closed = true
	}
}

func nonEmpty(v, fallback string) string {
	if v != "" {
		return v
	}
	return fallback
}