# SPDX-FileCopyrightText: ☭ Emery Hemingway
# SPDX-License-Identifier: Unlicense

import std/[endians, streams, strutils]
import bigints
import ./values

proc readVarint(s: Stream): uint =
  var
    shift = 0
    c = uint s.readUint8
  while (c and 0x80) == 0x80:
    result = result or ((c and 0x7f) shl shift)
    inc(shift, 7)
    c = uint s.readUint8
  result = result or (c shl shift)

proc decodePreserves*(s: Stream; E = void): Preserve[E] =
  ## Decode a Preserves value from a binary-encoded stream.
  if s.atEnd: raise newException(IOError, "End of Preserves stream")
  const endMarker = 0x84
  let tag = s.readUint8()
  case tag
  of 0x80: result = Preserve[E](kind: pkBoolean, bool: false)
  of 0x81: result = Preserve[E](kind: pkBoolean, bool: true)
  of 0x85:
    discard decodePreserves(s, E)
    result = decodePreserves(s, E)
  of 0x86:
    result = decodePreserves(s, E)
    result.embedded = true
  of 0x87:
    var N: int
    let n = int s.readUint8()
    case n
    of 4:
      result = Preserve[E](kind: pkFloat)
      var buf: uint32
      N = s.readData(addr buf, sizeof(buf))
      bigEndian32(addr result.float, addr buf)
    of 8:
      result = Preserve[E](kind: pkDouble)
      var buf: uint64
      N = s.readData(addr buf, sizeof(buf))
      bigEndian64(addr result.double, addr buf)
    else:
      raise newException(IOError, "unhandled IEEE754 value of " & $n & " bytes")
    if N != n: raise newException(IOError, "short read")
  of 0xb0:
    var n = int s.readVarint()
    if n <= sizeof(int):
      result = Preserve[E](kind: pkRegister)
      if n > 0:
        var
          buf: array[sizeof(int), byte]
          off = buf.len - n
        if s.readData(addr buf[off], n) != n:
          raise newException(IOError, "short read")
        if off > 0:
          var fill: uint8 = if (buf[off] and 0x80) == 0x80'u8: 0xff else: 0x00'u8
          for i in 0..<off: buf[i] = fill
        when buf.len == 4:
          bigEndian32(addr result.register, addr buf[0])
        elif buf.len == 8:
          bigEndian64(addr result.register, addr buf[0])
        else: {.error: "int size " & $buf.len & " not supported here".}
    else:
      result = Preserve[E](kind: pkBigInt)
      var buf = newSeq[byte](n)
      if s.readData(addr buf[0], buf.len) != n:
        raise newException(IOError, "short read")
      if (buf[0] and 0x80) == 0x80:
        for i, b in buf: buf[i] = not b
        result.bigint.fromBytes(buf, bigEndian)
        result.bigint = -(result.bigint.succ)
      else:
        result.bigint.fromBytes(buf, bigEndian)
  of 0xb1:
    var data = newString(s.readVarint())
    if data.len > 0:
      let n = s.readData(unsafeAddr data[0], data.len)
      if n != data.len:
        raise newException(IOError, "short read")
    result = Preserve[E](kind: pkString, string: data)
  of 0xb2:
    var data = newSeq[byte](s.readVarint())
    if data.len > 0:
      let n = s.readData(addr data[0], data.len)
      if n != data.len:
        raise newException(IOError, "short read")
    result = Preserve[E](kind: pkByteString, bytes: data)
  of 0xb3:
    var data = newString(s.readVarint())
    if data.len > 0:
      let n = s.readData(addr data[0], data.len)
      if n != data.len:
        raise newException(IOError, "short read")
    result = Preserve[E](kind: pkSymbol, symbol: Symbol data)
  of 0xb4:
    result = Preserve[E](kind: pkRecord)
    var label = decodePreserves(s, E)
    while s.peekUint8() != endMarker:
      result.record.add decodePreserves(s, E)
    result.record.add(move label)
    discard s.readUint8()
  of 0xb5:
    result = Preserve[E](kind: pkSequence)
    while s.peekUint8() != endMarker:
      result.sequence.add decodePreserves(s, E)
    discard s.readUint8()
  of 0xb6:
    result = Preserve[E](kind: pkSet)
    while s.peekUint8() != endMarker:
      incl(result, decodePreserves(s, E))
    discard s.readUint8()
  of 0xb7:
    result = Preserve[E](kind: pkDictionary)
    while s.peekUint8() != endMarker:
      result[decodePreserves(s, E)] = decodePreserves(s, E)
    discard s.readUint8()
  of endMarker:
    raise newException(ValueError, "invalid Preserves stream")
  else:
    raise newException(ValueError, "invalid Preserves tag byte 0x" & tag.toHex(2))

proc decodePreserves*(s: string; E = void): Preserve[E] =
  ## Decode a string of binary-encoded Preserves.
  decodePreserves(s.newStringStream, E)

proc decodePreserves*(s: seq[byte]; E = void): Preserve[E] =
  ## Decode a byte-string of binary-encoded Preserves.
  decodePreserves(cast[string](s), E)

type BufferedDecoder* = object
  ## Type for buffering binary Preserves before decoding.
  stream: StringStream
  appendPosition, decodePosition, maxSize: int

proc newBufferedDecoder*(maxSize = 4096): BufferedDecoder =
  ## Create a new `newBufferedDecoder`.
  runnableExamples:
    var
      buf = newBufferedDecoder()
      bin = encode(parsePreserves("<foobar>"))
    buf.feed(bin[0..2])
    buf.feed(bin[3..bin.high])
    var (success, pr) = decode(buf)
    assert success
    assert $pr == "<foobar>"
  BufferedDecoder(
      stream: newStringStream(newStringOfCap(maxSize)),
      maxSize: maxSize,
    )

proc feed*(dec: var BufferedDecoder; buf: pointer; len: int) =
  assert len > 0
  if dec.maxSize > 0 and dec.maxSize < (dec.appendPosition + len):
    raise newException(IOError, "BufferedDecoder at maximum buffer size")
  dec.stream.setPosition(dec.appendPosition)
  dec.stream.writeData(buf, len)
  inc(dec.appendPosition, len)
  assert dec.appendPosition == dec.stream.getPosition()

proc feed*[T: byte|char](dec: var BufferedDecoder; data: openarray[T]) =
  if data.len > 0:
    dec.feed(unsafeAddr data[0], data.len)

proc decode*(dec: var BufferedDecoder; E = void): (bool, Preserve[E]) =
  ## Decode from `dec`. If decoding fails the internal position of the
  ## decoder does not advance.
  if dec.appendPosition > 0:
    assert(dec.decodePosition < dec.appendPosition)
    dec.stream.setPosition(dec.decodePosition)
    try:
      result[1] = decodePreserves(dec.stream, E)
      result[0] = true
      dec.decodePosition = dec.stream.getPosition()
      if dec.decodePosition == dec.appendPosition:
        dec.stream.setPosition(0)
        dec.stream.data.setLen(0)
        dec.appendPosition = 0
        dec.decodePosition = 0
    except IOError:
      discard