alice-backend/station/main.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"encoding/binary"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gorilla/websocket"
)

const rate = 16000
const frameSamples = 320
const hop = 160
const captureBlock = hop * 3 * 16

type template struct {
	Frames [][]float64 `json:"frames"`
}
type message struct {
	Type string `json:"type"`
	Code string `json:"code,omitempty"`
}

var hann [frameSamples]float64
var coefficients [20]float64

func init() {
	for i := range hann {
		hann[i] = 0.5 - 0.5*math.Cos(2*math.Pi*float64(i)/float64(frameSamples-1))
	}
	for b := range coefficients {
		freq := 150 * math.Pow(7000.0/150.0, float64(b)/19)
		coefficients[b] = 2 * math.Cos(2*math.Pi*freq/rate)
	}
}

// features calculates log spectral energy in coarse bands, normalized against overall level.
// It is a small experimental template matcher, not a production speech model.
func features(samples []int16) []float64 {
	const bands = 20
	out := make([]float64, bands)
	for b := 0; b < bands; b++ {
		coeff := coefficients[b]
		var a, c float64
		for i, s := range samples {
			v := float64(s)*hann[i] + coeff*a - c
			c = a
			a = v
		}
		power := a*a + c*c - coeff*a*c
		out[b] = math.Log(1 + math.Max(0, power))
	}
	var mean float64
	for _, v := range out {
		mean += v
	}
	mean /= bands
	for i := range out {
		out[i] -= mean
	}
	return out
}
func rms(s []int16) float64 {
	var v float64
	for _, x := range s {
		v += float64(x) * float64(x)
	}
	return math.Sqrt(v / float64(len(s)))
}
func distance(a, b [][]float64) float64 {
	if len(a) == 0 || len(b) == 0 {
		return math.Inf(1)
	}
	prev := make([]float64, len(b)+1)
	for i := 1; i < len(prev); i++ {
		prev[i] = math.Inf(1)
	}
	for _, x := range a {
		cur := make([]float64, len(b)+1)
		cur[0] = math.Inf(1)
		for j, y := range b {
			var cost float64
			for k := range x {
				d := x[k] - y[k]
				cost += d * d
			}
			cost = math.Sqrt(cost / float64(len(x)))
			cur[j+1] = cost + math.Min(prev[j], math.Min(prev[j+1], cur[j]))
		}
		prev = cur
	}
	return prev[len(b)] / float64(len(a)+len(b))
}
func loadTemplates(path string) ([]template, error) {
	f, e := os.Open(path)
	if e != nil {
		return nil, e
	}
	defer f.Close()
	b, e := io.ReadAll(io.LimitReader(f, 2<<20+1))
	if len(b) > 2<<20 {
		return nil, errors.New("template file too large")
	}
	if e != nil {
		return nil, e
	}
	var t []template
	e = json.Unmarshal(b, &t)
	if e != nil {
		return nil, e
	}
	if len(t) == 0 || len(t) > 10 {
		return nil, errors.New("no templates")
	}
	for _, x := range t {
		if len(x.Frames) == 0 || len(x.Frames) > 400 {
			return nil, errors.New("empty template")
		}
		for _, f := range x.Frames {
			if len(f) != 20 {
				return nil, errors.New("invalid feature dimensions")
			}
		}
	}
	return t, nil
}

// Convert each 480-frame/8-channel block to 160 mono samples. Only physical
// mic channel 0 is used, never the odd channels (which may have other meanings).
func downsample(raw []byte) []int16 {
	n := len(raw) / (8 * 2 * 3)
	out := make([]int16, n)
	for i := 0; i < n; i++ {
		var sum int32
		for j := 0; j < 3; j++ {
			off := (i*3 + j) * 16
			sum += int32(int16(binary.LittleEndian.Uint16(raw[off:])))
		}
		out[i] = int16(sum / 3)
	}
	return out
}
func pcm(s []int16) []byte {
	b := make([]byte, len(s)*2)
	for i, x := range s {
		binary.LittleEndian.PutUint16(b[i*2:], uint16(x))
	}
	return b
}
func capture(device string) (*exec.Cmd, io.ReadCloser, error) {
	c := exec.Command("arecord", "-q", "-D", device, "-t", "raw", "-f", "S16_LE", "-r", "48000", "-c", "8", "--period-time=10000", "--buffer-time=80000")
	c.Stderr = os.Stderr
	terminateWithClient(c)
	p, e := c.StdoutPipe()
	if e != nil {
		return nil, nil, e
	}
	if e = c.Start(); e != nil {
		return nil, nil, e
	}
	return c, p, nil
}

