#### 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/>.

####
#### Very Basic Ogg Demuxer
####

module RemiAudio::Demuxers
  # A very basic Ogg demuxer.  This implementation can read pages, and can
  # rewind to the very beginning of an Ogg file, but cannot seek yet.
  class Ogg
    MAGIC = "OggS"
    MAGIC_BYTES = MAGIC.to_bytes
    PAGE_VERSION = 0 # Mandated to always be zero by the specifications

    @[Flags]
    enum PageType : UInt8
      Continuation = 0x01u8
      StreamStart  = 0x02u8
      StreamEnd    = 0x04u8
    end

    class Error < Exception
    end

    class BadVersion < Error
    end

    class MissingMarkerError < Error
    end

    @io : IO
    @packetsDecoded : Int64 = 0i64
    getter? atPage : Bool = false
    getter startPos : Int64 = 0i64
    getter type : PageType = PageType::None
    getter granulePos : UInt64 = 0u64
    getter serial : UInt32 = 0u32
    getter seqNumber : UInt32 = 0u32
    getter crc : UInt32 = 0u32
    @lastSegment : UInt32 = 0u32
    @segmentLengths : Array(UInt8) = Array(UInt8).new(257, 0u8) # One extra for safety
    @curSeg : Int32 = 0
    getter pageNumber : UInt64 = 0u64

    def packetsDecoded : Int64
      @packetsDecoded - 1
    end

    def packetsDecoded=(@packetsDecoded : Int64)
    end

    def initialize(@io : IO, upto : Int? = nil)
      # Possibly read the first segment
      advancePage(upto)
    end

    def rewind : Nil
      @io.pos = 0
      @packetsDecoded = 0
      advancePage(nil)
    end

    @[AlwaysInline]
    private def findNextMarker(upto : Int?) : Nil
      start = @io.pos
      loop do
        return if @io.read_string(4) == MAGIC
        @io.read_byte
        if upto && (@io.pos - start >= upto)
          raise MissingMarkerError.new("Could not find an Ogg page marker before #{upto} bytes")
        end
      end
    end

    private def advancePage(upto : Int?) : Nil
      findNextMarker(upto) unless @io.read_string(4) == MAGIC
      if byt = @io.read_byte
        if byt != PAGE_VERSION
          raise BadVersion.new("Invalid Ogg version (byte position: #{@io.pos - 1})")
        end
      else
        raise IO::EOFError.new("Unexpected end of stream (could not find Ogg version)")
      end

      @startPos = @io.pos.to_i64! - 5
      @type = PageType.from_value(@io.read_byte || raise IO::EOFError.new("Could not read Ogg page type"))
      @granulePos = @io.readUInt64
      @serial = @io.readUInt32
      @seqNumber = @io.readUInt32
      @crc = @io.readUInt32
      @curSeg = 0
      @lastSegment = 0
      if numSegs = @io.read_byte
        numSegs.times do |_|
          if len = @io.read_byte
            @segmentLengths[@lastSegment] = len
            @lastSegment += 1
          else
            raise IO::EOFError.new("Could not read segment length: unexpected end of stream")
          end
        end
      else
        raise IO::EOFError.new("Could not read num segments: unexpected end of stream")
      end
      @pageNumber += 1
      @atPage = true
    rescue IO::EOFError
      @atPage = false
    end

    @[AlwaysInline]
    protected def advanceSegment(upto : Int?) : Nil
      if @atPage
        @curSeg += 1

        # Advance the page's current segment by one, and check if we've reached
        # the end of this page's segments.
        advancePage(upto) if @curSeg == @lastSegment
      else
        # No page yet, read the first one
        advancePage(upto)
      end
    rescue IO::EOFError
      @atPage = false
    end

    def nextPacket(upto : Int? = nil) : Bytes
      ret = IO::Memory.new
      lastLen : UInt8 = 0
      len : UInt8 = 0
      pos : Int32 = 0
      buf : Bytes = Bytes.new(256)

      while @atPage
        lastLen = len
        len = @segmentLengths[@curSeg]
        if len == 0
          break if lastLen == 255 # Packet was a multiple of 255?
        end

        pos = @io.read_fully(buf[...len])
        if pos == len
          ret.write(buf[...pos])
        else
          raise Error.new("Could not read enough data for segment #{@curSeg} " \
                          "(wanted #{len} bytes, got #{pos}")
        end

        # If the segment is less than 255 bytes, end of packet
        break if len < 255

        # Goto the next segment
        advanceSegment(upto)
      end # while @atPage

      # Goto the next segment, then return
      advanceSegment(upto)
      @packetsDecoded += 1
      ret.to_slice
    end
  end
end
