//go:build !disable_format_aac && !disable_codec_libfdk_aac
// +build !disable_format_aac,!disable_codec_libfdk_aac

package aac

import (
	"bytes"
	"errors"
	"fmt"
	"git.gammaspectra.live/S.O.N.G/Kirika/audio"
	"git.gammaspectra.live/S.O.N.G/Kirika/cgo"
	"git.gammaspectra.live/S.O.N.G/go-fdkaac/fdkaac"
	aac_adts "github.com/edgeware/mp4ff/aac"
	"github.com/edgeware/mp4ff/mp4"
	"io"
	"time"
	"unsafe"
)

type Format struct {
}

func NewFormat() Format {
	return Format{}
}

func (f Format) Name() string {
	return "aac"
}

func (f Format) Description() string {
	return "libfdk-aac (S.O.N.G/go-fdkaac)"
}

func decodeFrame(decoder *fdkaac.AacDecoder, r io.Reader) ([]float32, error) {
	pcm, err := tryDecodeFrame(decoder)
	if err != nil {
		return nil, err
	}
	if pcm != nil {
		return pcm, err
	}

	header, _, err := aac_adts.DecodeADTSHeader(r)

	if err != nil {
		return nil, err
	}

	data := make([]byte, header.PayloadLength)

	if _, err = io.ReadFull(r, data); err != nil {
		return nil, err
	}

	var n int
	fullData := append(header.Encode(), data...)
	if n, err = decoder.Fill(fullData); n != 0 {
		return nil, errors.New("buffer under read")
	}
	if err != nil {
		return nil, err
	}

	return decodeFrame(decoder, r)
}

func tryDecodeFrame(decoder *fdkaac.AacDecoder) ([]float32, error) {
	pcm, err := decoder.Decode()

	if err != nil {
		return nil, err
	}

	if pcm != nil {
		return cgo.Int32ToFloat32(cgo.BytesToInt32(pcm, 16), 16), nil
	}

	return nil, nil
}

func decodeFrameMP4(decoder *fdkaac.AacDecoder, demuxer *mp4Decoder) (result []float32, err error) {
	pcm, err := tryDecodeFrame(decoder)
	if err != nil {
		return nil, err
	}
	if pcm != nil {
		return pcm, err
	}

	samples := demuxer.Read()
	if samples == nil {
		return nil, io.EOF
	}
	var n int
	for _, sample := range samples {
		if n, err = decoder.Fill(sample); n != 0 {
			return nil, errors.New("buffer under read")
		}
		if err != nil {
			return nil, err
		}
		pcm, err = tryDecodeFrame(decoder)
		if pcm != nil {
			result = append(result, pcm...)
		}
	}

	return result, nil
}

func (f Format) Open(r io.ReadSeekCloser) (audio.Source, error) {
	decoder := fdkaac.NewAacDecoder()

	mp4Demuxer, err := tryDecodeMP4(r)
	if err != nil { //try ADTS
		r.Seek(0, io.SeekStart)
		err = decoder.InitAdts()
		if err != nil {
			return audio.Source{}, err
		}

		buf, err := decodeFrame(decoder, r)

		if err != nil {
			decoder.Close()
			return audio.Source{}, err
		}

		newChannel := make(chan []float32)

		go func() {
			defer close(newChannel)
			defer decoder.Close()

			if len(buf) > 0 {
				newChannel <- buf
			}

			for {
				buf, err = decodeFrame(decoder, r)

				if err != nil {
					return
				}

				if len(buf) > 0 {
					newChannel <- buf
				}
			}
		}()

		return audio.Source{
			Channels:   decoder.NumChannels(),
			SampleRate: decoder.SampleRate(),
			Blocks:     newChannel,
		}, nil
	} else {
		return audio.Source{}, fmt.Errorf("unsupported format mp4")
		err = decoder.InitRaw(mp4Demuxer.cookie)
		if err != nil {
			return audio.Source{}, err
		}

		buf, err := decodeFrameMP4(decoder, mp4Demuxer)

		if err != nil {
			decoder.Close()
			return audio.Source{}, err
		}

		newChannel := make(chan []float32)

		go func() {
			defer close(newChannel)
			defer decoder.Close()

			if len(buf) > 0 {
				newChannel <- buf
			}

			for {
				buf, err = decodeFrameMP4(decoder, mp4Demuxer)

				if err != nil {
					return
				}

				if len(buf) > 0 {
					newChannel <- buf
				}
			}
		}()

		return audio.Source{
			Channels:   decoder.NumChannels(),
			SampleRate: decoder.SampleRate(),
			Blocks:     newChannel,
		}, nil
	}

}

