import scala.reflect.macros._
import scala.language.experimental.macros
import scala.annotation.StaticAnnotation

class freeze extends StaticAnnotation {
  def macroTransform(annottees: Any*): Any = macro freeze.annotationImpl
}

object freeze {
  def mangleName(name: String) = "freeze$" + name

  def annotationImpl(c: blackbox.Context)(annottees: c.Expr[Any]*): c.Expr[Any] = {
    import c.universe._
    import c.universe.Flag._

    val inputs = annottees.map(_.tree).toList
    val (annottee, expandees) = inputs match {
      case (param: ValDef) :: (rest @ (_ :: _)) => (param, rest)
      case (param: TypeDef) :: (rest @ (_ :: _)) => (param, rest)
      case _ => (EmptyTree, inputs)
    }

    if (!annottee.isEmpty) {
      c.abort(c.enclosingPosition, "@freeze annotation on invalid argument")
    }

    val transformedExpandee = expandees map {
      case ClassDef(modifiers, typename, tparams, Template(parents, self, body)) => {
        case class Mangled(mangledName: String, modifiers: Modifiers, typeName: TypeName)
        var nameMangling = Map[String, Mangled]()
        var parameters = Set[String]()

        // Do a pass to find all of the names that need to be rewritten
        body foreach {
          definition => definition match {
            case ValDef(modifiers, TermName(name), tpt, rhs) => {
              val Modifiers(flags, privateWithin, annotations) = modifiers

              if ((flags.## & MUTABLE.##) == MUTABLE.##) {
                if (tpt.isEmpty) {
                  c.abort(definition.pos, "@freeze requires explicit type declarations for all vars")
                }

                nameMangling = nameMangling + (name -> Mangled(mangleName(name), modifiers, tpt.asInstanceOf[Ident].name.asInstanceOf[TypeName]))

                if ((flags.## & PARAMACCESSOR.##) == PARAMACCESSOR.##) {
                  parameters = parameters + name
                }
              }
            }

            case _ => {}
          }
        }

        // Hack, assume that the default constructor is always the first one in the AST
        var foundConstructor = false

        // Perform the mangling and add generated getters
        val transformed = body map {
          case ValDef(modifiers, TermName(name), tpt, rhs) => {
            nameMangling.get(name).map(_.mangledName) match {
              case Some(properName) => {
                val Modifiers(flags, privateWithin, annotations) = modifiers
                ValDef(Modifiers(flags | SYNTHETIC | FINAL, privateWithin, annotations), TermName(properName), tpt, rhs)
              }

              case None => {
                ValDef(modifiers, TermName(name), tpt, rhs)
              }
            }
          }

          case defDef: DefDef => {
            val DefDef(modifiers, name, tparams, vparamss, tpt, rhs) = defDef

            if (!foundConstructor && name == termNames.CONSTRUCTOR) {
              foundConstructor = true

              if (vparamss.size != 1) {
                c.abort(defDef.pos, "<init> has multiple vparamss")
              }

              val fixedParameters = vparamss(0) map {
                parameter => {
                  val ValDef(modifiers, TermName(name), tpt, rhs) = parameter

                  val finalName = nameMangling.get(name).map(_.mangledName).getOrElse(name)
                  ValDef(modifiers, TermName(finalName), tpt, rhs)
                }
              }

              DefDef(modifiers, name, tparams, List(fixedParameters), tpt, rhs)
            } else {
              defDef
            }
          }
        }

        var getters = Seq[DefDef]()
        var setters = Seq[DefDef]()

        nameMangling.toTraversable.foreach {
          case (original, mangled) => {
            val getterName = TermName(original)
            val setterName = TermName(original + "_$eq")
            val originalType = mangled.typeName
            val mangledName = TermName(mangled.mangledName)

            // Modifiers we care about: final, private, protected, sealed
            val interestingModifiers = Set(FINAL, PRIVATE, PROTECTED, SEALED)
            val declaredModifiers = interestingModifiers.filter(mangled.modifiers.hasFlag _)

            var modifiers = declaredModifiers.fold(NoFlags)((lhs: FlagSet, rhs: FlagSet) => lhs | rhs)
            val implicitGate = TypeName("freeze$ImplicitGate")

            getters = q"$modifiers def $getterName: $originalType = $mangledName" +: getters
            setters =
              q"""
                $modifiers def $setterName(rhs: $originalType) = {
                  $mangledName = rhs
                }
              """ +: setters
          }
        }

        println(getters)
        println(setters)
        val clone = {
          q"def clone(mutator: $typename => Any): $typename = macro freeze.cloneImpl[$typename]"
        }

        val traitType = AppliedTypeTree(Ident(TypeName("freeze$Frozen")), List(Ident(typename)))
        ClassDef(modifiers, typename, tparams, Template(parents ++ List(traitType), self, transformed ++ getters ++ Seq(clone)))
      }

      case unknown => {
        println("unknown! " + unknown)
        unknown
      }
    }
    c.Expr[Any](Block(transformedExpandee, Literal(Constant(()))))
  }

  def cloneImpl[T: c.WeakTypeTag](c: blackbox.Context)(mutator: c.Expr[T => T]): c.Expr[T] = {
    import c.universe._
    import c.universe.Flag._

    val prefix = c.prefix
    println(showRaw(mutator.tree))
    mutator.tree match {
      case q"(..$params) => $exprs" => {
        val implicitName = TermName(c.freshName("freeze$"))
        val implicitObject = Ident(TermName("freeze$ImplicitKey"))
        val implicitDef = ValDef(Modifiers(IMPLICIT), implicitName, TypeTree(), implicitObject)

        val block =
          q"""
            {
              implicit val $implicitName = $implicitObject
              $exprs
            }
          """

        println("params: " + showRaw(params))
        println("exprs: " + showRaw(exprs))
        println("result: " + showRaw(block))

        val expr = c.Expr[T](q"((..$params) => $block)($prefix)")
        println("expr: " + expr)
        expr
      }

      case _ => {
        c.abort(mutator.tree.pos, "clone must be called with a function literal")
        null.asInstanceOf[c.Expr[T]]
      }
    }
  }
}