package main.utils

object ByteUtils {
  trait SizedNumeric[T] {
    def manifest: Manifest[T]
    def byteSize: Int
    def bitSize: Int
    def fromLong(l: Long): T
    def toLong(l: T): Long
    def buildFrom[U : SizedNumeric](from: U): T = fromLong(from.toLong)
    def truncate[U : SizedNumeric](from: U): T = buildFrom(from)
    override def toString = s"SizedNumeric[${manifest}]"
  }

  trait BitOps[T] {
    def @>>>(t: T, x: Int): T
    def @<<(t: T, x: Int): T
    def @|(a: T, b: T): T
    def @&(a: T, b: T): T
    def @+(a: T, b: T): T
  }

  implicit val byteOps: BitOps[Byte] = new BitOps[Byte] {
    def @>>>(t: Byte, x: Int): Byte = (t >>> x).toByte
    def @<<(t: Byte, x: Int): Byte = (t << x).toByte
    def @|(a: Byte, b: Byte): Byte = (a | b).toByte
    def @&(a: Byte, b: Byte): Byte = (a & b).toByte
    def @+(a: Byte, b: Byte): Byte = (a + b).toByte
  }

  implicit val intOps: BitOps[Int] = new BitOps[Int] {
    def @>>>(t: Int, x: Int): Int = t >>> x
    def @<<(t: Int, x: Int): Int = (t << x)
    def @|(a: Int, b: Int): Int = a | b
    def @&(a: Int, b: Int): Int = a & b
    def @+(a: Int, b: Int): Int = a + b
  }

  implicit class WithBitOpts[T : BitOps](x: T) {
    def @>>>(b: Int): T = implicitly[BitOps[T]].@>>>(x, b)
    def @<<(b: Int): T = implicitly[BitOps[T]].@<<(x, b)
    def @|(b: T): T = implicitly[BitOps[T]].@|(x, b)
    def @&(b: T): T = implicitly[BitOps[T]].@&(x, b)
    def @+(b: T): T = implicitly[BitOps[T]].@+(x, b)
  }

  implicit class SizedWithByteInfo[T : SizedNumeric](x: T) {
    def byteSize = implicitly[SizedNumeric[T]].byteSize
    def bitSize = implicitly[SizedNumeric[T]].bitSize
    def toLong = implicitly[SizedNumeric[T]].toLong(x)
    def liftedTo[U : SizedNumeric]: U = {
      val uSize = implicitly[SizedNumeric[U]]
      val tSize = implicitly[SizedNumeric[T]]
      if (uSize.byteSize < tSize.byteSize) {
        throw new RuntimeException(s"Target size ($uSize) is smaller than the source size ($tSize)")
      }

      uSize.buildFrom(x)
    }
    def truncatedTo[U : SizedNumeric]: U = {
      implicitly[SizedNumeric[U]].truncate(x)
    }
  }

  implicit class RichArrayOfBytes(b: Array[Byte]) {
    def to[T : SizedNumeric : BitOps]: T = {
      val numeric = implicitly[SizedNumeric[T]]
      if(b.length < numeric.byteSize) {
        throw new RuntimeException("Byte input is not long enough")
      }

      var out: Long = b(0)
      for(i <- 1 until numeric.byteSize) {
        out << 8
        out = b(i) | out
      }
      numeric.fromLong(out)
    }
  }
  def sizedNumeric[T : Manifest](bytes: Int)(from: Long => T)(to: T => Long) = new SizedNumeric[T] {
    def manifest = implicitly[Manifest[T]]
    def byteSize = bytes
    def bitSize = bytes * 8
    def toLong(t: T) = to(t)
    def fromLong(l: Long) = from(l)
  }

  implicit val intSized = sizedNumeric(4) { x: Long => x.toInt } { _.toLong }
  implicit val shortized = sizedNumeric(2) { x: Long => x.toShort } { _.toLong }
  implicit val byteSized = sizedNumeric(1) { x: Long => x.toByte } { _.toLong }

  def toArrayBuf[T : SizedNumeric : BitOps](x: T): Array[Byte] = {
    val buf = new Array[Byte](x.byteSize)
    for(i <- 0 until x.byteSize) {
      val shiftAmount: Int = (x.byteSize - i - 1) << 3
      buf(i) = (x @>>> shiftAmount).truncatedTo[Byte]
    }
    buf
  }

  def toBinaryString[T : SizedNumeric](x: T)(implicit a: BitOps[T]): String = {
    val buf = new StringBuilder(x.bitSize)
    for(i <- 0 until x.bitSize) {
      val shiftAmount: Int = x.bitSize - i - 1
      buf.append((x @>>> shiftAmount) @& 0x01.liftedTo[T])
    }
    buf.toString()
  }
}