// Continuously drain the microphone even when the network or user is slow. Drop
// old frames rather than sending stale audio captured before activation.
type liveReader struct {
	chunks    chan []byte
	remainder []byte
	err       error
}

func newLiveReader(r io.Reader) *liveReader {
	l := &liveReader{chunks: make(chan []byte, 2)}
	go func() {
		for {
			b := make([]byte, captureBlock)
			if _, e := io.ReadFull(r, b); e != nil {
				l.err = e
				close(l.chunks)
				return
			}
			select {
			case l.chunks <- b:
			default:
				select {
				case <-l.chunks:
				default:
				}
				l.chunks <- b
			}
		}
	}()
	return l
}
func (l *liveReader) discardPending() {
	l.remainder = nil
	for {
		select {
		case _, ok := <-l.chunks:
			if !ok {
				return
			}
		default:
			return
		}
	}
}
func (l *liveReader) Read(p []byte) (int, error) {
	if len(l.remainder) == 0 {
		b, ok := <-l.chunks
		if !ok {
			return 0, l.err
		}
		l.remainder = b
	}
	n := copy(p, l.remainder)
	l.remainder = l.remainder[n:]
	return n, nil
}

var mutePath string
var errMuted = errors.New("microphone muted or mute state unavailable")

func unmuted() bool {
	b, e := os.ReadFile(mutePath)
	return e == nil && strings.TrimSpace(string(b)) == "1"
}
func cue(device string) error {
	samples := make([]int16, 2400)
	for i := range samples {
		samples[i] = int16(12000 * math.Sin(2*math.Pi*880*float64(i)/24000) * math.Sin(math.Pi*float64(i)/float64(len(samples))))
	}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	c := exec.CommandContext(ctx, "aplay", "-q", "-D", device, "-t", "raw", "-f", "S16_LE", "-r", "24000", "-c", "1")
	c.Stdin = bytes.NewReader(pcm(samples))
	c.Stderr = os.Stderr
	terminateWithClient(c)
	return c.Run()
}
func waitRecord(r io.Reader, max time.Duration, silence time.Duration, threshold float64) ([][]float64, error) {
	var samples []int16
	var frames [][]float64
	start := time.Now()
	last := start
	heard := false
	buf := make([]byte, hop*3*16)
	for time.Since(start) < max {
		if _, e := io.ReadFull(r, buf); e != nil {
			return nil, e
		}
		if !unmuted() {
			samples = nil
			frames = nil
			heard = false
			last = time.Now()
			start = time.Now()
			continue
		}
		s := downsample(buf)
		samples = append(samples, s...)
		if len(samples) >= frameSamples {
			frame := samples[:frameSamples]
			if rms(frame) >= threshold {
				heard = true
				last = time.Now()
			}
			if heard {
				frames = append(frames, features(frame))
			}
			samples = samples[hop:]
		}
		if len(frames) >= 250 {
			return frames, nil
		}
		if heard && time.Since(last) >= silence {
			return frames, nil
		}
	}
	if !heard {
		return nil, errors.New("no speech detected")
	}
	return frames, nil
}
func enroll(r io.Reader, path string, count int, threshold float64, playback string) error {
	var all []template
	in := bufio.NewReader(os.Stdin)
	for i := 0; i < count; i++ {
		fmt.Printf("Sample %d/%d: press Enter, then say the wake phrase.\n", i+1, count)
		if _, e := in.ReadString('\n'); e != nil {
			return e
		}
		for !unmuted() {
			time.Sleep(100 * time.Millisecond)
		}
		if e := cue(playback); e != nil {
			return e
		}
		if live, ok := r.(*liveReader); ok {
			live.discardPending()
		}
		f, e := waitRecord(r, 4*time.Second, 350*time.Millisecond, threshold)
		if e != nil {
			return e
		}
		all = append(all, template{f})
		fmt.Printf("Recorded %d frames.\n", len(f))
	}
	b, e := json.Marshal(all)
	if e != nil {
		return e
	}
	return os.WriteFile(path, b, 0600)
}
func connect(url, token string) (*websocket.Conn, error) {
	h := http.Header{}
	h.Set("Authorization", "Bearer "+token)
	d := websocket.Dialer{HandshakeTimeout: 15 * time.Second, TLSClientConfig: &tls.Config{MinVersion: tls.VersionTLS12}}
	c, resp, e := d.Dial(url, h)
	if e != nil {
		if resp != nil {
			return nil, fmt.Errorf("websocket handshake HTTP %d: %w", resp.StatusCode, e)
		}
		return nil, e
	}
	c.SetReadLimit(4 << 20)
	_ = c.SetReadDeadline(time.Now().Add(20 * time.Second))
	var m message
	if e = c.ReadJSON(&m); e != nil {
		c.Close()
		return nil, e
	}
	if m.Type != "ready" {
		c.Close()
		return nil, fmt.Errorf("expected ready, got %q", m.Type)
	}
	_ = c.SetReadDeadline(time.Time{})
	return c, nil
}
func exchange(url, token string, live *liveReader, fileAudio []byte, threshold float64, max, silence time.Duration, playback string) error {
	if fileAudio == nil && !unmuted() {
		return errMuted
	}
	c, e := connect(url, token)
	if e != nil {
		return e
	}
	defer c.Close()
	if fileAudio == nil && !unmuted() {
		return errMuted
	}
	_ = c.SetWriteDeadline(time.Now().Add(15 * time.Second))
	if e = c.WriteJSON(message{Type: "start"}); e != nil {
		return e
	}
	if fileAudio != nil {
		for len(fileAudio) > 0 {
			n := 3200
			if n > len(fileAudio) {
				n = len(fileAudio)
			}
			_ = c.SetWriteDeadline(time.Now().Add(15 * time.Second))
			if e = c.WriteMessage(websocket.BinaryMessage, fileAudio[:n]); e != nil {
				return e
			}
			fileAudio = fileAudio[n:]
		}
	} else {
		if e = cue(playback); e != nil {
			return e
		}
		live.discardPending()
		// Keep the onset in a short local buffer. It contains only post-cue audio.
		const prerollSamples = rate / 5
		pre := make([]int16, 0, prerollSamples+hop)
		started := time.Now()
		lastVoice := started
		heard := false
		var voiced int
		buf := make([]byte, captureBlock)
		for time.Since(started) < max {
			if _, e = io.ReadFull(live, buf); e != nil {
				return e
			}
			if !unmuted() {
				return errMuted
			}
			s := downsample(buf)
			if rms(s) >= threshold {
				voiced++
				lastVoice = time.Now()
			} else {
				voiced = 0
			}
			if !heard {
				pre = append(pre, s...)
				if len(pre) > prerollSamples {
					copy(pre, pre[len(pre)-prerollSamples:])
					pre = pre[:prerollSamples]
				}
				if voiced < 4 {
					continue
				}
				heard = true
				s = pre
			}
			_ = c.SetWriteDeadline(time.Now().Add(15 * time.Second))
			if e = c.WriteMessage(websocket.BinaryMessage, pcm(s)); e != nil {
				return e
			}
			if heard && time.Since(lastVoice) >= silence {
				break
			}
		}
		if !unmuted() {
			return errMuted
		}
		if !heard {
			return errors.New("no speech detected after cue")
		}
	}
	_ = c.SetWriteDeadline(time.Now().Add(15 * time.Second))
	if e = c.WriteJSON(message{Type: "end"}); e != nil {
		return e
	}
	player := exec.Command("aplay", "-q", "-D", playback, "-t", "raw", "-f", "S16_LE", "-r", "24000", "-c", "1")
	player.Stderr = os.Stderr
	terminateWithClient(player)
	w, e := player.StdinPipe()
	if e != nil {
		return e
	}
	if e = player.Start(); e != nil {
		return e
	}
	defer func() {
		w.Close()
		if player.Process != nil {
			player.Process.Kill()
		}
		player.Wait()
	}()
	_ = c.SetReadDeadline(time.Now().Add(90 * time.Second))
	for {
		typ, data, e := c.ReadMessage()
		if e != nil {
			return e
		}
		if typ == websocket.BinaryMessage {
			if _, e = w.Write(data); e != nil {
				return e
			}
			continue
		}
		if typ == websocket.TextMessage {
			var m message
			if e = json.Unmarshal(data, &m); e != nil {
				return e
			}
			switch m.Type {
			case "done":
				if e = w.Close(); e != nil {
					return e
				}
				e = player.Wait()
				player.Process = nil
				return e
			case "error":
				return fmt.Errorf("server error: %s", m.Code)
			}
		}
	}
}
func main() {
	mode := flag.String("mode", "ptt", "ptt, enroll, wake, or file")
	muteGPIO := flag.String("mute-gpio", "/sys/class/gpio/gpio414/value", "physical microphone mute value, 1 = unmuted")
	input := flag.String("input", "", "16 kHz mono PCM16LE file for file mode (max 30 seconds)")
	device := flag.String("capture", "hw:0,1", "ALSA capture device")
	playback := flag.String("playback", "default", "ALSA playback device")
	url := flag.String("url", "wss://alice.pepega.systems/v1/voice", "server websocket URL")
	path := flag.String("templates", "wake-templates.json", "local enrollment file")
	th := flag.Float64("voice-threshold", 350, "speech RMS threshold")
	match := flag.Float64("match-threshold", 1.3, "experimental DTW maximum distance")
	oneTurn := flag.Bool("once", false, "exit wake mode after one successful response")
	count := flag.Int("samples", 4, "enrollment samples")
	max := flag.Duration("max", 12*time.Second, "maximum request duration")
	silence := flag.Duration("silence", 1200*time.Millisecond, "end after speech silence")
	flag.Parse()
	mutePath = *muteGPIO
	if *th <= 0 || *max <= 0 || *silence <= 0 || *count < 1 || *count > 10 {
		log.Fatal("invalid threshold, duration, or sample count")
	}
	token := strings.TrimSpace(os.Getenv("DEVICE_TOKEN"))
	if *mode != "enroll" && token == "" {
		log.Fatal("DEVICE_TOKEN is required")
	}
	var ts []template
	var e error
	if *mode == "wake" {
		ts, e = loadTemplates(*path)
		if e != nil {
			log.Fatal(e)
		}
	} else if *mode != "ptt" && *mode != "enroll" && *mode != "file" {
		log.Fatal("unknown mode")
	}
	if *mode == "file" {
		f, e := os.Open(*input)
		if e != nil {
			log.Fatal(e)
		}
		b, e := io.ReadAll(io.LimitReader(f, 30*rate*2+1))
		f.Close()
		if e != nil {
			log.Fatal(e)
		}
		if len(b) == 0 || len(b) > 30*rate*2 || len(b)%2 != 0 {
			log.Fatal("input must contain at most 30 seconds of PCM16LE and have an even nonzero size")
		}
		if e = exchange(*url, token, nil, b, *th, *max, *silence, *playback); e != nil {
			log.Fatal(e)
		}
		return
	}
	if *mode == "ptt" {
		fmt.Println("Press Enter to start speaking.")
		if _, e = bufio.NewReader(os.Stdin).ReadString('\n'); e != nil {
			log.Fatal(e)
		}
	}
	cmd, r, e := capture(*device)
	if e != nil {
		log.Fatal(e)
	}
	var once sync.Once
	cleanup := func() { once.Do(func() { r.Close(); cmd.Process.Kill(); cmd.Wait() }) }
	defer cleanup()
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	go func() { <-sig; cleanup(); os.Exit(0) }()
	live := newLiveReader(r)
	if *mode == "enroll" {
		if e = enroll(live, *path, *count, *th, *playback); e != nil {
			cleanup()
			log.Fatal(e)
		}
		return
	}
	if *mode == "ptt" {
		for {
			for !unmuted() {
				time.Sleep(100 * time.Millisecond)
			}
			e = exchange(*url, token, live, nil, *th, *max, *silence, *playback)
			if errors.Is(e, errMuted) {
				log.Println("microphone muted; waiting for unmute")
				continue
			}
			if e != nil {
				cleanup()
				log.Fatal(e)
			}
			return
		}
	}

	// Detect whole speech segments locally before opening a voice connection.
	log.Println("Listening locally for enrolled phrase")
	raw := make([]byte, hop*3*16)
	var samples []int16
	var f [][]float64
	active := false
	quiet := 0
	for {
		if _, e = io.ReadFull(live, raw); e != nil {
			cleanup()
			log.Fatal(e)
		}
		if !unmuted() {
			samples = nil
			f = nil
			active = false
			quiet = 0
			continue
		}
		s := downsample(raw)
		samples = append(samples, s...)
		if len(samples) < frameSamples {
			continue
		}
		v := rms(samples[:frameSamples])
		frame := samples[:frameSamples]
		if v >= *th {
			active = true
			quiet = 0
		}
		samples = samples[hop:]
		if active {
			f = append(f, features(frame))
			if v < *th {
				quiet++
			}
			if quiet >= 35 || len(f) >= 250 {
				best := math.Inf(1)
				for _, t := range ts {
					if d := distance(f, t.Frames); d < best {
						best = d
					}
				}
				log.Printf("wake candidate distance=%.3f", best)
				if best <= *match {
					log.Println("wake detected")
					if e = exchange(*url, token, live, nil, *th, *max, *silence, *playback); e != nil {
						log.Printf("session: %v", e)
					} else if *oneTurn {
						return
					}
				}
				f = nil
				active = false
				quiet = 0
				samples = nil
			}
		}
	}
}