ai-operator/internal/media/asteriskws/client.go

package asteriskws

import (
	"context"
	"errors"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"ai-operator/internal/media"
)

const MaxMessageBytes = 65500

type ClientConfig struct {
	URL               string
	ConnectionID      string
	Codec             media.Codec
	DialTimeout       time.Duration
	ReadLimitBytes    int64
	BasicAuthUser     string
	BasicAuthPassword string
}

type Client struct {
	cfg       ClientConfig
	conn      *websocket.Conn
	events    chan ControlEvent
	audio     chan media.AudioChunk
	mu        sync.RWMutex
	writeMu   sync.Mutex
	stats     media.Stats
	paused    bool
	frameSize int
	ptime     time.Duration
	closed    chan struct{}
	closeOnce sync.Once
}

func NewClient(cfg ClientConfig) *Client {
	if cfg.DialTimeout == 0 {
		cfg.DialTimeout = 5 * time.Second
	}
	if cfg.ReadLimitBytes == 0 {
		cfg.ReadLimitBytes = MaxMessageBytes
	}
	return &Client{cfg: cfg, events: make(chan ControlEvent, 32), audio: make(chan media.AudioChunk, 32), closed: make(chan struct{})}
}
func (c *Client) Connect(ctx context.Context) error {
	d := websocket.Dialer{HandshakeTimeout: c.cfg.DialTimeout, Subprotocols: []string{"media"}}
	h := http.Header{}
	if c.cfg.BasicAuthUser != "" {
		req, _ := http.NewRequest(http.MethodGet, c.cfg.URL, nil)
		req.SetBasicAuth(c.cfg.BasicAuthUser, c.cfg.BasicAuthPassword)
		h = req.Header
	}
	conn, _, err := d.DialContext(ctx, c.cfg.URL, h)
	if err != nil {
		return err
	}
	conn.SetReadLimit(c.cfg.ReadLimitBytes)
	c.conn = conn
	go c.readLoop()
	return nil
}
func (c *Client) Close(ctx context.Context) error {
	c.closeOnce.Do(func() {
		close(c.closed)
		if c.conn != nil {
			_ = c.conn.Close()
		}
		close(c.events)
		close(c.audio)
	})
	return nil
}
func (c *Client) Events() <-chan ControlEvent    { return c.events }
func (c *Client) Audio() <-chan media.AudioChunk { return c.audio }
func (c *Client) Stats() media.Stats             { c.mu.RLock(); defer c.mu.RUnlock(); return c.stats }
func (c *Client) SendCommand(ctx context.Context, command string) error {
	if c.conn == nil {
		return errors.New("media websocket not connected")
	}
	return c.writeMessage(websocket.TextMessage, []byte(command))
}
func (c *Client) GetStatus(ctx context.Context) error { return c.SendCommand(ctx, "GET_STATUS") }
func (c *Client) StartMediaBuffering(ctx context.Context) error {
	return c.SendCommand(ctx, "START_MEDIA_BUFFERING")
}
func (c *Client) StopMediaBuffering(ctx context.Context, id string) error {
	return c.SendCommand(ctx, "STOP_MEDIA_BUFFERING "+id)
}
func (c *Client) FlushMedia(ctx context.Context) error { return c.SendCommand(ctx, "FLUSH_MEDIA") }
func (c *Client) MarkMedia(ctx context.Context, id string) error {
	return c.SendCommand(ctx, "MARK_MEDIA "+id)
}
func (c *Client) Hangup(ctx context.Context) error { return c.SendCommand(ctx, "HANGUP") }
func (c *Client) SendAudio(ctx context.Context, data []byte) error {
	if c.conn == nil {
		return errors.New("media websocket not connected")
	}
	for len(data) > 0 {
		c.mu.RLock()
		paused := c.paused
		c.mu.RUnlock()
		if paused {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(10 * time.Millisecond):
				continue
			}
		}
		n, ptime := c.nextAudioFrameSize(len(data))
		if err := c.writeMessage(websocket.BinaryMessage, data[:n]); err != nil {
			return err
		}
		c.addOutbound(n)
		data = data[n:]
		if len(data) > 0 && ptime > 0 {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(ptime):
			}
		}
	}
	return nil
}

func (c *Client) writeMessage(messageType int, data []byte) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	return c.conn.WriteMessage(messageType, data)
}

func (c *Client) nextAudioFrameSize(remaining int) (int, time.Duration) {
	c.mu.RLock()
	frameSize := c.frameSize
	ptime := c.ptime
	c.mu.RUnlock()
	if frameSize <= 0 || frameSize > MaxMessageBytes {
		frameSize = MaxMessageBytes
		ptime = 0
	}
	if remaining < frameSize {
		return remaining, ptime
	}
	return frameSize, ptime
}
func (c *Client) readLoop() {
	defer c.Close(context.Background())
	for {
		mt, data, err := c.conn.ReadMessage()
		if err != nil {
			return
		}
		switch mt {
		case websocket.TextMessage:
			ev := ParseControlEvent(string(data))
			c.recordEvent(ev)
			select {
			case c.events <- ev:
			default:
			}
		case websocket.BinaryMessage:
			now := time.Now().UTC()
			chunk := media.AudioChunk{Data: data, Codec: c.cfg.Codec, Timestamp: now}
			c.recordAudio(len(data), now)
			select {
			case c.audio <- chunk:
			default:
			}
		}
	}
}
func (c *Client) recordEvent(ev ControlEvent) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.stats.TextEvents++
	switch ev.Type {
	case ControlMediaXOFF:
		c.paused = true
		c.stats.XOFFCount++
	case ControlMediaXON:
		c.paused = false
		c.stats.XONCount++
	case ControlMediaStart:
		ms := ev.MediaStart()
		c.stats.MediaStart = ms
		if ms != nil && ms.OptimalFrameSize > 0 {
			c.frameSize = ms.OptimalFrameSize
		}
		if ms != nil && ms.PTimeMS > 0 {
			c.ptime = time.Duration(ms.PTimeMS) * time.Millisecond
		}
	}
}
func (c *Client) recordAudio(n int, t time.Time) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.stats.InboundFrames++
	c.stats.InboundBytes += int64(n)
	if c.stats.FirstInboundAudioAt == nil {
		tt := t
		c.stats.FirstInboundAudioAt = &tt
	}
	tt := t
	c.stats.LastInboundAudioAt = &tt
}
func (c *Client) addOutbound(n int) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.stats.OutboundFrames++
	c.stats.OutboundBytes += int64(n)
}