CassandraMacros.scala

package com.github.lucasrpb.cassandra

import com.datastax.driver.core.querybuilder.Insert
import com.datastax.driver.core.{AbstractGettableData, GettableByNameData, Row}

import scala.collection.mutable

/**
  * Created by lucasrpb on 6/17/16.
  * Incomplete for some types (but for other types the implementation is analogous)
  */
package object macros {

  import com.datastax.driver.core.{UDTValue, UserType}

  import scala.language.experimental.{macros => Macros}
  import scala.reflect.macros.blackbox.Context
  import scala.reflect.runtime.universe._

  trait CaseClassUDTConverter[T] {
    def convert(obj: T)(implicit types: Map[String, UserType]): UDTValue
    def convert(udt: GettableByNameData): T
  }

  trait RowCaseClassConverter[T] {
      def convert(row: GettableByNameData): T
  }

  trait CaseClassInsertConverter[T] {
      def convert(obj: T, table: String): Insert
  }

  trait CassandraValueConverter[T] {
      def convert(v: T): Any
  }

  protected def getCassandraType(c: Context)(tpe: c.Type, v: c.Tree, attr: c.Tree, udt: c.Tree)(types: c.Tree): c.Tree = {

    import c.universe._

    val s = tpe.typeSymbol

    tpe match {
      case t if t <:< typeOf[String] => {
        q"""
          val a = $attr.asInstanceOf[String]
          if(a != null) $udt.asInstanceOf[UDTValue].setString(a, $v)
          $v
        """
      }
      case t if t <:< typeOf[Int] => {
        q"""
          val a = $attr.asInstanceOf[String]
          if(a != null) $udt.asInstanceOf[UDTValue].setInt(a, $v)
          $v
        """
      }
      case t if t <:< typeOf[Option[_]] => {
        val t0 = t.typeArgs(0)

        q"""
           val v = if($v.isEmpty) null else ${getCassandraType(c)(t0, q"$v.get", attr, udt)(types)}

           val a = $attr.asInstanceOf[String]

           if(a != null && v == null){
              $udt.asInstanceOf[UDTValue].setToNull(a)
           }

           v
        """
      }
      case t if s.isClass && s.asClass.isCaseClass => {
        q"""
             val c = implicitly[CaseClassUDTConverter[$tpe]]
             val value = c.convert($v)($types)

             val a = $attr.asInstanceOf[String]
             if(a != null) $udt.asInstanceOf[UDTValue].setUDTValue(a, value)

             value
            """
      }
      case t if t <:< typeOf[Map[_, _]] => {

        val t0 = t.typeArgs(0)
        val t1 = t.typeArgs(1)

        q"""
          val map = new java.util.HashMap[String, Any]()

          $v.foreach(e => map.put(${getCassandraType(c)(t0, q"e._1", q"null", q"null")(types)}),
            ${getCassandraType(c)(t1, q"e._2", q"null", q"null")(types)})

          val a = $attr.asInstanceOf[String]
          if(a != null) $udt.asInstanceOf[UDTValue].setMap(a, map)

          map
        """
      }
      case t if t <:< typeOf[Iterable[_]] => {

        val t0 = t.typeArgs(0)

        q"""
          val list = new java.util.ArrayList[Any]()

          $v.foreach(e => list.add(${getCassandraType(c)(t0, q"e", q"null", q"null")(types)}))

          val a = $attr.asInstanceOf[String]
          if(a != null) $udt.asInstanceOf[UDTValue].setList(a, list)

          list
        """
      }
      case _ => throw new RuntimeException(s"Type ${tpe} not supported!")
    }
  }

  implicit def caseClassToUDT[T]: macros.CaseClassUDTConverter[T] = macro caseClassToUDTImpl[T]

  def caseClassToUDTImpl[T: c.WeakTypeTag](c: Context): c.Expr[CaseClassUDTConverter[T]] = {

    import c.universe._

    val tpe = weakTypeOf[T]
    val s = tpe.typeSymbol
    val name = s.name.toString
    val fullName = s.fullName.toString
    val companion = s.companion

    val params = tpe.companion.decl(TermName("apply")).asMethod.paramLists.head

    val code1 = mutable.Buffer[Tree]()
    val code2 = mutable.Buffer[Tree]()

    params.map(a => {
      val name = a.name.toTermName
      val sname = name.toString
      val atpe = a.typeSignature

      val transient = a.annotations.exists(_.tree.tpe =:= typeOf[transient])

      if(!transient) {
        code1 += q"${getCassandraType(c)(atpe, q"obj.$name", q"$sname", q"udt")(q"types")}"
      }

      code2 += q"""${getScalaType(c)(atpe, q"null", q"$sname", q"udt", q"$transient")}"""
    })

    c.Expr[CaseClassUDTConverter[T]](
      q"""
         import com.datastax.driver.core.{UDTValue, UserType, Row, GettableByNameData}
         import scala.collection.mutable

          new CaseClassUDTConverter[$tpe]{
              override def convert(obj: $tpe)(implicit types: Map[String, UserType]): UDTValue = {

                val udt = types.get($name).getOrElse(throw new RuntimeException(String.format("Type %s not mapped!",
                  $fullName))).newValue()

                ..$code1

                udt
              }

              override def convert(udt: GettableByNameData): $tpe = $companion(..$code2)
          }
       """)
  }

  implicit def rowToCaseClass[T]: macros.RowCaseClassConverter[T] = macro rowToCaseClassImpl[T]

  def rowToCaseClassImpl[T: c.WeakTypeTag](c: Context): c.Expr[RowCaseClassConverter[T]] = {

    import c.universe._

    val tpe = weakTypeOf[T]
    val s = tpe.typeSymbol
    val name = s.name.toString
    val fullName = s.fullName.toString
    val companion = s.companion

    val params = tpe.companion.decl(TermName("apply")).asMethod.paramLists.head

    val code = params.map(a => {

      val name = a.name.toTermName
      val sname = name.toString
      val atpe = a.typeSignature

      val transient = a.annotations.exists(_.tree.tpe =:= typeOf[transient])

      q"""${getScalaType(c)(atpe, q"null", q"$sname", q"row", q"$transient")}"""
    })

    c.Expr[RowCaseClassConverter[T]](
      q"""
        import com.datastax.driver.core.{UDTValue, UserType, Row, GettableByNameData}
        import scala.collection.mutable

        new RowCaseClassConverter[$tpe]{
            override def convert(row: GettableByNameData): $tpe = $companion(..$code)
        }
       """)
  }

  implicit def caseClassToInsert[T]: macros.CaseClassInsertConverter[T] = macro caseClassToInsertImpl[T]

  def caseClassToInsertImpl[T: c.WeakTypeTag](c: Context): c.Expr[CaseClassInsertConverter[T]] = {

    import c.universe._

    val tpe = weakTypeOf[T]
    val s = tpe.typeSymbol
    val name = s.name.toString
    val fullName = s.fullName.toString
    val companion = s.companion

    val params = tpe.companion.decl(TermName("apply")).asMethod.paramLists.head

    val code = params.filter(!_.annotations.exists(_.tree.tpe =:= typeOf[transient])).map(a => {

      val name = a.name.toTermName
      val sname = name.toString
      val atpe = a.typeSignature

      q"""
        stm.value(String.format("\"%s\"", $sname), ${getCassandraType(c)(atpe, q"o.$name", q"null", q"null")(q"types")})
      """
    })

    c.Expr[CaseClassInsertConverter[T]](
      q"""
        import com.datastax.driver.core.{UDTValue, UserType, Row, GettableByNameData}
        import com.datastax.driver.core.querybuilder.{QueryBuilder, Insert}
        import scala.collection.mutable

        new CaseClassInsertConverter[$tpe]{
            override def convert(o: $tpe, table: String): Insert = {
                  val stm = QueryBuilder.insertInto(table)

                  $code

                  stm
              }
        }
       """)
  }

  protected def getScalaType(c: Context)(tpe: c.Type, v: c.Tree, attr: c.Tree, row: c.Tree, transient: c.Tree): c.Tree = {

    import c.universe._

    val s = tpe.typeSymbol

    tpe match {
      case t if t <:< typeOf[String] => {
        q"""
            val a = $attr.asInstanceOf[String]
            if($transient) null else if(a != null) $row.asInstanceOf[GettableByNameData].getString(a)
              else $v.asInstanceOf[$tpe]
            """
      }
      case t if t <:< typeOf[Int] => {
        q"""
              val a = $attr.asInstanceOf[String]
              if($transient) 0 else if(a != null) $row.asInstanceOf[GettableByNameData].getInt(a)
                else $v.asInstanceOf[$tpe]
            """
      }
      case t if t <:< typeOf[Option[_]] => {
        val t0 = t.typeArgs(0)
        q"""
          if($transient.asInstanceOf[Boolean]) null else {
              val a = $attr.asInstanceOf[String]

              val v = if(a != null) $row.asInstanceOf[GettableByNameData].getObject(a).asInstanceOf[Any]
                  else $v.asInstanceOf[Any]

              if(v == null) None else Some(${getScalaType(c)(t0, q"v", q"null", q"null", q"false")}.asInstanceOf[$t0])
            }
        """
      }
      case t if s.isClass && s.asClass.isCaseClass => {
        q"""
              val a = $attr.asInstanceOf[String]

              if($transient) null else {
                val c = implicitly[CaseClassUDTConverter[$tpe]]
                c.convert(if(a != null) $row.asInstanceOf[GettableByNameData].getUDTValue(a) else $v.asInstanceOf[UDTValue])
              }
            """
      }
      case t if t <:< typeOf[Map[_, _]] => {

        val t0 = t.typeArgs(0)
        val t1 = t.typeArgs(1)

        q"""
             val a = $attr.asInstanceOf[String]
             val values = if(a != null) $row.asInstanceOf[GettableByNameData].getObject(a)
                .asInstanceOf[java.util.Map[Any, Any]].entrySet.iterator else $v.asInstanceOf[java.util.Map[Any, Any]]
                .entrySet.iterator

             val map = scala.collection.mutable.Map[$t0, $t1]()

             while(values.hasNext){
                val e = values.next
                val k = ${getScalaType(c)(t0, q"e.getKey", q"null", q"null", q"false")}.asInstanceOf[$t0]
                val value = ${getScalaType(c)(t1, q"e.getValue", q"null", q"null", q"false")}.asInstanceOf[$t1]

                map += k -> value
             }

             map.toMap
          """
      }
      case t if t <:< typeOf[Iterable[_]] => {

        val t0 = t.typeArgs(0)

        q"""
            val a = $attr.asInstanceOf[String]
            val values = if(a != null) $row.asInstanceOf[GettableByNameData].getObject(a).asInstanceOf[java.util.List[Any]]
              else $v.asInstanceOf[java.util.List[Any]]

            val list = scala.collection.mutable.Buffer[$t0]()
            val len = values.size

            for(i<-0 until len){
              val value = values.get(i)
              list += ${getScalaType(c)(t0, q"value", q"null", q"null", q"false")}.asInstanceOf[$t0]
            }

            list.toSeq
          """
      }
      case _ => throw new RuntimeException(s"Type ${tpe} not supported!")
    }
  }

  implicit def toCassandraValue[T](implicit types: Map[String, UserType]): macros.CassandraValueConverter[T] =
    macro toCassandraValueImpl[T]

  def toCassandraValueImpl[T: c.WeakTypeTag](c: Context)(types: c.Expr[Map[String, UserType]]):
    c.Expr[CassandraValueConverter[T]] = {

    import c.universe._

    val tpe = weakTypeOf[T]

    c.Expr[CassandraValueConverter[T]](
      q"""
          import com.datastax.driver.core.{UDTValue, UserType}

          new CassandraValueConverter[$tpe] {
              override def convert(v: $tpe): Any = ${getCassandraType(c)(tpe, q"v", q"null", q"null")(types.tree)}
          }
       """)
  }

}