package audio

import (
	"git.gammaspectra.live/S.O.N.G/Kirika/vector"
	"sync/atomic"
)

type int16Source struct {
	BitDepth   int
	SampleRate int
	Channels   int
	Blocks     chan []int16
	locked     atomic.Bool
}

func newInt16Source(bitDepth, sampleRate, channels int) TypedSource[int16] {
	if bitDepth <= 0 || bitDepth > 16 {
		return nil
	}
	return &int16Source{
		BitDepth:   bitDepth,
		SampleRate: sampleRate,
		Channels:   channels,
		Blocks:     make(chan []int16),
	}
}

func (s *int16Source) Split(n int) []Source {
	return Split[int16](s, n)
}

func (s *int16Source) New() TypedSource[int16] {
	return newInt16Source(s.BitDepth, s.SampleRate, s.Channels)
}

func (s *int16Source) GetBlocks() chan []int16 {
	if s.Locked() {
		return nil
	}
	if !s.locked.Swap(true) {
		return s.Blocks
	}
	return nil
}

func (s *int16Source) Close() {
	close(s.Blocks)
}

func (s *int16Source) GetSampleRate() int {
	return s.SampleRate
}

func (s *int16Source) GetChannels() int {
	return s.Channels
}

func (s *int16Source) GetBitDepth() int {
	return s.BitDepth
}

func (s *int16Source) Locked() bool {
	return s.locked.Load()
}

func (s *int16Source) Unlock() {
	s.locked.Store(false)
}

func (s *int16Source) SwapBlocks(blocks chan []int16) chan []int16 {
	oldBlocks := s.Blocks
	s.Blocks = blocks
	return oldBlocks
}

func (s *int16Source) GetFormat() SourceFormat {
	return SourceInt16
}

func (s *int16Source) ToFloat32() TypedSource[float32] {
	if s.Locked() {
		return nil
	}

	source := newFloat32Source(s.BitDepth, s.SampleRate, s.Channels)

	go func() {
		defer source.Close()
		for block := range s.GetBlocks() {
			source.IngestFloat32(vector.Int16ToFloat32(block, s.BitDepth))
		}
	}()

	return source
}

func (s *int16Source) ToInt16() TypedSource[int16] {
	return s
}

func (s *int16Source) ToInt32(bitDepth int) TypedSource[int32] {
	if s.Locked() {
		return nil
	}

	source := newInt32Source(bitDepth, s.SampleRate, s.Channels)

	go func() {
		diff := s.BitDepth - bitDepth

		if diff >= 0 {
			defer source.Close()
			for block := range s.GetBlocks() {
				buf := make([]int32, len(block))
				for i := range block {
					buf[i] = int32(block[i]) >> diff
				}
				source.IngestInt32(buf, bitDepth)
			}
		} else {
			diff = -diff
			defer source.Close()
			for block := range s.GetBlocks() {
				buf := make([]int32, len(block))
				for i := range block {
					buf[i] = int32(block[i]) << diff
				}
				source.IngestInt32(buf, bitDepth)
			}
		}
	}()

	return source
}

func (s *int16Source) IngestFloat32(buf []float32) {
	s.Blocks <- vector.Float32ToInt16(buf)
}

func (s *int16Source) IngestInt8(buf []int8, bitDepth int) {
	block := make([]int16, len(buf))
	for i := range buf {
		block[i] = int16(buf[i]) << (s.BitDepth - bitDepth)
	}
	s.Blocks <- block
}

func (s *int16Source) IngestInt16(buf []int16, bitDepth int) {
	if bitDepth == s.BitDepth {
		s.Blocks <- buf
	} else {
		block := make([]int16, len(buf))
		for i := range buf {
			block[i] = buf[i] << (s.BitDepth - bitDepth)
		}
		s.Blocks <- block
	}
}

func (s *int16Source) IngestInt24(buf []byte, bitDepth int) {
	block := make([]int16, len(buf)/3)

	for i := 0; i < len(buf); i += 3 {
		sample := uint32(buf[i])
		sample += uint32(buf[i+1]) << 8
		sample += uint32(buf[i+2]) << 16

		//sign extend
		block[i/3] = int16(int32(sample<<8) >> (bitDepth - s.BitDepth + 8))
	}
	s.Blocks <- block
}

func (s *int16Source) IngestInt32(buf []int32, bitDepth int) {
	block := make([]int16, len(buf))
	for i := range buf {
		block[i] = int16(buf[i] >> (bitDepth - s.BitDepth))
	}
	s.Blocks <- block
}

func (s *int16Source) IngestNative(buf []int16, bitDepth int) {
	s.IngestInt16(buf, bitDepth)
}