#### RemiAudio
#### Copyright (C) 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 "socket"
require "../dsp/ditherer"

####
#### Driver to Transmit Audio over TCP
####

module RemiAudio::Drivers::Tcp
  # The `TcpDevice` allows for raw audio to be sent over a TCP connection.
  #
  # Audio input is always expected to be 32-bit floating point, but a
  # `TcpDevice` can be configured to convert this to other formats before
  # transmitting the audio over the connection.  This can optionally be done
  # with a `RemiAudio::DSP::Ditherer` for higher quality.
  class TcpDevice < AudioDevice
    # The remote hostname.
    getter host : String

    # The remote port.
    getter port : UInt16

    # The format that audio will be converted to before transmission over the
    # TCP connection.
    getter outFormat : RemiAudio::SampleFormat

    ##
    ## Various buffers used when converting audio.
    ##

    @f64Buf : Array(Float64)?
    @i64Buf : Array(Int64)?
    @i32Buf : Array(Int32)?
    @i16Buf : Array(Int16)?
    @i8Buf : Array(Int8)?
    @u8Buf : Array(UInt8)?

    # When non-nil, then dithering will be used to convert the audio to the
    # chosen `#outFormat` before transmission.  If this is `nil`, then no
    # dithering is performaned during the conversion.  If the output format is
    # already `RemiAudio::SampleFormat:F32`, then this does nothing.
    property ditherer : RemiAudio::DSP::Ditherer?

    # Whether or not to use noise shaping when using a `#ditherer`.
    property? noiseShaping : Bool = false

    # The remote connection.
    @sock : TCPSocket?

    # The actual function that does the conversion and writing to the socket.
    @writeFn : Proc(Array(Float32)|Slice(Float32), Nil)

    # Creates a new `TcpDevice` instance that will send audio to the given host
    # and port.  The *outFormat` parameter is used to determine what format the
    # audio is converted to before being transmitted over the connection; the
    # input format is always 32-bit floating point.
    #
    # The *inSampleRate* and *inChannel` parameters are stored in `#sampleRate`
    # and `#channels`, respectively, and are just for informational purposes.
    # They are otherwise unused by this class.
    def initialize(inSampleRate : Int, @outFormat : RemiAudio::SampleFormat, inChannels : Int,
                   *, @host : String, @port : UInt16)
      @sampleRate = inSampleRate.to_u32.as(UInt32)
      @channels = inChannels.to_u8.as(UInt8)

      case @outFormat
      in .f64?
        @writeFn = ->writeF64(Array(Float32)|Slice(Float32))
        @bitDepth = 64
      in .f32?
        @writeFn = ->writeIdentity(Array(Float32)|Slice(Float32))
        @bitDepth = 32
      in .i64?
        @writeFn =  ->writeI64(Array(Float32)|Slice(Float32))
        @bitDepth = 64
      in .i32?
        @writeFn = ->writeI32(Array(Float32)|Slice(Float32))
        @bitDepth = 32
      in .i24?
        @writeFn = ->writeI24(Array(Float32)|Slice(Float32))
        @bitDepth = 24
      in .i16?
        @writeFn = ->writeI16(Array(Float32)|Slice(Float32))
        @bitDepth = 16
      in .i8?
        @writeFn = ->writeI8(Array(Float32)|Slice(Float32))
        @bitDepth = 8
      in .u8?
        @writeFn = ->writeU8(Array(Float32)|Slice(Float32))
        @bitDepth = 8
      end
    end

    # :inherit:
    def start : Nil
      raise AudioDeviceError.new("Attempted to start a device twice") if @started
      @sock = TCPSocket.new(@host, @port)
      @sock.not_nil!.keepalive = true
      @sock.not_nil!.send_buffer_size = @expectedBufferSize.to_i32
      @started = true
    end

    # :inherit:
    def bufferSize=(value : Int) : Nil
      if value > Int32::MAX
        raise ArgumentError.new("Buffer size too large for a TCP driver")
      end
      super
      @sock.try { |con| con.send_buffer_size = @expectedBufferSize.to_i32 }
    end

    # :inherit:
    def stop : Nil
      @sock.try(&.close)
      @started = false
    end

    # :inherit:
    def writeBuffer(buf : Array(Float32)|Slice(Float32)) : Nil
      {% unless flag?(:remiaudio_wd40) %}
        raise AudioDeviceError.new("Device not started") unless @started
        unless buf.size == @expectedBufferSize
          raise AudioDeviceError.new("Buffer was the incorrect size: #{buf.size} != #{@expectedBufferSize}")
        end
        RemiLib.assert(!@sock.nil?)
      {% end %}

      @writeFn.call(buf)
    end

    ############################################################################
    ###
    ### Actual conversiona nd transmission methods
    ###

    # Used when no conversion is necessary.
    protected def writeIdentity(buf : Array(Float32)|Slice(Float32)) : Nil
      @sock.not_nil!.write(Slice.new(buf.to_unsafe.unsafe_as(Pointer(UInt8)), buf.size * sizeof(Float32)))
    end

    protected def writeF64(buf : Array(Float32)|Slice(Float32)) : Nil
      outBuf = @f64Buf
      if outBuf.nil? || outBuf.size != buf.size
        outBuf = @f64Buf = Array(Float64).new(buf.size) do |i|
          buf.unsafe_fetch(i).to_f64!
        end
      else
        buf.size.times do |i|
          outBuf.unsafe_put(i, buf.unsafe_fetch(i).to_f64!)
        end
      end

      @sock.not_nil!.write(Slice.new(outBuf.to_unsafe.unsafe_as(Pointer(UInt8)), outBuf.size * sizeof(Float64)))
    end

    protected def writeI64(buf : Array(Float32)|Slice(Float32)) : Nil
      outBuf = @i64Buf
      if outBuf.nil? || outBuf.size != buf.size
        outBuf = @i64Buf = Array(Int64).new(buf.size) do |i|
          (buf.unsafe_fetch(i) * Int64::MAX).to_i64!
        end
      else
        buf.size.times do |i|
          outBuf.unsafe_put(i, (buf.unsafe_fetch(i) * Int64::MAX).to_i64!)
        end
      end

      @sock.not_nil!.write(Slice.new(outBuf.to_unsafe.unsafe_as(Pointer(UInt8)), outBuf.size * sizeof(Int64)))
    end

    protected def writeI32(buf : Array(Float32)|Slice(Float32)) : Nil
      outBuf = @i32Buf
      if outBuf.nil? || outBuf.size != buf.size
        outBuf = @i32Buf = Array(Int32).new(buf.size) do |i|
          (buf.unsafe_fetch(i) * Int32::MAX).to_i32!
        end
      else
        buf.size.times do |i|
          outBuf.unsafe_put(i, (buf.unsafe_fetch(i) * Int32::MAX).to_i32!)
        end
      end

      @sock.not_nil!.write(Slice.new(outBuf.to_unsafe.unsafe_as(Pointer(UInt8)), outBuf.size * sizeof(Int32)))
    end

    protected def writeI24(buf : Array(Float32)|Slice(Float32)) : Nil
      {% begin %}
        {% i24Max = (2 ** 23) - 1 %}
        {% i24Min = -(2 ** 23) %}

        # We can just use @i32Buf
        outBuf = @i32Buf
        if dith = @ditherer
          if outBuf.nil? || outBuf.size != buf.size
            outBuf = @i32Buf = Array(Int32).new(buf.size) do |i|
              dith.dither(buf.unsafe_fetch(i).as(Sample), :F32, 24, useNoiseShaping: @noiseShaping).as(Int32)
            end
          else
            buf.size.times do |i|
              outBuf.unsafe_put(
                i,
                dith.dither(buf.unsafe_fetch(i).as(Sample), :F32, 24, useNoiseShaping: @noiseShaping).as(Int32)
              )
            end
          end
        else
          if outBuf.nil? || outBuf.size != buf.size
            outBuf = @i32Buf = Array(Int32).new(buf.size) do |i|
              (buf.unsafe_fetch(i) * {{i24Max}}).to_i32!.clamp({{i24Min}}, {{i24Max}})
            end
          else
            buf.size.times do |i|
              outBuf.unsafe_put(
                i,
                (buf.unsafe_fetch(i) * {{i24Max}}).to_i32!.clamp({{i24Min}}, {{i24Max}})
              )
            end
          end
        end

        @sock.not_nil!.write(Slice.new(outBuf.to_unsafe.unsafe_as(Pointer(UInt8)), outBuf.size * sizeof(Int32)))
      {% end %}
    end

    protected def writeI16(buf : Array(Float32)|Slice(Float32)) : Nil
      outBuf = @i16Buf
      if dith = @ditherer
        if outBuf.nil? || outBuf.size != buf.size
          outBuf = @i16Buf = Array(Int16).new(buf.size) do |i|
            dith.dither(buf.unsafe_fetch(i).as(Sample), :F32, 16, useNoiseShaping: @noiseShaping).as(Int16)
          end
        else
          buf.size.times do |i|
            outBuf.unsafe_put(
              i,
              dith.dither(buf.unsafe_fetch(i).as(Sample), :F32, 16, useNoiseShaping: @noiseShaping).as(Int16)
            )
          end
        end
      else
        if outBuf.nil? || outBuf.size != buf.size
          outBuf = @i16Buf = Array(Int16).new(buf.size) do |i|
            (buf.unsafe_fetch(i) * Int16::MAX).to_i16!
          end
        else
          buf.size.times do |i|
            outBuf.unsafe_put( i, (buf.unsafe_fetch(i) * Int16::MAX).to_i16! )
          end
        end
      end

      @sock.not_nil!.write(Slice.new(outBuf.to_unsafe.unsafe_as(Pointer(UInt8)), outBuf.size * sizeof(Int16)))
    end

    protected def writeI8(buf : Array(Float32)|Slice(Float32)) : Nil
      outBuf = @i8Buf
      if dith = @ditherer
        if outBuf.nil? || outBuf.size != buf.size
          outBuf = @i8Buf = Array(Int8).new(buf.size) do |i|
            dith.dither(buf.unsafe_fetch(i).as(Sample), :F32, 8, useNoiseShaping: @noiseShaping).as(Int8)
          end
        else
          buf.size.times do |i|
            outBuf.unsafe_put(
              i,
              dith.dither(buf.unsafe_fetch(i).as(Sample), :F32, 8, useNoiseShaping: @noiseShaping).as(Int8)
            )
          end
        end
      else
        if outBuf.nil? || outBuf.size != buf.size
          outBuf = @i8Buf = Array(Int8).new(buf.size) do |i|
            (buf.unsafe_fetch(i) * Int8::MAX).to_i8!
          end
        else
          buf.size.times do |i|
            outBuf.unsafe_put( i, (buf.unsafe_fetch(i) * Int8::MAX).to_i8! )
          end
        end
      end

      @sock.not_nil!.write(Slice.new(outBuf.to_unsafe.unsafe_as(Pointer(UInt8)), outBuf.size * sizeof(Int8)))
    end

    protected def writeU8(buf : Array(Float32)|Slice(Float32)) : Nil
      outBuf = @u8Buf
      if dith = @ditherer
        if outBuf.nil? || outBuf.size != buf.size
          outBuf = @u8Buf = Array(UInt8).new(buf.size) do |i|
            i8Smp : Int16 = dith.dither(buf.unsafe_fetch(i).as(Sample), :F32, 8, useNoiseShaping: @noiseShaping).as(Int8).to_i16!
            (i8Smp + 128).to_u8!
          end
        else
          buf.size.times do |i|
            i8Smp = dith.dither(buf.unsafe_fetch(i).as(Sample), :F32, 8, useNoiseShaping: @noiseShaping).as(Int8).to_i16!
            outBuf.unsafe_put(i, (i8Smp + 128).to_u8!)
          end
        end
      else
        if outBuf.nil? || outBuf.size != buf.size
          outBuf = @u8Buf = Array(UInt8).new(buf.size) do |i|
            i8Smp = (buf.unsafe_fetch(i) * Int8::MAX).to_i8!.to_i16!
            (i8Smp + 128).to_u8!
          end
        else
          buf.size.times do |i|
            i8Smp = (buf.unsafe_fetch(i) * Int8::MAX).to_i8!.to_i16!
            outBuf.unsafe_put(i, (i8Smp + 128).to_u8!)
          end
        end
      end

      @sock.not_nil!.write(Slice.new(outBuf.to_unsafe.unsafe_as(Pointer(UInt8)), outBuf.size * sizeof(UInt8)))
    end
  end
end