func (f Format) Encode(source audio.Source, writer io.WriteCloser, options map[string]interface{}) error {
	var bitrate = 128
	var isHEv2 bool
	var format = "adts"

	if options != nil {
		var val interface{}
		var ok bool
		var intVal int
		var strVal string
		if val, ok = options["bitrate"]; ok {
			if strVal, ok = val.(string); ok {
				switch strVal {
				case "320k":
					bitrate = 320
				case "256k":
					bitrate = 256
				case "192k":
					bitrate = 192
				case "128k":
					bitrate = 128

				default:
					return fmt.Errorf("unknown setting bitrate=%s", strVal)
				}
			} else if intVal, ok = val.(int); ok {
				bitrate = intVal
			}
		}
		if val, ok = options["mode"]; ok {
			if strVal, ok = val.(string); ok {
				switch strVal {
				case "lc":
					isHEv2 = false
				case "hev2":
					isHEv2 = true

				default:
					return fmt.Errorf("unknown setting mode=%s", strVal)
				}
			}
		}
		if val, ok = options["format"]; ok {
			if strVal, ok = val.(string); ok {
				format = strVal
			}
		}
	}

	muxingMode := fdkaac.MuxingModeADTS

	if format == "adts" {
		muxingMode = fdkaac.MuxingModeADTS
	} else if format == "mp4" {
		muxingMode = fdkaac.MuxingModeRAW
	} else if format == "adif" {
		muxingMode = fdkaac.MuxingModeADIF
	} else {
		return fmt.Errorf("unsupported format %s", format)
	}

	encoder := fdkaac.NewAacEncoder()
	if isHEv2 {
		err := encoder.InitHEv2(source.Channels, source.SampleRate, bitrate*1024, muxingMode)
		if err != nil {
			return err
		}
	} else {
		err := encoder.InitLc(source.Channels, source.SampleRate, bitrate*1024, muxingMode)
		if err != nil {
			return err
		}
	}
	defer encoder.Close()

	frameSize := encoder.FrameSize() * encoder.Channels()

	if format == "mp4" {
		init := mp4.CreateEmptyInit()
		init.AddEmptyTrack(uint32(source.SampleRate), "audio", "en")
		trackId := init.Moov.Mvhd.NextTrackID - 1
		trak := init.Moov.Trak

		objType := aac_adts.AAClc
		if isHEv2 {
			objType = aac_adts.HEAACv2
		}

		{
			stsd := trak.Mdia.Minf.Stbl.Stsd
			asc := &aac_adts.AudioSpecificConfig{
				ObjectType:           byte(objType),
				ChannelConfiguration: byte(source.Channels),
				SamplingFrequency:    source.SampleRate,
				ExtensionFrequency:   0,
				SBRPresentFlag:       false,
				PSPresentFlag:        false,
			}
			switch objType {
			case aac_adts.HEAACv1:
				asc.ExtensionFrequency = 2 * source.SampleRate
				asc.SBRPresentFlag = true
			case aac_adts.HEAACv2:
				asc.ExtensionFrequency = 2 * source.SampleRate
				asc.SBRPresentFlag = true
				asc.ChannelConfiguration = 1
				asc.PSPresentFlag = true
			}

			buf := &bytes.Buffer{}
			err := asc.Encode(buf)
			if err != nil {
				return err
			}
			ascBytes := buf.Bytes()
			esds := mp4.CreateEsdsBox(ascBytes)
			mp4a := mp4.CreateAudioSampleEntryBox("mp4a",
				uint16(asc.ChannelConfiguration),
				16, uint16(source.SampleRate), esds)
			stsd.AddChild(mp4a)
		}

		init.Encode(writer)

		var seqNumber uint32
		var packetsWritten uint64
		var outputBuffer [][]byte

		segmentDuration := time.Millisecond * 100

		outputSegment := func() {
			seg := mp4.NewMediaSegment()
			frag, _ := mp4.CreateFragment(seqNumber, trackId)
			seg.AddFragment(frag)

			for _, b := range outputBuffer {
				frag.AddFullSampleToTrack(mp4.FullSample{
					Sample: mp4.Sample{
						Dur:  uint32(frameSize),
						Size: uint32(len(b)),
					},
					DecodeTime: uint64(frameSize) * packetsWritten,
					Data:       b,
				}, trackId)
				packetsWritten++
			}

			seg.Encode(writer)
			seqNumber++
			outputBuffer = nil
		}

		outputPacket := func(packet []byte) {
			outputBuffer = append(outputBuffer, packet)

			if time.Duration(float64(time.Second)*(float64(frameSize*len(outputBuffer))/float64(source.SampleRate))) >= segmentDuration {
				outputSegment()
			}
		}

		var buffer []int16
		for block := range source.Blocks {

			buffer = append(buffer, cgo.Float32ToInt16(block)...)

			for len(buffer) >= frameSize {
				sl := buffer[:frameSize]
				buf, err := encoder.Encode(unsafe.Slice((*byte)(unsafe.Pointer(&sl[0])), len(sl)*2))

				if err != nil {
					return err
				}

				if len(buf) > 0 {
					outputPacket(buf)
				}

				buffer = buffer[frameSize:]
			}
		}

		if len(buffer) > 0 {
			//pad
			buffer = append(buffer, make([]int16, frameSize-len(buffer))...)
			buf, err := encoder.Encode(unsafe.Slice((*byte)(unsafe.Pointer(&buffer[0])), len(buffer)*2))

			if err != nil {
				return err
			}

			if len(buf) > 0 {
				outputPacket(buf)
			}
		}

		//Do flush
		for {
			buf, err := encoder.Flush()

			if err != nil {
				return err
			}

			if len(buf) > 0 {
				outputPacket(buf)
			} else {
				break
			}
		}

		if len(outputBuffer) > 0 {
			outputSegment()
		}
	} else {

		var buffer []int16
		for block := range source.Blocks {

			buffer = append(buffer, cgo.Float32ToInt16(block)...)

			for len(buffer) >= frameSize {
				sl := buffer[:frameSize]
				buf, err := encoder.Encode(unsafe.Slice((*byte)(unsafe.Pointer(&sl[0])), len(sl)*2))

				if err != nil {
					return err
				}

				if len(buf) > 0 {
					_, err = writer.Write(buf)
					if err != nil {
						return err
					}
				}

				buffer = buffer[frameSize:]
			}
		}

		if len(buffer) > 0 {
			//pad
			buffer = append(buffer, make([]int16, frameSize-len(buffer))...)
			buf, err := encoder.Encode(unsafe.Slice((*byte)(unsafe.Pointer(&buffer[0])), len(buffer)*2))

			if err != nil {
				return err
			}

			if len(buf) > 0 {
				_, err = writer.Write(buf)
				if err != nil {
					return err
				}
			}
		}

		//Do flush
		for {
			buf, err := encoder.Flush()

			if err != nil {
				return err
			}

			if len(buf) > 0 {
				_, err = writer.Write(buf)
				if err != nil {
					return err
				}
			} else {
				break
			}
		}
	}

	return nil
}

func (f Format) Identify(peek []byte, extension string) bool {
	return extension == "aac" || extension == "adts"
}