package main

// silero-vad, the speech detector that replaces the energy threshold (V-487).
//
// Why an energy threshold is not a voice activity detector. It answers "is
// this frame loud", and a fan, a door and a television are all loud. mavwaked
// sends every utterance it accepts to speech-to-text and then to the daemon,
// so a false trigger is a turn Maven takes on something nobody said to her.
// Silero answers "is this frame speech", which is the question.
//
// It is 2.3MB of ONNX and runs on one CPU core in real time. That is not an
// aside: this is the one model in the system that may never be offloaded or
// gated on GPU admission, because a wake path that waits on a card is not a
// wake path.
//
// Nil is a working value. Without -vad-model the daemon runs the energy VAD
// exactly as it did before this file existed.

import (
	"fmt"
	"sync"

	ort "github.com/yalue/onnxruntime_go"
)

const (
	// sileroWindow — samples per inference at 16kHz. The model is fixed at
	// 512 and does not accept another size, which is why this file
	// re-chunks rather than reusing the 480-sample capture frame. main.go
	// used to claim the two matched; that was true of silero v4.
	sileroWindow = 512

	// sileroContext — samples of the previous window prepended to each
	// inference, as the reference implementation does. Without it the first
	// milliseconds of every window are judged with no history and speech
	// onsets score low.
	sileroContext = 64

	// sileroState — the LSTM state carried between windows, [2][1][128].
	sileroStateDim = 128

	// defaultSileroThreshold — probability above which a window is speech.
	// 0.5 is the reference default. Raising it costs speech onsets, which
	// are the quietest part of an utterance.
	defaultSileroThreshold = 0.5
)

// sileroVAD holds one ONNX session and the streaming state around it. It is
// fed 30ms capture frames and answers per frame, buffering across calls
// because 480 samples never line up with a 512-sample window.
type sileroVAD struct {
	mu      sync.Mutex
	session *ort.DynamicAdvancedSession

	pending []float32               // samples not yet part of a full window
	context [sileroContext]float32  // tail of the previous window
	state   []float32               // [2][1][128], carried between windows
	last    float64                 // most recent probability, held between windows
	sr      []int64
}

// newSileroVAD loads the graph. The ONNX environment is initialised here when
// nothing else has done it, because mavwaked has no embedder to do it first.
func newSileroVAD(modelPath, libPath string) (*sileroVAD, error) {
	if !ort.IsInitialized() {
		if libPath != "" {
			ort.SetSharedLibraryPath(libPath)
		}
		if err := ort.InitializeEnvironment(); err != nil {
			return nil, fmt.Errorf("silero: onnx runtime: %w", err)
		}
	}
	s, err := ort.NewDynamicAdvancedSession(modelPath,
		[]string{"input", "state", "sr"}, []string{"output", "stateN"}, nil)
	if err != nil {
		return nil, fmt.Errorf("silero: load %s: %w", modelPath, err)
	}
	return &sileroVAD{
		session: s,
		state:   make([]float32, 2*sileroStateDim),
		sr:      []int64{16000},
	}, nil
}

// Speech reports whether the frame carries speech, and the probability behind
// that answer. A frame that completes no window inherits the previous
// probability, so the caller sees one answer per frame either way.
func (s *sileroVAD) Speech(frame []int16, threshold float64) (bool, float64) {
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, v := range frame {
		s.pending = append(s.pending, float32(v)/32768.0)
	}
	for len(s.pending) >= sileroWindow {
		p, err := s.infer(s.pending[:sileroWindow])
		if err != nil {
			// A failed inference must not silence the microphone. Hold the
			// last answer and let the next window try again.
			break
		}
		s.last = p
		s.pending = s.pending[sileroWindow:]
	}
	return s.last >= threshold, s.last
}

// infer runs one window and rolls the state and the context forward.
func (s *sileroVAD) infer(window []float32) (float64, error) {
	in := make([]float32, sileroContext+sileroWindow)
	copy(in, s.context[:])
	copy(in[sileroContext:], window)

	inT, err := ort.NewTensor(ort.NewShape(1, int64(len(in))), in)
	if err != nil {
		return 0, err
	}
	defer inT.Destroy()
	stT, err := ort.NewTensor(ort.NewShape(2, 1, sileroStateDim), s.state)
	if err != nil {
		return 0, err
	}
	defer stT.Destroy()
	srT, err := ort.NewTensor(ort.NewShape(1), s.sr)
	if err != nil {
		return 0, err
	}
	defer srT.Destroy()

	out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1))
	if err != nil {
		return 0, err
	}
	defer out.Destroy()
	next, err := ort.NewEmptyTensor[float32](ort.NewShape(2, 1, sileroStateDim))
	if err != nil {
		return 0, err
	}
	defer next.Destroy()

	if err := s.session.Run(
		[]ort.Value{inT, stT, srT},
		[]ort.Value{out, next},
	); err != nil {
		return 0, err
	}
	copy(s.state, next.GetData())
	copy(s.context[:], in[len(in)-sileroContext:])
	return float64(out.GetData()[0]), nil
}

// Reset drops the streaming state. Called at every utterance boundary and
// after barge-in, so echo-era history never scores the next sentence.
func (s *sileroVAD) Reset() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.pending = s.pending[:0]
	s.context = [sileroContext]float32{}
	for i := range s.state {
		s.state[i] = 0
	}
	s.last = 0
}

// Close releases the session.
func (s *sileroVAD) Close() error {
	if s == nil || s.session == nil {
		return nil
	}
	return s.session.Destroy()
}