main.go

package main

import (
	"bytes"
	"context"
	"encoding/binary"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"os"
	"os/signal"
	"strings"
	"sync"
	"time"
	"unsafe"

	"github.com/go-ole/go-ole"
	"github.com/moutend/go-hook/keyboard"
	"github.com/moutend/go-wav"
	"github.com/moutend/go-wca"
)

var version = "latest"
var revision = "latest"
var offset int
var tx = time.Now()
var mainBuffer []float64
var subBuffer []float64
var bufferingWG sync.WaitGroup

type FilenameFlag struct {
	Value string
}

func (f *FilenameFlag) Set(value string) (err error) {
	if !strings.HasSuffix(value, ".wav") {
		err = fmt.Errorf("specify WAVE audio file (*.wav)")
		return
	}
	f.Value = value
	return
}

func (f *FilenameFlag) String() string {
	return f.Value
}

func main() {
	var err error
	if err = run(os.Args); err != nil {
		log.Fatal(err)
	}
	fmt.Println("Successfully done")
}

func run(args []string) (err error) {
	var filenameFlag FilenameFlag
	var versionFlag bool
	var audio = &wav.File{}
	var file []byte

	f := flag.NewFlagSet(args[0], flag.ExitOnError)
	f.Var(&filenameFlag, "input", "Specify WAVE format audio (e.g. music.wav)")
	f.Var(&filenameFlag, "i", "Alias of --input")
	f.BoolVar(&versionFlag, "version", false, "Show version")
	f.Parse(args[1:])

	if versionFlag {
		fmt.Printf("%s-%s\n", version, revision)
		return
	}
	if filenameFlag.Value == "" {
		return
	}
	if file, err = ioutil.ReadFile(filenameFlag.Value); err != nil {
		return
	}
	if err = wav.Unmarshal(file, audio); err != nil {
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	errChan := make(chan error, 1)
	keyboardChan := make(chan keyboard.KBDLLHOOKSTRUCT, 1)
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	drumData := audio.Float64s()

	var isInterrupted bool
	var renderingWG sync.WaitGroup

	renderingWG.Add(1)
	go func() {
		errChan <- renderSharedTimerDriven(ctx, audio)
		renderingWG.Done()
	}()
	go keyboard.Notify(ctx, keyboardChan)
	for {
		if isInterrupted {
			break
		}
		select {
		case err = <-errChan:
			return
		case <-signalChan:
			isInterrupted = true
			cancel()
			fmt.Println("Interrupted by SIGINT")
		case v := <-keyboardChan:
			if v.Flags != 0 {
				continue
			}
			tx = time.Now()
			length := len(drumData)

			bufferingWG.Add(1)
			subBuffer = make([]float64, length)
			copy(subBuffer, drumData)
			for i, j := offset, 0; i < len(mainBuffer); i, j = i+1, j+1 {
				subBuffer[j] += mainBuffer[i]
				subBuffer[j] /= 2
				if subBuffer[j] > 1.0 {
					subBuffer[j] = 1.0
				} else if subBuffer[j] < -1.0 {
					subBuffer[j] = -1.0
				}
			}
			mainBuffer = subBuffer
			offset = 0
			bufferingWG.Done()
		}
	}
	renderingWG.Wait()
	err = <-errChan
	return
}

func renderSharedTimerDriven(ctx context.Context, audio *wav.File) (err error) {
	if err = ole.CoInitializeEx(0, ole.COINIT_APARTMENTTHREADED); err != nil {
		return
	}
	defer ole.CoUninitialize()

	var de *wca.IMMDeviceEnumerator
	if err = wca.CoCreateInstance(wca.CLSID_MMDeviceEnumerator, 0, wca.CLSCTX_ALL, wca.IID_IMMDeviceEnumerator, &de); err != nil {
		return
	}
	defer de.Release()

	var mmd *wca.IMMDevice
	if err = de.GetDefaultAudioEndpoint(wca.ERender, wca.EConsole, &mmd); err != nil {
		return
	}
	defer mmd.Release()

	var ps *wca.IPropertyStore
	if err = mmd.OpenPropertyStore(wca.STGM_READ, &ps); err != nil {
		return
	}
	defer ps.Release()

	var pv wca.PROPVARIANT
	if err = ps.GetValue(&wca.PKEY_Device_FriendlyName, &pv); err != nil {
		return
	}

	var ac *wca.IAudioClient
	if err = mmd.Activate(wca.IID_IAudioClient, wca.CLSCTX_ALL, nil, &ac); err != nil {
		return
	}
	defer ac.Release()

	var wfx *wca.WAVEFORMATEX
	if err = ac.GetMixFormat(&wfx); err != nil {
		return
	}
	defer ole.CoTaskMemFree(uintptr(unsafe.Pointer(wfx)))

	if wfx.WFormatTag != wca.WAVE_FORMAT_PCM {
		wfx.WFormatTag = 1
		wfx.NSamplesPerSec = uint32(audio.SamplesPerSec())
		wfx.WBitsPerSample = uint16(audio.BitsPerSample())
		wfx.NChannels = uint16(audio.Channels())
		wfx.NBlockAlign = uint16(audio.BlockAlign())
		wfx.NAvgBytesPerSec = uint32(audio.AvgBytesPerSec())
		wfx.CbSize = 0
	}

	var defaultPeriod int64
	var minimumPeriod int64
	var renderingPeriod time.Duration
	var latency uint32 = 16 // millisecond

	if err = ac.GetDevicePeriod(&defaultPeriod, &minimumPeriod); err != nil {
		return
	}
	if err = ac.Initialize(wca.AUDCLNT_SHAREMODE_SHARED, 0, latency*10000, 0, wfx, nil); err != nil {
		return
	}

	renderingPeriod = time.Duration(int(defaultPeriod) * 100)
	fmt.Printf("Default rendering period: %d ms\n", renderingPeriod/time.Millisecond)

	var bufferFrameSize uint32
	if err = ac.GetBufferSize(&bufferFrameSize); err != nil {
		return
	}
	fmt.Printf("Allocated buffer size: %d\n", bufferFrameSize)

	var arc *wca.IAudioRenderClient
	if err = ac.GetService(wca.IID_IAudioRenderClient, &arc); err != nil {
		return
	}
	defer arc.Release()

	if err = ac.Start(); err != nil {
		return
	}
	fmt.Println("Start rendering audio with shared-timer-driven mode")
	fmt.Println("Press Ctrl-C to stop rendering")

	time.Sleep(renderingPeriod)

	var data *byte
	var padding uint32
	var availableFrameSize uint32
	var b *byte
	var isPlaying bool = true

	scale := math.Pow(2.0, float64(wfx.WBitsPerSample)-1.0)

	for {
		if !isPlaying {
			break
		}
		select {
		case <-ctx.Done():
			isPlaying = false
			break
		default:
			if err = ac.GetCurrentPadding(&padding); err != nil {
				return
			}
			availableFrameSize = bufferFrameSize - padding
			if availableFrameSize == 0 {
				continue
			}
			if err = arc.GetBuffer(availableFrameSize, &data); err != nil {
				return
			}

			start := unsafe.Pointer(data)
			lim := int(availableFrameSize) * int(wfx.NBlockAlign)
			samples := int(availableFrameSize) * int(wfx.NChannels)
			rawData := make([]byte, lim)
			var f float64

			remaining := len(mainBuffer) - offset
			if remaining < samples {
				samples = remaining
			}

			bufferingWG.Wait()

			for i := 0; i < samples; i++ {
				f = mainBuffer[offset]
				s := int16(scale * f)
				buf := new(bytes.Buffer)
				binary.Write(buf, binary.LittleEndian, s)
				bs := buf.Bytes()
				rawData[i*2] = bs[0]
				rawData[i*2+1] = bs[1]
				offset++
			}
			for n := 0; n < lim; n++ {
				b = (*byte)(unsafe.Pointer(uintptr(start) + uintptr(n)))
				*b = rawData[n]
			}

			if err = arc.ReleaseBuffer(availableFrameSize, 0); err != nil {
				return
			}
			time.Sleep(renderingPeriod)
		}
	}

	time.Sleep(time.Duration(latency) * time.Millisecond)

	if err = ac.Stop(); err != nil {
		return
	}
	fmt.Println("Stop rendering loopback audio")
	return
}