// Package audio talks to the default input and output devices and presents
// 16 kHz mono PCM to the station.
package audio

import (
	"context"
	"encoding/binary"
	"math"

	"github.com/gen2brain/malgo"
)

const stationRate = 16000

// Start opens the default capture device and the default playback device and
// converts them to and from 16 kHz mono. It returns when both devices are
// running. They close when ctx ends.
func Start(ctx context.Context, capture func([]int16), playback func() []int16) error {
	mctx, err := malgo.InitContext(nil, malgo.ContextConfig{}, nil)
	if err != nil {
		return err
	}

	capCfg := malgo.DefaultDeviceConfig(malgo.Capture)
	capCfg.SampleRate = stationRate
	capCfg.Capture.Format = malgo.FormatS16
	capCfg.Capture.Channels = 1

	playCfg := malgo.DefaultDeviceConfig(malgo.Playback)
	playCfg.SampleRate = stationRate
	playCfg.Playback.Format = malgo.FormatS16
	playCfg.Playback.Channels = 1

	var capRes, playRes converter
	var playBuf []int16

	capDev, err := malgo.InitDevice(mctx.Context, capCfg, malgo.DeviceCallbacks{
		Data: func(_, input []byte, framecount uint32) {
			if len(input) == 0 {
				return
			}
			channels := channelCount(len(input), framecount)
			pcm := capRes.Convert(downmix(bytesToS16(input), channels))
			if len(pcm) > 0 {
				capture(pcm)
			}
		},
	})
	if err != nil {
		mctx.Free()
		return err
	}

	playDev, err := malgo.InitDevice(mctx.Context, playCfg, malgo.DeviceCallbacks{
		Data: func(output, _ []byte, framecount uint32) {
			if len(output) == 0 {
				return
			}
			channels := channelCount(len(output), framecount)
			playBuf = append(playBuf, playRes.Convert(playback())...)
			need := int(framecount)
			var frame []int16
			if len(playBuf) >= need {
				frame = playBuf[:need]
				rest := append([]int16(nil), playBuf[need:]...)
				playBuf = rest
			} else {
				frame = playBuf
				playBuf = nil
			}
			writeS16(output, frame, channels)
		},
	})
	if err != nil {
		capDev.Uninit()
		mctx.Free()
		return err
	}

	capRate := int(capDev.SampleRate())
	if capRate == 0 {
		capRate = stationRate
	}
	playRate := int(playDev.SampleRate())
	if playRate == 0 {
		playRate = stationRate
	}
	capRes = converter{from: capRate, to: stationRate}
	playRes = converter{from: stationRate, to: playRate}

	if err := capDev.Start(); err != nil {
		playDev.Uninit()
		capDev.Uninit()
		mctx.Free()
		return err
	}
	if err := playDev.Start(); err != nil {
		capDev.Stop()
		playDev.Uninit()
		capDev.Uninit()
		mctx.Free()
		return err
	}

	go func() {
		<-ctx.Done()
		_ = playDev.Stop()
		_ = capDev.Stop()
		playDev.Uninit()
		capDev.Uninit()
		mctx.Free()
	}()
	return nil
}

func channelCount(bytes int, framecount uint32) int {
	channels := 1
	if framecount > 0 {
		channels = bytes / int(framecount) / 2
	}
	if channels < 1 {
		return 1
	}
	return channels
}

func bytesToS16(b []byte) []int16 {
	n := len(b) / 2
	out := make([]int16, n)
	for i := range out {
		out[i] = int16(binary.NativeEndian.Uint16(b[i*2:]))
	}
	return out
}

func writeS16(dst []byte, samples []int16, channels int) {
	if channels < 1 {
		return
	}
	frames := len(dst) / (channels * 2)
	for i := 0; i < frames; i++ {
		var s int16
		if i < len(samples) {
			s = samples[i]
		}
		for ch := 0; ch < channels; ch++ {
			off := (i*channels + ch) * 2
			binary.NativeEndian.PutUint16(dst[off:], uint16(s))
		}
	}
}

func downmix(samples []int16, channels int) []int16 {
	if channels <= 1 {
		return samples
	}
	n := len(samples) / channels
	out := make([]int16, n)
	for i := range out {
		var sum int
		for ch := 0; ch < channels; ch++ {
			sum += int(samples[i*channels+ch])
		}
		out[i] = int16(sum / channels)
	}
	return out
}

type converter struct {
	from int
	to   int
	pos  float64
	buf  []int16
}

func (c *converter) Convert(in []int16) []int16 {
	if c.from == 0 || c.to == 0 || c.from == c.to {
		return in
	}
	c.buf = append(c.buf, in...)
	step := float64(c.from) / float64(c.to)
	var out []int16
	for c.pos+1 < float64(len(c.buf)) {
		i := int(c.pos)
		frac := c.pos - float64(i)
		sample := (1-frac)*float64(c.buf[i]) + frac*float64(c.buf[i+1])
		out = append(out, int16(math.Round(sample)))
		c.pos += step
	}
	drop := int(c.pos)
	if drop > 0 && drop <= len(c.buf) {
		c.buf = append([]int16(nil), c.buf[drop:]...)
		c.pos -= float64(drop)
	}
	return out
}