scodec.scala

package com.rr.experiment

import org.specs2.ScalaCheck
import org.specs2.mutable._

import org.scalacheck._

import scalaz._

import scodec._
import scodec.bits._
import scodec.codecs._

import shapeless._

import java.nio.charset.Charset
import java.util.UUID

object BasicSpecs extends Specification with ScalaCheck {
  import poly._
  import ops.hlist._

  implicit def const(bits: Int): Codec[Unit] = scodec.codecs.constant(BitVector fromInt bits.toByte)

  // heaven forgive me
  implicit val unitMonoid: Monoid[Unit] = new Monoid[Unit] {
    def zero = ()
    def append(left: Unit, right: => Unit): Unit = ()
  }

  "version 1" should {
    "code a point2" in check { (x: Int, y: Int) =>
      implicit val codec = point2(Version.v1)

      val p = Point2(x, y)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point2](bits)
      } yield result

      rest must beEmpty
      p2 mustEqual p
    }

    implicit val arbEnum3: Arbitrary[Enum3] = Arbitrary(Gen.oneOf(Enum3.A, Enum3.B, Enum3.C))

    "code an Enum3" in check { e: Enum3 =>
      implicit val codec = enum3(Version.v1)

      val \/-((rest, e2)) = for {
        bits <- Codec.encode(e)
        result <- Codec.decode[Enum3](bits)
      } yield result

      rest must beEmpty
      e2 mustEqual e
    }
  }

  "version 2" should {
    "code a point3 via version1" in check { (x: Int, y: Int, z: Int) =>
      implicit val codec = point3(Version.v1)

      val p = Point3(x, y, z)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point3](bits)
      } yield result

      rest must beEmpty

      p2.x mustEqual p.x
      p2.y mustEqual p.y
      p2.z mustEqual 0
    }

    "code a point3" in check { (x: Int, y: Int, z: Int) =>
      implicit val codec = point3(Version.v2)

      val p = Point3(x, y, z)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point3](bits)
      } yield result

      rest must beEmpty
      p2 mustEqual p
    }

    implicit val arbEnum4: Arbitrary[Enum4] = Arbitrary(Gen.oneOf(Enum4.A, Enum4.B, Enum4.C, Enum4.D))

    "code an Enum4 via version1" in check { e: Enum4 =>
      implicit val codec = enum4(Version.v1)

      val \/-((rest, e2)) = for {
        bits <- Codec.encode(Some(e))
        result <- Codec.decode[Option[Enum4]](bits)
      } yield result

      rest must beEmpty

      if (e == Enum4.D) {
        e2 must beNone
      } else {
        e2 must beSome(e)
      }
    }

    "code an Enum4" in check { e: Enum4 =>
      implicit val codec = enum4(Version.v2)

      val \/-((rest, e2)) = for {
        bits <- Codec.encode(Some(e))
        result <- Codec.decode[Option[Enum4]](bits)
      } yield result

      rest must beEmpty
      e2 must beSome(e)
    }
  }

  "version 3" should {
    "code a point3d via version1" in check { (x: Double, y: Double, z: Double) =>
      implicit val codec = point3d(Version.v1)

      val p = Point3D(x, y, z)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point3D](bits)
      } yield result

      rest must beEmpty

      // of course it's lossy, Double => Int?!?!
      p2.x mustEqual p.x.toInt.toDouble
      p2.y mustEqual p.y.toInt.toDouble
      p2.z mustEqual 0d
    }

    "code a point3d via version2" in check { (x: Double, y: Double, z: Double) =>
      implicit val codec = point3d(Version.v2)

      val p = Point3D(x, y, z)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point3D](bits)
      } yield result

      rest must beEmpty

      // of course it's lossy, Double => Int?!?!
      p2.x mustEqual p.x.toInt.toDouble
      p2.y mustEqual p.y.toInt.toDouble
      p2.z mustEqual p.z.toInt.toDouble
    }

    "code a point3d" in check { (x: Double, y: Double, z: Double) =>
      implicit val codec = point3d(Version.v3)

      val p = Point3D(x, y, z)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point3D](bits)
      } yield result

      rest must beEmpty
      p2 mustEqual p
    }

    implicit val arbEnum4: Arbitrary[Enum4] = Arbitrary(Gen.oneOf(Enum4.A, Enum4.B, Enum4.C, Enum4.D))

    "code an Enum4 via version1" in check { e: Enum4 =>
      implicit val codec = enum4(Version.v1)

      val \/-((rest, e2)) = for {
        bits <- Codec.encode(Some(e))
        result <- Codec.decode[Option[Enum4]](bits)
      } yield result

      rest must beEmpty

      if (e == Enum4.D) {
        e2 must beNone
      } else {
        e2 must beSome(e)
      }
    }

    "code an Enum4 via version2" in check { e: Enum4 =>
      implicit val codec = enum4(Version.v2)

      val \/-((rest, e2)) = for {
        bits <- Codec.encode(Some(e))
        result <- Codec.decode[Option[Enum4]](bits)
      } yield result

      rest must beEmpty
      e2 must beSome(e)
    }

    "code an Enum4" in check { e: Enum4 =>
      implicit val codec = enum4(Version.v3)

      val \/-((rest, e2)) = for {
        bits <- Codec.encode(Some(e))
        result <- Codec.decode[Option[Enum4]](bits)
      } yield result

      rest must beEmpty
      e2 must beSome(e)
    }
  }

  "version 4" should {
    "code a point2d via version1" in check { (x: Double, y: Double) =>
      implicit val codec = point2d(Version.v1)

      val p = Point2D(x, y)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point2D](bits)
      } yield result

      rest must beEmpty

      // of course it's lossy, Double => Int?!?!
      p2.x mustEqual p.x.toInt.toDouble
      p2.y mustEqual p.y.toInt.toDouble
    }

    "code a point2d via version2" in check { (x: Double, y: Double) =>
      implicit val codec = point2d(Version.v2)

      val p = Point2D(x, y)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point2D](bits)
      } yield result

      rest must beEmpty

      // of course it's lossy, Double => Int?!?!
      p2.x mustEqual p.x.toInt.toDouble
      p2.y mustEqual p.y.toInt.toDouble
    }

    "code a point2d via version3" in check { (x: Double, y: Double) =>
      implicit val codec = point2d(Version.v3)

      val p = Point2D(x, y)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point2D](bits)
      } yield result

      rest must beEmpty

      p2.x mustEqual p.x
      p2.y mustEqual p.y
    }

    "code a point2d" in check { (x: Double, y: Double) =>
      implicit val codec = point2d(Version.v4)

      val p = Point2D(x, y)
      val \/-((rest, p2)) = for {
        bits <- Codec.encode(p)
        result <- Codec.decode[Point2D](bits)
      } yield result

      rest must beEmpty
      p2 mustEqual p
    }

    implicit val arbUUID: Arbitrary[UUID] = Arbitrary(for {
      high <- Arbitrary.arbitrary[Long]
      low <- Arbitrary.arbitrary[Long]
    } yield new UUID(high, low))

    "code stuff" in check { (uuid: UUID, strings: List[String]) =>
      implicit val codec = stuff(Version.v4)

      val s = Stuff(uuid, strings)
      val \/-((rest, s2)) = for {
        bits <- Codec.encode(s)
        result <- Codec.decode[Stuff](bits)
      } yield result

      rest must beEmpty
      s2 mustEqual s
    }

    "code a list of stuff" in check { data: List[(UUID, List[String])] =>
      implicit val codec = list(stuff(Version.v4))

      val stuffs = data map Stuff.tupled
      val \/-((rest, stuffs2)) = for {
        bits <- Codec.encode(stuffs)
        result <- Codec.decode[List[Stuff]](bits)
      } yield result

      rest must beEmpty
      stuffs2 mustEqual stuffs
    }
  }

  "recover" should {
    "always code a value" in check { (i: Int, b: Boolean) =>
      val codec = recover(constant(BitVector fromInt i))

      val \/-((rest, b2)) = for {
        bits <- codec.encode(b)
        result <- Codec.decode(bits)(codec)
      } yield result

      rest must beEmpty
      b2 must beTrue
    }

    "decode an empty value as false" in check { i: Int =>
      val codec = recover(constant(BitVector fromInt i))
      val \/-((rest, b)) = Codec.decode(BitVector.empty)(codec)

      rest must beEmpty
      b must beFalse
    }

    "decode the wrong value as false and backtrack" in check { (i1: Int, i2: Int) =>
      if (i1 != i2 && i1 >= 0 && i2 >= 0) {
        val codec = recover(constant(BitVector fromInt i1))
        val \/-((rest, b)) = Codec.decode(BitVector fromInt i2)(codec)

        rest mustEqual (BitVector fromInt i2)
        b must beFalse
      } else {
        ok
      }
    }
  }

  "optional" should {
    "produce the target value on true" in check { i: Int =>
      val codec = optional(provide(true), int32)

      val \/-((rest, b)) = codec.decode(BitVector fromInt i)

      rest must beEmpty
      b must beSome(i)
    }

    "produce none on false" in check { i: Int =>
      val codec = optional(provide(false), int32)

      val \/-((rest, b)) = codec.decode(BitVector fromInt i)

      rest mustEqual (BitVector fromInt i)
      b must beNone
    }
  }

  def optional[A](guard: Codec[Boolean], target: Codec[A]): Codec[Option[A]] =
    either(guard, provide(()), target).xmap[Option[A]]({ _.toOption }, { _ map { \/-(_) } getOrElse -\/(()) })

  def withDefault[A](opt: Codec[Option[A]], default: Codec[A]): Codec[A] = {
    val paired = opt flatZip {
      case Some(a) => provide(a)
      case None => default
    }

    paired.xmap[A]({ _._2 }, { a => (Some(a), a) })
  }

  def recover(codec: Codec[Unit]): Codec[Boolean] = new Codec[Boolean] {
    def encode(a: Boolean): String \/ BitVector = codec.encode(())

    def decode(buffer: BitVector): String \/ (BitVector, Boolean) = {
      codec.decode(buffer).toOption map {
        case (rest, _) => \/-(rest, true)
      } getOrElse \/-(buffer, false)      // backtrack on failure
    }
  }

  // like recover, but produces the original buffer rather than the consumed one when true
  def lookahead(codec: Codec[Unit]): Codec[Boolean] = new Codec[Boolean] {
    def encode(a: Boolean): String \/ BitVector = codec.encode(())

    def decode(buffer: BitVector): String \/ (BitVector, Boolean) = {
      codec.decode(buffer).toOption map {
        case (_, _) => \/-(buffer, true)     // backtrack on success!
      } getOrElse \/-(buffer, false)      // backtrack on failure
    }
  }

  // greedy RLE
  def list[A](codec: Codec[A]): Codec[List[A]] =
    variableSizeBytes(int32, repeated(codec)).xmap({ _.toList }, { _.toVector })

  def rleString(implicit cs: Charset): Codec[String] =
    variableSizeBytes(int32, string)

  implicit class RichCodec[A](self: Codec[A]) {
    // this exists in scodec 1.2.1
    def unit(zero: A): Codec[Unit] = self.xmap[Unit]({ _ => () }, { _ => zero })
  }

  implicit class ListCodec[L <: HList](self: Codec[L]) {
    def pmap[L2 <: HList](p: Poly)(implicit m: Mapper.Aux[p.type, L, L2], m2: Mapper.Aux[p.type, L2, L]): Codec[L2] =
      self.xmap({ _ map p }, { _ map p })
  }

  ////////////////////////////////////////////////////////////////////

  sealed trait Version

  object Version {
    case object v1 extends Version
    case object v2 extends Version
    case object v3 extends Version
    case object v4 extends Version
  }

  object i2d extends Poly1 {
    implicit def caseInt = at[Int] { _.toDouble }
    implicit def caseDouble = at[Double] { _.toInt }
  }

  val point2: Version => Codec[Point2] = {
    case Version.v1 => (int32 :: int32).as
  }

  val point3: Version => Codec[Point3] = {
    case Version.v1 => (int32 :: int32 :: provide(0)).as
    case Version.v2 => (int32 :: int32 :: int32).as
  }

  val point3d: Version => Codec[Point3D] = {
    case Version.v1 => (int32 :: int32 :: provide(0)).pmap(i2d).as
    case Version.v2 => (int32 :: int32 :: int32).pmap(i2d).as
    case Version.v3 => (double :: double :: double).as
  }

  val point2d: Version => Codec[Point2D] = {
    case Version.v1 => (int32 :: int32).pmap(i2d).as
    case Version.v2 => (int32 :: int32 <~ int32.unit(0)).pmap(i2d).as
    case Version.v3 => (double :: double <~ double.unit(0)).as
    case Version.v4 => (double :: double).as
  }

  val enum3: Version => Codec[Enum3] = {
    case Version.v1 => {
      discriminated[Enum3].by(uint8)
        .typecase(0, provide(Enum3.A))
        .typecase(1, provide(Enum3.B))
        .typecase(2, provide(Enum3.C))
    }
  }

  // we're assuming that we need to handle the missing case in the application layer
  val enum4: Version => Codec[Option[Enum4]] = {
    case Version.v1 => {
      discriminated[Option[Enum4]].by(uint8)
        .subcaseP(0)({ case s @ Some(Enum4.A) => s })(provide(Some(Enum4.A)))
        .subcaseP(1)({ case s @ Some(Enum4.B) => s })(provide(Some(Enum4.B)))
        .subcaseP(2)({ case s @ Some(Enum4.C) => s })(provide(Some(Enum4.C)))
        .subcaseP(3)({ case None | Some(Enum4.D) => None })(provide(None))
    }

    case Version.v2 | Version.v3 => {
      discriminated[Option[Enum4]].by(uint8)
        .subcaseP(0)({ case s @ Some(Enum4.A) => s })(provide(Some(Enum4.A)))
        .subcaseP(1)({ case s @ Some(Enum4.B) => s })(provide(Some(Enum4.B)))
        .subcaseP(2)({ case s @ Some(Enum4.C) => s })(provide(Some(Enum4.C)))
        .subcaseP(3)({ case s @ Some(Enum4.D) => s })(provide(Some(Enum4.D)))
    }
  }

  val stuff: Version => Codec[Stuff] = {
    case Version.v4 => (uuid :: list(rleString(Charset.defaultCharset))).as
  }

  // the datatypes below represent different versions of the *same* datatype; they would not coexist

  // v1
  case class Point2(x: Int, y: Int)

  sealed trait Enum3

  object Enum3 {
    case object A extends Enum3
    case object B extends Enum3
    case object C extends Enum3
  }

  // v2
  case class Point3(x: Int, y: Int, z: Int)

  sealed trait Enum4

  object Enum4 {
    case object A extends Enum4
    case object B extends Enum4
    case object C extends Enum4
    case object D extends Enum4
  }

  // v3
  case class Point3D(x: Double, y: Double, z: Double)

  // v3 also includes Enum4

  // v4
  case class Point2D(x: Double, y: Double)

  case class Stuff(uuid: UUID, strings: List[String])
}