#### RemiAudio
#### Copyright (C) 2022-2024 Remilia Scarlet <remilia@posteo.jp>
####
#### This program is free software: you can redistribute it and/or modify it
#### under the terms of the GNU Affero General Public License as published by
#### the Free Software Foundation, either version 3 of the License, or (at your
#### option) any later version.
####
#### This program is distributed in the hope that it will be useful, but WITHOUT
#### ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
#### FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public
#### License for more details.
####
#### You should have received a copy of the GNU Affero General Public License
#### along with this program.  If not, see <https://www.gnu.org/licenses/>.
require "./audiofile"

module RemiAudio::Formats
  ###
  ### Constants
  ###

  # :nodoc:
  AU_MAGIC = ".snd"

  # The minimum supported sample rate for Au files.
  AU_MIN_SAMPLE_RATE = 8000

  # The maximum supported sample rate for Au files.
  AU_MAX_SAMPLE_RATE = 352800

  # The minimum supported number of channels for Au files.
  AU_MIN_CHANNELS = 1

  # The maximum supported number of channels for Au files.
  AU_MAX_CHANNELS = 2

  # An array of all the bit depths in Au files that are supported by this
  # library for integer samples.
  AU_SUPPORTED_INT_BIT_DEPTHS = [8, 16, 24, 32, 64]

  # An array of all the bit depths in Au files that are supported by this
  # library for float samples.
  AU_SUPPORTED_FLOAT_BIT_DEPTHS = [32, 64]

  ###
  ### Enumerations
  ###

  # The supported encodings for an Au file.
  enum AuEncoding

    # µLaw encoded samples.
    MuLaw = 1

    # 8-bit signed linear PCM samples
    Lpcm8bit

    # 16-bit signed linear PCM samples
    Lpcm16bit

    # 24-bit signed linear PCM samples
    Lpcm24bit

    # 32-bit signed linear PCM samples
    Lpcm32bit

    # 32-bit IEEE Float samples
    Float32

    # 64-bit IEEE Float samples
    Float64

    # ALaw encoded samples.
    ALaw = 27

    def self.from_value(value : Int) : self
      from_value?(value) || raise UnsupportedEncodingError.new("Unsupported Au encoding")
    end

    # Gets the `SampleFormat` that describes this encoding.  If this is an
    # encoding that does not describe a PCM/Float encoding, or an encoding not
    # directly supported by RemiAudio, this raises an
    # `UnsupportedEncodingError`.
    @[AlwaysInline]
    def toSampleFormat : SampleFormat
      self.toSampleFormat? || raise UnsupportedEncodingError.new("Cannot represents #{self} as a SampleFormat")
    end

    # Attempts to get the `SampleFormat` that describes this encoding, or
    # `nil` if it's a non PCM/Float encoding, or an encoding not directly
    # supported by RemiAudio.
    @[AlwaysInline]
    def toSampleFormat? : SampleFormat?
      case self
      when .lpcm8bit?  then SampleFormat::I8
      when .lpcm16bit? then SampleFormat::I16
      when .lpcm24bit? then SampleFormat::I24
      when .lpcm32bit? then SampleFormat::I32
      when .float32?   then SampleFormat::F32
      when .float64?   then SampleFormat::F64
      else nil
      end
    end

    # Returns the bit depth for a single sample using this encoding.
    @[AlwaysInline]
    def getBitDepth : UInt8
      case self
      in .lpcm8bit?, .mu_law?, .a_law? then 8u8
      in .lpcm16bit? then 16u8
      in .lpcm24bit? then 24u8
      in .lpcm32bit?, .float32? then 32u8
      in .float64? then 64u8
      end
    end

    # Returns the size in bytes for a single sample using this encoding.
    @[AlwaysInline]
    def getByteSize : Int32
      case self
      in .lpcm8bit?, .mu_law?, .a_law? then 1
      in .lpcm16bit? then 2
      in .lpcm24bit?, .lpcm32bit?, .float32? then 4
      in .float64? then 8
      end
    end

    # Returns the value to use when converting a value in this encoding to a
    # `::Float64`, or `nil` if converting from this encoding is not supported.
    def getFloat64Div? : Float64?
      {% begin %}
        case self
        when .lpcm8bit?  then ::Int8::MAX.to_f64!
        when .lpcm16bit? then ::Int16::MAX.to_f64!
        when .lpcm24bit? then ({{ (2.0 ** 23) - 1 }}).to_f64!
        when .lpcm32bit? then ::Int32::MAX.to_f64!
        when .float32?   then 1.0
        when .float64?   then 1.0
        else nil
        end
      {% end %}
    end

    # Returns the value to use when converting a value in this encoding to a
    # `::Float32`, or `nil` if converting from this encoding is not supported.
    def getFloat32Div? : Float32?
      {% begin %}
        case self
        when .lpcm8bit?  then ::Int8::MAX.to_f32!
        when .lpcm16bit? then ::Int16::MAX.to_f32!
        when .lpcm24bit? then ({{ (2.0 ** 23) - 1 }}).to_f32!
        when .lpcm32bit? then ::Int32::MAX.to_f32!
        when .float32?   then 1.0f32
        when .float64?   then 1.0f32
        else nil
        end
      {% end %}
    end

    # Returns the value to use when converting a value in this encoding to a
    # `::Float64`.  This will raise an `UnsupportedEncodingError` if converting
    # from this encoding is not supported.
    @[AlwaysInline]
    def getFloat64Div : Float64
      getFloat64Div? ||
        raise UnsupportedEncodingError.new("This Au encoding is currently unsupported by this library: #{self}")
    end

    # Returns the value to use when converting a value in this encoding to a
    # `::Float32`.  This will raise an `UnsupportedEncodingError` if converting
    # from this encoding is not supported.
    @[AlwaysInline]
    def getFloat32Div : Float32
      getFloat32Div? ||
        raise UnsupportedEncodingError.new("This Au encoding is currently unsupported by this library: #{self}")
    end

    # Checks that *array* is the correct type for this encoding.  If it is, this
    # returns `true`, otherwise it returns `false`.
    @[AlwaysInline]
    def checkArray(array : SampleData) : Bool
      case self
      in .mu_law?, .a_law? then array.is_a?(Array(::Int8))
      in .lpcm8bit? then array.is_a?(Array(::Int8))
      in .lpcm16bit? then array.is_a?(Array(::Int16))
      in .lpcm24bit? then array.is_a?(Array(::Int32))
      in .lpcm32bit? then array.is_a?(Array(::Int32))
      in .float32? then array.is_a?(Array(::Float32))
      in .float64? then array.is_a?(Array(::Float64))
      end
    end

    # Creates a new array of the given size that is appropriate for this
    # encoding.
    @[AlwaysInline]
    def makeArray(size : Int32) : SampleData
      case self
      in .mu_law?, .a_law? then Array(UInt8).new(size, 0u8).as(SampleData)
      in .lpcm8bit? then  Array(Int8).new(size, 0i8).as(SampleData)
      in .lpcm16bit? then Array(Int16).new(size, 0i16).as(SampleData)
      in .lpcm24bit?, .lpcm32bit? then Array(Int32).new(size, 0i32).as(SampleData)
      in .float32? then Array(::Float32).new(size, 0.0f32).as(SampleData)
      in .float64? then Array(::Float64).new(size, 0.0f64).as(SampleData)
      end
    end
  end

  ##############################################################################
  ###
  ### AuFile Class
  ###

  # A virtual representation of an Au file (.au/.snd).
  #
  # Au files originated from Sun and store data in liner PCM, IEEE Float, µLaw,
  # or ALaw format.  Their data is stored in big endian format.
  #
  # The official specification supports more encodings, some of which are
  # specific to certain platforms.  This library does not support these.  See
  # `AuEncoding` for a list of supported encodings.
  class AuFile < AudioFile
    # The `AuEncoding` that the samples in this instance uses.
    getter encoding : AuEncoding = AuEncoding::Lpcm16bit

    # The annotation string for this instance, or `nil` if there is none.
    getter note : String?

    # Creates a new, blank `AuFile` that is backed by a fresh `::IO::Memory`.
    # This will immediately write a skeleton Au header into `#io`
    def initialize(*, @sampleRate : UInt32 = 44100u32, @channels : UInt32 = 2u32,
                   @encoding : AuEncoding = AuEncoding::Lpcm16bit, @note : String? = nil)
      @bitDepth = @encoding.getBitDepth
      @blockSize = @encoding.getByteSize.to_u16
      checkInternalValues
      assignFuncs

      @io = IO::Memory.new(0)
      @origIoPos = 0

      # Write the header
      writeHeader
    end

    # Creates a new, blank `AuFile` that is backed by the given `::IO`.  This
    # will immediately write a skeleton Au header into *io*.
    def initialize(@io : IO, *, @sampleRate : UInt32 = 44100u32, @channels : UInt32 = 2u32,
                   @encoding : AuEncoding = AuEncoding::Lpcm16bit, @note : String? = nil)
      @bitDepth = @encoding.getBitDepth
      @blockSize = @encoding.getByteSize.to_u16
      checkInternalValues
      assignFuncs

      @origIoPos = @io.pos.to_u64!

      # Write the header
      writeHeader
    end

    protected def initialize(*, _fromIo : IO)
      @io = _fromIo
      @origIoPos = @io.pos.to_u64
      readHeader
      checkInternalValues
      assignFuncs
      @io.pos = @dataStartsAt
    end

    # Creates a new `AuFile` by reading the existing Au data from *io*.  Samples
    # will be streamed from *io* as needed.
    def self.open(io : IO) : AuFile
      AuFile.new(_fromIo: io)
    end

    # Creates a new `AuFile` by reading the existing Au data from *io*, then
    # yields it to the block.  Samples will be streamed from *io* as needed.
    # This will **NOT** automatically close the `AuFile` before returning in
    # case you have passed an *io* that cannot does not support writing.
    def self.open(io : IO, &)
      au = AuFile.open(io)
      yield au
    end

    # Creates a new `AuFile` by reading the existing Au data from the given
    # file.  The file is always opened with the mode `"r+b"`.  Samples will be
    # streamed from the file as needed.
    def self.open(filename : Path|String) : AuFile
      AuFile.new(_fromIo: File.open(filename, "r+b"))
    end

    # Creates a new `AuFile` by reading the existing Au data from the file, then
    # yields it to the block.  The file is always opened with the mode `"r+b"`.
    # Samples will be streamed from *io* as needed.  This **will** automatically
    # close the `AuFile` and underlying `::IO` before returning.
    def self.open(filename : Path|String, &)
      au = AuFile.open(filename)
      yield au
    ensure
      au.try &.close
    end

    # Creates a new `AuFile` that that is backed by *io* and returns the new
    # instance.  A skeleton Au header data is immediately written into *io*.
    def self.create(io : IO, *, sampleRate : Int = 44100, channels : Int = 2,
                    encoding : AuEncoding = AuEncoding::Lpcm16bit, note : String? = nil) : AuFile
      AuFile.new(io, sampleRate: sampleRate.to_u32, channels: channels.to_u32, encoding: encoding, note: note)
    end

    # Creates a new `AuFile` that that is backed by *io* and yields the new
    # instance to the block.  A skeleton Au header data is immediately written
    # into *io*.  The `AuFile` instance will be flushed and closed once the
    # block exits.
    def self.create(io : IO, *, sampleRate : Int = 44100, channels : Int = 2,
                    encoding : AuEncoding = AuEncoding::Lpcm16bit, note : String? = nil, &) : AuFile
      au = AuFile.new(io, sampleRate: sampleRate.to_u32, channels: channels.to_u32, encoding: encoding, note: note)
      yield au
    ensure
      au.try &.flush
      au.try &.close
    end

    # Creates a new `AuFile` that is backed by a new file on disk.
    #
    # This always opens the file with the mode `"w+b"` (i.e. existing files are
    # overwritten/truncated).  A skeleton Au header data is immediately written
    # into the file.
    def self.create(filename : Path|String, *, sampleRate : Int = 44100, channels : Int = 2,
                    encoding : AuEncoding = AuEncoding::Lpcm16bit, note : String? = nil) : AuFile
      file = File.open(filename, "w+b")
      begin
        AuFile.new(file, sampleRate: sampleRate.to_u32, channels: channels.to_u32, encoding: encoding, note: note)
      rescue err : Exception
        file.try &.close
        raise err
      end
    end

    # Creates a new `AuFile` that is backed by a new file on disk, then yields
    # it to the block.  This automatically closes the `AuFile` and the
    # underlying file before the block returns.
    #
    # This always opens the file with the mode `"w+b"` (i.e. existing files are
    # overwritten/truncated).  A skeleton Au header data is immediately written
    # into the file.
    def self.create(filename : Path|String, *, sampleRate : Int = 44100, channels : Int = 2,
                    encoding : AuEncoding = AuEncoding::Lpcm16bit, note : String? = nil, &)
      file = File.open(filename, "w+b")
      begin
        au = AuFile.new(file, sampleRate: sampleRate.to_u32, channels: channels.to_u32, encoding: encoding, note: note)
        begin
          yield au
        ensure
          au.close
        end
      rescue err : Exception
        file.try &.close
        raise err
      end
    end

    ##############################################################################

    # Checks that the internal values are consistent.
    protected def checkInternalValues : Nil
      case @encoding
      when .mu_law?, .a_law?
        raise AudioFileError.new("Bad bit depth for µLaw/ALaw encoding") unless @bitDepth == 8
      when .lpcm8bit?
        raise AudioFileError.new("Bad bit depth for 8-bit LPCM encoding") unless @bitDepth == 8
      when .lpcm16bit?
        raise AudioFileError.new("Bad bit depth for 16-bit LPCM encoding") unless @bitDepth == 16
      when .lpcm24bit?
        raise AudioFileError.new("Bad bit depth for 24-bit LPCM encoding") unless @bitDepth == 24
      when .lpcm32bit?
        raise AudioFileError.new("Bad bit depth for 32-bit LPCM encoding") unless @bitDepth == 32
      when .float32?
        raise AudioFileError.new("Bad bit depth for 32-bit float encoding") unless @bitDepth == 32
      when .float64?
        raise AudioFileError.new("Bad bit depth for 64-bit float encoding") unless @bitDepth == 64
      else
        raise UnsupportedEncodingError.new("This Au encoding is currently unsupported by this library")
      end

      # Check some values
      unless @sampleRate >= AU_MIN_SAMPLE_RATE && @sampleRate <= AU_MAX_SAMPLE_RATE
        raise AudioFileError.new("Unsupported sample rate: #{@sampleRate}")
      end

      unless @channels >= AU_MIN_CHANNELS && @channels <= AU_MAX_CHANNELS
        raise AudioFileError.new("Unsupported number of channels: #{@channels}")
      end
    end

    # Assigns the internal functions according to our encoding.
    protected def assignFuncs : Nil
      case @encoding
      in .lpcm8bit?
        @sampleReadFn = ->fastReadSampleInt8
        @samplesReadFn = ->fastReadSamplesInt8(SampleData)
        @sampleWriteFn = ->fastWriteSampleInt8(Sample)
        @samplesWriteFn = ->fastWriteSamplesInt8(SampleData)
        @sampleToF64Fn = ->int8ToF64(Sample)
        @f64ToSampleFn = ->f64ToInt8(Float64)
        @sampleToF32Fn = ->int8ToF32(Sample)
        @f32ToSampleFn = ->f32ToInt8(Float32)
      in .lpcm16bit?
        @sampleReadFn = ->fastReadSampleInt16BE
        @samplesReadFn = ->fastReadSamplesInt16BE(SampleData)
        @sampleWriteFn = ->fastWriteSampleInt16BE(Sample)
        @samplesWriteFn = ->fastWriteSamplesInt16BE(SampleData)
        @sampleToF64Fn = ->int16ToF64(Sample)
        @f64ToSampleFn = ->f64ToInt16(Float64)
        @sampleToF32Fn = ->int16ToF32(Sample)
        @f32ToSampleFn = ->f32ToInt16(Float32)
      in .lpcm24bit?
        @sampleReadFn = ->fastReadSampleInt24BEFourByte
        @samplesReadFn = ->fastReadSamplesInt24BEFourByte(SampleData)
        @sampleWriteFn = ->fastWriteSampleInt24BEFourByte(Sample)
        @samplesWriteFn = ->fastWriteSamplesInt24BEFourByte(SampleData)
        @sampleToF64Fn = ->int24ToF64(Sample)
        @f64ToSampleFn = ->f64ToInt24(Float64)
        @sampleToF32Fn = ->int24ToF32(Sample)
        @f32ToSampleFn = ->f32ToInt24(Float32)
      in .lpcm32bit?
        @sampleReadFn = ->fastReadSampleInt32BE
        @samplesReadFn = ->fastReadSamplesInt32BE(SampleData)
        @sampleWriteFn = ->fastWriteSampleInt32BE(Sample)
        @samplesWriteFn = ->fastWriteSamplesInt32BE(SampleData)
        @sampleToF64Fn = ->int32ToF64(Sample)
        @f64ToSampleFn = ->f64ToInt32(Float64)
        @sampleToF32Fn = ->int32ToF32(Sample)
        @f32ToSampleFn = ->f32ToInt32(Float32)
      in .float32?
        @sampleReadFn = ->fastReadSampleFloat32BE
        @samplesReadFn = ->fastReadSamplesFloat32BE(SampleData)
        @sampleWriteFn = ->fastWriteSampleFloat32BE(Sample)
        @samplesWriteFn = ->fastWriteSamplesFloat32BE(SampleData)
        @sampleToF64Fn = ->float32ToF64(Sample)
        @f64ToSampleFn = ->f64ToFloat32(Float64)
        @sampleToF32Fn = ->float32ToF32(Sample)
        @f32ToSampleFn = ->f32ToFloat32(Float32)
      in .float64?
        @sampleReadFn = ->fastReadSampleFloat64BE
        @samplesReadFn = ->fastReadSamplesFloat64BE(SampleData)
        @sampleWriteFn = ->fastWriteSampleFloat64BE(Sample)
        @samplesWriteFn = ->fastWriteSamplesFloat64BE(SampleData)
        @sampleToF64Fn = ->float64ToF64(Sample)
        @f64ToSampleFn = ->f64ToFloat64(Float64)
        @sampleToF32Fn = ->float32ToF32(Sample)
        @f32ToSampleFn = ->f32ToFloat32(Float32)
      in .mu_law?, .a_law?
        @sampleReadFn = ->fastReadSampleUInt8
        @samplesReadFn = ->fastReadSamplesUInt8(SampleData)
        @sampleWriteFn = ->fastWriteSampleUInt8(Sample)
        @samplesWriteFn = ->fastWriteSamplesUInt8(SampleData)
        @sampleToF64Fn = ->AudioFile.cannotConvertToFloat64(Sample)
        @f64ToSampleFn = ->AudioFile.cannotConvertToSample(Float64)
        @sampleToF32Fn = ->AudioFile.cannotConvertToFloat32(Sample)
        @f32ToSampleFn = ->AudioFile.cannotConvertToSample(Float32)
      end
    end

    # Reads the header data.
    protected def readHeader : Nil
      # Read the magic bytes and check that it's an Au file.
      unless @io.read_string(AU_MAGIC.size) == AU_MAGIC
        raise AudioFileError.new("Not an Au file")
      end

      # Read the rest of the header information.  Note that Au files use big
      # endian throughout.
      @dataStartsAt = @io.readUInt32.byte_swap.to_u64!
      @dataSizeInBytes = @io.readUInt32.byte_swap.to_u64!
      @encoding = AuEncoding.from_value(@io.readUInt32.byte_swap)
      @sampleRate = @io.readUInt32.byte_swap
      @channels = @io.readUInt32.byte_swap

      # Set bit depth and get the block size
      @bitDepth = @encoding.getBitDepth
      @blockSize = @encoding.getByteSize.to_u16

      # Read the annotation, if any.  If there is no annotation, then this field
      # will simply be four zero bytes.
      initialChar = @io.readUInt32
      unless initialChar == 0
        # There is an annotation we need to read.  We have the first four bytes
        # already. But it may be three bytes or less.
        @note = case
                when ((initialChar & 0xFF000000) >> 24) == 0
                  # Odd case, empty string, ignore junk bytes after
                  ""

                when ((initialChar & 0x00FF0000) >> 16) == 0
                  # Single character
                  String.build { |str| str.write_byte(((initialChar & 0xFF000000) >> 24).to_u8!) }

                when ((initialChar & 0x0000FF00) >> 8) == 0
                  # Two characters
                  String.build do |str|
                    str.write_byte(((initialChar & 0xFF000000) >> 24).to_u8!)
                    str.write_byte(((initialChar & 0x00FF0000) >> 16).to_u8!)
                  end

                when (initialChar & 0x000000FF) == 0
                  # Three characters
                  String.build do |str|
                    str.write_byte(((initialChar & 0xFF000000) >> 24).to_u8!)
                    str.write_byte(((initialChar & 0x00FF0000) >> 16).to_u8!)
                    str.write_byte(((initialChar & 0x0000FF00) >> 8).to_u8!)
                  end

                else
                  # Read a full null-terminated string that is four bytes or more.
                  String.build do |str|
                    # Write the first four bytes we already have
                    str.write_byte(((initialChar & 0xFF000000) >> 24).to_u8!)
                    str.write_byte(((initialChar & 0x00FF0000) >> 16).to_u8!)
                    str.write_byte(((initialChar & 0x0000FF00) >> 8).to_u8!)
                    str.write_byte( (initialChar & 0x000000FF).to_u8!)

                    # Write more bytes until we hit null.
                    while (byte = @io.readUInt8) != 0
                      str.write_byte(byte)
                    end
                  end
                end
      end

      # Determine the number of samples we have.
      @numSamples = @dataSizeInBytes.to_u64.tdiv(@blockSize)
      raise AudioFileError.new("Bad number of samples, some are missing") unless @numSamples % @channels == 0
    end

    # Writes the header data.
    protected def writeHeader : Nil
      @io << AU_MAGIC
      @io.writeUInt32BE(69u32) # Temporary value for data offset
      @io.writeUInt32BE(0xFFFFFFFF_u32) # Temporary value for data size in bytes
      @io.writeUInt32BE(@encoding.value.to_u32)
      @io.writeUInt32BE(@sampleRate)
      @io.writeUInt32BE(@channels.to_u32!)

      if @note && !@note.not_nil!.empty?
        str = @note.not_nil!
        @io << str

        if str.size < 4
          (4 - str.size).times { |_| @io.write_byte(0) }
        else
          @io.write_byte(0)
        end
      else
        @io.writeUInt32BE(0u32)
      end

      @io.flush # Needed for the pos below to work right
      @dataStartsAt = @io.pos.to_u64
    end

    # Flushes the stream, then adjusts the header data with the correct data
    # size and offset.
    protected def finishHeader : Nil
      raise AudioFileError.new("Bad number of samples, some are missing") unless @numSamples % @channels == 0
      @io.flush
      @io.withExcursion(4) do
        @io.writeUInt32BE(@dataStartsAt.to_u32) # Data offset

        finalNumBytes : UInt32 = 0u32
        begin
          finalNumBytes = (@numSamples * @blockSize).to_u32
          @io.writeUInt32BE(finalNumBytes) # Data size in bytes
        rescue ::OverflowError
          raise AudioFileError.new("Too many samples for an Au file")
        end
      end
    end

    ##############################################################################

    # :inherited:
    def sampleFormat : SampleFormat
      @encoding.toSampleFormat || raise AudioFileError.new("Cannot represent encoding #{@encoding} as a SampleFormat")
    end

    # :inherited:
    def read(dest : Array(Float64)) : Int32
      temp : SampleData = @encoding.makeArray(dest.size)
      num : Int32 = readSamples(temp)
      num.times do |i|
        dest[i] = @sampleToF64Fn.call(temp[i])
      end
      num
    end

    # :inherited:
    def read(dest : Array(Float32)) : Int32
      temp : SampleData = @encoding.makeArray(dest.size)
      num : Int32 = readSamples(temp)
      num.times do |i|
        dest[i] = @sampleToF32Fn.call(temp[i])
      end
      num
    end

    # :inherited:
    def readSamplesToEnd : SampleData
      curSample = self.pos
      toRead : UInt64 = @numSamples - curSample
      ret : SampleData = @encoding.makeArray(toRead.to_i32)
      @samplesReadFn.call(ret)
      ret
    end

    # :inherited:
    def flush : Nil
      return unless @changed
      finishHeader
      @io.flush
    end

    def copyTo(dest : IO) : Nil
      flush
      @io.withExcursion(@origIoPos) do
        IO.copy(@io, dest)
      end
    end
  end
end
